fix

Signed-off-by: Yu Chin Fabian Lim <flim@sg.ibm.com>
foundation-model-stack · Aug 27, 2024 · 59907e3 · 59907e3
1 parent 329c4d0
commit 59907e3
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 8 deletions.
diff --git a/plugins/accelerated-moe/src/fms_acceleration_moe/framework_plugin_megablocks.py b/plugins/accelerated-moe/src/fms_acceleration_moe/framework_plugin_megablocks.py
@@ -168,21 +168,16 @@ def get_callbacks_and_ready_for_train(
             accelerator is not None
             and getattr(accelerator.state, "fsdp_plugin", None) is not None
         ):
-            # - use an internal function call to get the no split 
+            # - use an internal function call to get the no split
             # module names, which are typically layers
             _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
 
 

diff --git a/scripts/benchmarks/benchmark.py b/scripts/benchmarks/benchmark.py
@@ -698,7 +698,7 @@ def prepare_arguments(args, benchmark_dataset: BenchmarkDataset):
         scenario = ScenarioMatrix(scenario_config, acceleration_config_map)
 
         if (
-            not args.run_only_scenarios and 
+            not args.run_only_scenarios
             and scenarios.slow
         ):
             # unfiltered runs omit all "slow" marked scenarios