better init for parallelism_config

mosaicml · Sep 25, 2024 · 2b9664e · 2b9664e
1 parent 19f5477
commit 2b9664e
Showing 1 changed file with 1 addition and 3 deletions.
diff --git a/llmfoundry/command_utils/train.py b/llmfoundry/command_utils/train.py
@@ -530,9 +530,7 @@ def train(cfg: DictConfig) -> Trainer:
             tp_config['layer_plan'] |= strategy_layer_plan
 
     # Parallelism config
-    tp = TPConfig(**tp_config)
-    fsdp = FSDPConfig(**fsdp_config)
-    parallelism_config = ParallelismConfig(fsdp=fsdp, tp=tp)
+    parallelism_config = dict(fsdp=fsdp_config, tp=tp_config)
 
     # Optimizer
     optimizer_name: str = train_cfg.optimizer.pop('name')