separate out training for fault tolerance

[tushar00jain]

Summary:

• extract logic in train.py that was different for ft in separate functions
• override these functions in ft's train.py

Test Plan:
Run lighthouse and two replicas

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

TRAIN_FILE=torchtitan.experiments.ft.train CONFIG_FILE=./torchtitan/models/llama3_ft/train_configs/debug_model.toml CUDA_VISIBLE_DEVICES=0 NGPU=1 ./run_train.sh --parallelism.data_parallel_shard_degree=1 --fault_tolerance.enable --fault_tolerance.group_size=2 --fault_tolerance.replica_id=0


TRAIN_FILE=torchtitan.experiments.ft.train CONFIG_FILE=./torchtitan/models/llama3_ft/train_configs/debug_model.toml CUDA_VISIBLE_DEVICES=1 NGPU=1 ./run_train.sh --parallelism.data_parallel_shard_degree=1 --fault_tolerance.enable --fault_tolerance.group_size=2 --fault_tolerance.replica_id=1

[tianyu-l]

I feel the whole point is for ft trainer not to influence how main trainer API looks like. These functions are there purely because ft trainer can override, which sounds the opposite.

The point is that, ft and main cannot share all the code, so they probably shouldn't. E.g. the Flux trainer has its own logic. I feel the ft trainer can just rewrite, instead of trying to share as much as possible.

[tushar00jain]

batch_degree is being used a lot which will require pretty much just copy pasting all off init.py. lmk if you'd still prefer that

[fegin]

IMO, whether a change makes sense or not depending on whether the function is too intrusive. This doesn't make sense because this function is not needed for the regular trainer. I'm okay if this method contains the original logic in init():

        def get_batch_dim_info(parallel_dims):
            if parallel_dims.dp_enabled:
                batch_mesh = parallel_dims.get_mesh("batch")
                batch_degree, batch_rank = batch_mesh.size(), batch_mesh.get_local_rank()
            else:
                batch_degree, batch_rank = 1, 0
            return batch_degree, batch_rank

This is just my thought.

[fegin]

IMO, whether a change makes sense or not depending on whether the function is too intrusive. This doesn't make sense because this function is not needed for the regular trainer. I'm okay if this method contains the original logic in init():

        def get_batch_dim_info(parallel_dims):
            if parallel_dims.dp_enabled:
                batch_mesh = parallel_dims.get_mesh("batch")
                batch_degree, batch_rank = batch_mesh.size(), batch_mesh.get_local_rank()
            else:
                batch_degree, batch_rank = 1, 0
            return batch_degree, batch_rank

This is just my thought.

[fegin]

Overall, this PR introduces two extra methods, compute_global_losses and get_dp_info while removing all the FT from the main trainer. I think this is a balanced compromise. @tianyu-l any thought?

[fegin]

Forgot to delete?