Tensor Parallelism

[eitanturok]

Implement Tensor Parallelism (TP) in foundry.

To do:

• Make tp_strategy registry
• Check TP: if world_size == 1 or training an MoE, don't apply TP
• Get the same loss with fsdp and fsdp-tp
• Add Tests

Updates:
I compared training 125m param models for 100 steps on c4 with tp-fsdp VS fsdp.

• brown = tp-fsdp
• gray = fsdp

loss_train_total:

throughput_batches_per_sec:

memory_peak_reserved_mem:

It is okay that we don't see performance improvements here yet -- we'll get those later, in follow up PRs.

[eitanturok]

Currently, the ffn strategy gives different results when we train fsdp vs fsdp-tp.

See mcli runs:

• mpt-125m-tp-fsdp-IHKG5s
• mpt-125m-fsdp-8OgGjt

and here are their losses which are visibly different.

Currently investigating, though I think this may have to do more with my specific layer plan/strategy then with anything else.

[snarayan21]

Will review once loss discrepancy has been addressed. good to see it's at least mechanically working though

[eitanturok]

Should we include a tp-mpt-125m.yaml in the repo?

[dakinggg]

@eitanturok does checkpointing work now?

[mvpatel2000]

LGTM!

[mvpatel2000]

@eitanturok does checkpointing work now?

No and it won't with FSDPv1.

[dakinggg]

@mvpatel2000 @eitanturok ok lets leave the yaml out then

[dakinggg]

Also can we log a warning when using TP that checkpointing is known to not work?

[eitanturok]

@dakinggg I just added a warning that checkpointing does not work + give a link to the exact pytorch issue.

One of the tests verifies that the trainer works but it takes too long cause it downloads a dataset. So I will fix this and I think we will be good to go.