From 3a6dec61ffc9b9b23a1e3f163725712906fb85f2 Mon Sep 17 00:00:00 2001
From: Eitan Turok <eitan.turok@gmail.com>
Date: Wed, 25 Sep 2024 20:13:10 +0000
Subject: [PATCH] remove test for tp training b/c in composer

---
 tests/models/utils/test_tp_strategy.py | 48 --------------------------
 1 file changed, 48 deletions(-)

diff --git a/tests/models/utils/test_tp_strategy.py b/tests/models/utils/test_tp_strategy.py
index 0d4d04a023..0e0079ad7e 100644
--- a/tests/models/utils/test_tp_strategy.py
+++ b/tests/models/utils/test_tp_strategy.py
@@ -65,51 +65,3 @@ def test_ffn_tp_strategy_layer_plan():
             assert lp1.input_layouts == lp2.input_layouts
             assert lp1.output_layouts == lp2.output_layouts
             assert lp1.use_local_output == lp2.use_local_output
-
-
-@pytest.mark.gpu
-@pytest.mark.world_size(4)
-def test_tp_train(tiny_ft_dataloader: DataLoader):
-    ic(tiny_ft_dataloader)
-
-    # Actual layer plan
-    tp_config = {
-        'strategy': 'ffn',
-        }
-
-    model_cfg = {
-        'name': 'mpt_causal_lm',
-        'd_model': 128,
-        'n_heads': 4,
-        'n_layers': 3,
-        'expansion_ratio': 1,
-        'max_seq_len': 16,
-        'vocab_size': 50368,
-        'attn_config': {
-            'attn_impl': 'flash',
-        },
-    }
-    model = ComposerMPTCausalLM(**model_cfg)
-    layer_plan = build_tp_strategy(tp_config['strategy'], model)
-
-    # dataset = RandomClassificationDataset(size=8)
-    # dataloader = DataLoader(dataset, batch_size=2, sampler=dist.get_sampler(dataset))
-
-    trainer = Trainer(
-        model=model,
-        train_dataloader=tiny_ft_dataloader,
-        parallelism_config={
-            'tp': {
-                'layer_plan': layer_plan,
-                'tensor_parallel_degree': 2,
-            },
-            'fsdp': {},
-        },
-        max_duration='3ba',
-    )
-
-    trainer.fit()
-
-
-# if __name__ == '__main__':
-#     test_tp_train()