[cp] apply fsdp to model when CP is enabled without DP for correct loss and lower mem usage

[XilunWu]

Stack from ghstack (oldest at bottom):

• [cp] add option to choose kv shards rotation method #684
• -> [cp] apply fsdp to model when CP is enabled without DP for correct loss and lower mem usage #685

Summary
Previously CP forgot to shard the model via apply_fsdp when DP is not combined with CP. This leads to high peak memory usage and diverging loss.

Test

1. modify train_configs/llama3_8b.toml

steps = 20
context_parallel_degree = 8

2. run training on 8xH100 GPUs
   CONFIG_FILE="./train_configs/llama3_8b.toml" NGPU=8 LOG_RANK=0,1,2,3,4,5,6,7 ./run_llama_train.sh
   Before: CUDA OutOfMemory
   After: successful 20-steps training

[tianyu-l]

I wonder if we should take a more systematic approach for creating & naming of sub-meshes in parallel_dims.py:
There are two layers sub-meshes: atomic ones and derived ones

1. The atomic ones include those at minimal granularity: dp_shard, dp_replicate, tp, pp, cp.
2. The derived ones include: dp (for data loading), dpshard_cp (for FSDP param sharding).

We can create / skip the atomic ones first, and then always create the derived ones. E.g. if dp_shard enabled, dp_replicate not enabled, cp not enabled, then we still create dp and dpshard_cp by flattening dp_shard alone.

This way the code would be simpler and more readable. The creation in parallel_dims would also be less confusing.

[fegin]

After looking into the table, I feel we should use dp_shard_cp for fully_shard purpose and leave dp for the data loader. This will be more consistent.

I also think @tianyu-l comment is a good approach. But I'm not sure if this will become more complicated when MoE is involved.

[tianyu-l]

But I'm not sure if this will become more complicated when MoE is involved.

@fegin
I feel the idea could carry over to any new parallelisms -- always initialize the lowest level dimensions, and then gradually build up, with carefully chosen names though.

[tianyu-l]

lgtm!