enable Context Parallel

[XilunWu]

Stack from ghstack (oldest at bottom):

• -> enable Context Parallel #592

Summary
This PR adds DTensor Context Parallel (pytorch/pytorch#131351) option to torchtitan.

Instructions
Find the context_parallel_degree option in the experimental section in .toml file, assign the desired Context Parallel degree (i.e. the # of GPUs along the Context Parallel dimension).

Benchmark Settings
The benchmark uses train_configs/llama3_8b.toml with extra changes as below:

• norm_type = "fused_rmsnorm"
• mode = 'full' (using full activation checkpoint allows us to find the longest possible input sequence length on a device)

Benchmark Results

1. max seq_len scalability
   The first set of benchmark shows that with CP, adding more GPUs on the CP dimension allows longer input seq_len and this growth is linearly scalable. We use FSDP to shard the model and fix data_parallel_shard_degree to be 8, and we use 8, 16, 32, 64 H100 GPUs (i.e. CP degree=1,2,4,8) to measure the longest possible input seq_len w/o GPU OOM. We adjust the CP degree by modifying the context_parallel_degree in the experimental section in .toml file.

As Context Parallel degrees goes up (1, 2, 4, 8), the max seq_len also increases (32k, 80k, 144k, 300k)

While increasing Context Parallel degree increases the max seq_len, it also decreases the WPS on each device. This is the price paid for longer input

However, increasing CP degree doesn't significantly affect MFU.

2. loss curve convergence
   The second set of benchmark proves that our implementation is correct by showing the loss curve converges just as FSDP + TP. This experiment uses 64 H100 GPUs and fix TP degree and local batch size to 8 (tensor_parallel_degree=8, batch_size=8). As we increase CP degree, we decrease DP degree to make sure that CP * DP = 8. Besides, we also make sure seq_len increases such that seq_len = 8192 * CP.

The training loss curve matches among the 4 combination of parallelisms: DP=8; DP=4, CP=2; DP=2, CP=4; CP=8 (TP is fixed to 8).

The only observed difference is among the warm-up steps. This is due to our pick of the number of steps and picking an appropriate number of warm-up steps should smooth out this.

3. max seq_len vs. local WPS on a fixed number of GPUs
   The third set of benchmark gives users a guide on what CP degree to choose and what WPS to expect on a specific set of devices (e.g. number of devices, device types, etc). Here we use H100 as an example but it is expected to observer similar trade-off curves on other hardware.

Increasing CP degree allows longer input sequence length at the cost of reducing WPS on each device. The exception is pure CP -- the seq_len increase is very limited compared to DP=2, CP=4 on 8 GPUs. This is due to our implementation, and we're working on enabling performant CP to support ultra-long input sequence.

TODO:

1. An option for CP tensor rotate implementation between all-to-all and all-gather will be added once [CP] Implement AllGather based context parallelism pytorch#132820 is landed.
2. Investigate the limitation in pure CP, to enable ultra-long input sequence.

[fegin]

cc., @tianyu-l Want to understand is this okay?

For a general use case, we can also expand the CP to support stride-like feature.

[tianyu-l]

Could you please elaborate a bit on why this change was needed by CP?

[fegin]

@tianyu-l CP parallelize on the sequence dimension, anything related to the sequence dimension needs to be shard. So freqs_cis is the positional embedding and is required to be sharded according to the sequence length. So it is easier to support CP if everything has the same sequence length.

[tianyu-l]

Sounds reasonable to me. @awgu to confirm this is OK.

Also we need to add a note in docs/composability.md to clarify why this (model change) is needed. It can be addressed in a separate PR; in that case please create issue / leave TODO.

[tianyu-l]

Looks awesome! Still have a few questions, but I think it's mostly ready.

[tianyu-l]

Question: Looks like FSDP/HSDP + TP + CP is working. How about PP? We can also mention progress in the .md doc later.

[XilunWu]

Right, the next step is to test 4D/5D (w/ PP and HSDP)

[tianyu-l]

_flatten seems called twice, another one in parallel_dims.py. I wonder how this API works?

[tianyu-l]

another question is: given the current implementation, does it mean DP and CP have to be adjacent to each other in the device mesh? E.g. it seems we can't do [TP, CP, PP, DP] (from inner to outer) as in Llama 3.1 paper. Is that correct?

[XilunWu]

@awgu thanks for catching it! I forgot to paste the 5000-step loss curve screenshot but a zoom-in one for warm-up steps. Updated.

[tianyu-l]

While increasing Context Parallel degree increases the max seq_len, it also decreases the WPS on each device. This is the price paid for longer input. However, increasing CP degree doesn't significantly affect MFU.

By looking into how token/sec and MFU change when increasing seq_len, here's my explanation to this phenomenon.

There are two parts of flops computation
https://github.com/pytorch/torchtitan/blob/main/torchtitan/utils.py#L150

• (I) matmul scales linearly with number of tokens
• (II) attention scales quadratically with number of tokens

When doubling seq_len together with CP degree (also doubling number of GPUs since other degrees are fixed)

• (I) matmul computation doubles; (II) attention computation quadruples.
• Even CP doesn’t bring any overhead (e.g. extra comm exposure), we shoudn’t expect the same WPS per GPU because computation resource is doubled, but computation load is between 2x – 4x. The throughput per GPU should be between 0.5x and 1x, assuming perfect scaling.
• The MFU change depends on the actual config, more specifically between the ratio of (I) and (II). If (I) dominates, the MFU/WPS ratio wouldn’t change as we increase seq_len; if (II) dominates, the MFU/WPS ratio would increase linearly with seq_len. For 8B and seq_len around 32k/64k, (II) is larger than (I) so that’s why we are seeing WPS drop, but MFU stays (because the ratio increases).

My conclusion is that the perf scaling we observed makes sense and look good to me.

[tianyu-l]

LGTM! Thanks!

[XilunWu]

@tianyu-l actually, the 8 GPU test has a failure. This is the HSDP + CP case. I think the flatten logic has some bug when handling mesh["dp_cp"].

Flattening mesh["dp_replicate", "dp_shard", "cp"] into mesh["dp_cp"] is a workaround w/o actual cost. I think we can land this first with a TODO and put the __get_item__ call back once it's fixed in DeviceMesh.

[tmm1]

Just curious if you recall what the blocker for CUDNN_ATTENTION is

[tianyu-l]

@XilunWu

[XilunWu]

Hi @tmm1 it's simply cudnn attention has a different op signature. I'm adding support now and should be able to have the PR draft out by next week.