[FT] Support local_sgd / diloco in titan

[H-Huang]

Depends on torchft changes:

• Support DTensor params in local_sgd/diloco torchft#168
• Update DiLoCo backup device and docs torchft#170

This PR adds a new "FtTrainer" to the torchft experiments folder as well as a context manager which wraps around the train loop to run local sgd or diloco. It also adds a config property to set the training method.

To run (need 3 different terminals):

Start torchft lighthouse (terminal 1):

RUST_LOGS=debug RUST_BACKTRACE=1 torchft_lighthouse --min_replicas 2 --quorum_tick_ms 100 --join_timeout_ms 10000

Start replica 1 (terminal 2, update lighthouse URL):

TORCHFT_LIGHTHOUSE=<url> TORCHFT_MANAGER_PORT=29520 REPLICA_GROUP_ID=0 CUDA_VISIBLE_DEVICES=0,1,2,3 NGPU=4 ./run_train.sh --parallelism.data_parallel_shard_degree=4 --fault_tolerance.enable --fault_tolerance.group_size=2 --fault_tolerance.replica_id=0

Start replica 2 (terminal 3, update lighthouse URL):

TORCHFT_LIGHTHOUSE=<url> TORCHFT_MANAGER_PORT=29522 REPLICA_GROUP_ID=1 CUDA_VISIBLE_DEVICES=4,5,6,7 NGPU=4 ./run_train.sh --parallelism.data_parallel_shard_degree=4 --fault_tolerance.enable --fault_tolerance.group_size=2 --fault_tolerance.replica_id=1

[fegin]

It looks like we can share the training loop? We basically just add another context manager. Additionally, the training_method is added to the main JobConfig. so we shouldn't create an additional train.py.

If we don't want to expose this feature as the main feature yet, then we have to move training_method configuration to experiments/fault_tolerance/train.py as well. You can check torchtitan/tests/unit_tests/test_job_config.py to find out the example.

[fegin]

It seems that we are not using this variable. We don't need as semi_sync_training?

[H-Huang]

Yeah we dont need it. I can remove

[fegin]

semi_sync_method or synchronize_method are less confusing. trianing_method is pretty general. And can we specify that if the value is not set, we will use synchronized training?

[H-Huang]

semi_sync_method sounds good. I will also update that comment

[H-Huang]

Correct @fegin, we can share the training loop. My main worry was making the regular training loop harder to read by adding this, but yeah you are right that the configs are already in JobConfig.

The alternative is to ask the user to use torchtitan.experiments.fault_tolerance import FtTrainer and replace their trainer = FtTrainer(config), but im not sure if we currently support swapping in another trainer. Maybe your proposal is the right way to go

[d4l3k]

I think if we can do it with the same train loop it'd be nice -- maybe we can use ExitStack for cleaner optional registration of the context managers?

[fegin]

I'm okay either way. Since TorchFT is already in the core TorchTitan, I think it is okay to put the semi-sync to the main training loop.

If you want to put it in the experiment folder, custom_args_module is all you need.