How to use TorchFT?

[lkhphuc]

As far as I understand, TorchFT integration in TorchTitan will make enable HSDP to be fault tolerance, i.e some replica groups can fail, be restarted while the optimizer will keep stepping with the average gradient of the working replicate groups.

From the config, I understand when enable torchft, we would replace training.data_parallel_replicate_degree with fault_tolerance.group_size.
If so for NGPU=8,
run_train.sh --training.data_parallel_replicate_degree=2 --training.data_parallel_shard_degree=4.
is equivalent to as
run_train.sh --fault_tolerance.group_size=2 --training.data_parallel_replicate_degree=1 --training.data_parallel_shard_degree=4

However, running the second command give the error AssertionError: Invalid parallel dims: dp_replicate(1) * dp_shard(4) * cp(1) * tp(1) * pp(1) != WORLD_SIZE(8)

Can anyone help explain how to properly enable torchft with torchtitan? @fegin since you landed this feature.

Thanks in advance.

[fegin]

Will draft a document for using TorchFT.

[fegin]

#993