Skip to content

Commit

Permalink
remove test for tp training b/c in composer
Browse files Browse the repository at this point in the history
  • Loading branch information
eitanturok committed Sep 25, 2024
1 parent 5e58dbc commit 3a6dec6
Showing 1 changed file with 0 additions and 48 deletions.
48 changes: 0 additions & 48 deletions tests/models/utils/test_tp_strategy.py
Original file line number Diff line number Diff line change
Expand Up @@ -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()

0 comments on commit 3a6dec6

Please sign in to comment.