remove ic

mosaicml · Sep 30, 2024 · a6df71f · a6df71f
1 parent 81130c8
commit a6df71f
Show file tree

Hide file tree

Showing 3 changed files with 4 additions and 7 deletions.
diff --git a/llmfoundry/command_utils/train.py b/llmfoundry/command_utils/train.py
@@ -520,7 +520,7 @@ def train(cfg: DictConfig) -> Trainer:
         layer_plan = build_tp_strategies(strategy, model)
         tp_config = {
             'layer_plan': layer_plan,
-            'tensor_parallel_degree': tp_config['tensor_parallel_degree']
+            'tensor_parallel_degree': tp_config['tensor_parallel_degree'],
         }
 
     # Parallelism config

diff --git a/llmfoundry/utils/config_utils.py b/llmfoundry/utils/config_utils.py
@@ -542,12 +542,13 @@ def process_init_device(
         # Check tp_config has required fields
         if 'strategy' not in tp_config or 'tensor_parallel_degree' not in tp_config:
             raise ValueError(
-                "`tp_config` requires 'strategy' and 'tensor_parallel_degree' values. "
+                "`tp_config` requires 'strategy' and 'tensor_parallel_degree' values. ",
             )
 
         # Check we are not using tensor parallelism with MoEs
         if 'ffn_config' in model_cfg and model_cfg['ffn_config'].get(
-            'ffn_type', None
+            'ffn_type',
+            None,
         ) in ffns_with_megablocks:
             raise ValueError(
                 'Tensor Parallelism is not currently supported for MoE models.',

diff --git a/tests/tp/test_tp_strategies.py b/tests/tp/test_tp_strategies.py
@@ -163,7 +163,3 @@ def test_tp_train_with_moes():
         match='Tensor Parallelism is not currently supported for MoE models.',
     ):
         process_init_device(model_cfg, fsdp_cfg, tp_cfg)
-
-
-if __name__ == '__main__':
-    test_tp_train('ffn')