fix

Signed-off-by: Yu Chin Fabian Lim <[email protected]>

Author: fabianlim

plugins/accelerated-moe/src/fms_acceleration_moe/framework_plugin_megablocks.py

@@ -173,16 +173,11 @@ def get_callbacks_and_ready_for_train(
             _layers = model._get_no_split_modules('')
             accelerator.state.fsdp_plugin.ignored_modules = [
                 getattr(layer, name)
-                for name in moe_component_module_names
+                for name in self._moe_component_module_names
                 for layer in model.modules()
                 if layer.__class__.__name__ in _layers
             ]
 
-FSDP(
-    model,
-    ignored_modules=ignored_modules,
-)
-
         return callbacks