From a6df71f9b53d4ff8ebc1150b7adba063fba2d4af Mon Sep 17 00:00:00 2001 From: Eitan Turok Date: Mon, 30 Sep 2024 17:23:32 +0000 Subject: [PATCH] remove ic --- llmfoundry/command_utils/train.py | 2 +- llmfoundry/utils/config_utils.py | 5 +++-- tests/tp/test_tp_strategies.py | 4 ---- 3 files changed, 4 insertions(+), 7 deletions(-) diff --git a/llmfoundry/command_utils/train.py b/llmfoundry/command_utils/train.py index f3ef486d61..6749d3abdd 100644 --- 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 index 75f2d3ffdd..7adc593772 100644 --- 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 index 085e51787f..1d23096832 100644 --- 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')