|
| 1 | +# Dynamic Context Parallel |
| 2 | + |
| 3 | +Dynamic context parallelism (dynamic CP) runs every packed microbatch on a |
| 4 | +context-parallel group sized for its longest sequence, instead of splitting |
| 5 | +every microbatch across the full `context_parallel_size` group. It builds on |
| 6 | +[global-batch online packing](packed-sequences.md#global-batch-online-packing): |
| 7 | +the packing scheduler already sees the whole global batch across data-parallel |
| 8 | +ranks, and with dynamic CP it also picks a context-parallel group for each |
| 9 | +packed bin from the pool of `data_parallel_size x context_parallel_size` ranks. |
| 10 | + |
| 11 | +Dynamic CP needs the Megatron-Core `dev` submodule pin: |
| 12 | + |
| 13 | +```bash |
| 14 | +./scripts/switch_mcore.sh dev && uv sync |
| 15 | +``` |
| 16 | + |
| 17 | +`ConfigContainer.validate` reports what an older pin is missing. |
| 18 | + |
| 19 | +## When to Use It |
| 20 | + |
| 21 | +Static context parallelism is sized for the longest sequence a run must fit. |
| 22 | +On a variable-length stream most sequences are much shorter, and splitting them |
| 23 | +across the whole group adds context-parallel communication and smaller kernels |
| 24 | +without needing the memory. Dynamic CP pays off when both of these hold: |
| 25 | + |
| 26 | +- Static CP is required by memory for the long tail. A large vocabulary can |
| 27 | + make the logits buffer the limit, for example. |
| 28 | +- The stream has a meaningful share of sequences short enough to run on one GPU |
| 29 | + or a small group. |
| 30 | + |
| 31 | +When nearly every sequence needs the full group, dynamic CP matches static CP up |
| 32 | +to a small scheduling overhead. |
| 33 | + |
| 34 | +## How It Works |
| 35 | + |
| 36 | +`model.max_seqlen_per_dp_cp_rank` is the token capacity of one rank in one |
| 37 | +microbatch. With `default_dynamic_cp`, a sequence of at most that length runs on |
| 38 | +one GPU, a sequence of up to twice that length on a two-GPU group, and so on up |
| 39 | +to the whole pool. The scheduler forms bins longest first and gives each bin a |
| 40 | +power-of-two group, and every rank runs the same number of microbatches in a |
| 41 | +step. |
| 42 | + |
| 43 | +Bridge wires this through the global-batch packing path: |
| 44 | + |
| 45 | +- `model.dynamic_context_parallel=True` makes validation select the |
| 46 | + `default_dynamic_cp` scheduler instead of `dp_balanced`. |
| 47 | +- The packed batch fetch receives the dynamic-CP flag and returns a |
| 48 | + `PackedSeqParams` that carries the microbatch's context-parallel group, which |
| 49 | + attention and rotary embeddings use in place of the static group. |
| 50 | +- The per-sequence alignment multiple covers every group the pool can form: |
| 51 | + `2 x dp x cp`, and also `dp x cp x tp` under sequence parallelism. |
| 52 | + |
| 53 | +## Configuration |
| 54 | + |
| 55 | +```python |
| 56 | +cfg.dataset.enable_global_batch_packing = True |
| 57 | +cfg.dataset.dataloader_type = "single" |
| 58 | +cfg.model.context_parallel_size = 4 # static CP domain inside each data-parallel replica |
| 59 | +cfg.model.dynamic_context_parallel = True |
| 60 | +cfg.model.min_dynamic_context_parallel_size = 1 |
| 61 | +cfg.model.max_seqlen_per_dp_cp_rank = cfg.model.seq_length // 4 |
| 62 | +cfg.model.calculate_per_token_loss = True |
| 63 | +cfg.ddp.average_in_collective = False |
| 64 | +cfg.train.micro_batch_size = 1 |
| 65 | +``` |
| 66 | + |
| 67 | +Validation applies every global-batch packing rule and adds these: |
| 68 | + |
| 69 | +- The `dp x cp` pool is a power of two with at least two ranks. |
| 70 | +- `min_dynamic_context_parallel_size` is a power of two no larger than the pool. |
| 71 | +- `dist.use_decentralized_pg=False`, because Megatron-Core resolves dynamic |
| 72 | + groups from its global parallel state. |
| 73 | +- The pool times `max_seqlen_per_dp_cp_rank` covers `seq_length`. |
| 74 | + |
| 75 | +## Reference Recipe |
| 76 | + |
| 77 | +`qwen35_text_35b_a3b_sft_8gpu_gb200_bf16_dynamic_cp_config` fine-tunes |
| 78 | +Qwen3.5-35B-A3B on CoderForge-Preview SWE-Rebench trajectories with TP1, PP1, |
| 79 | +CP4, and EP8 on eight GB200 GPUs and a 32K-token cap. |
| 80 | +`examples/training_features/long_context/qwen35_35b_a3b_dynamic_cp.py` launches |
| 81 | +it. Add these overrides for the static-CP baseline on the same packing path: |
| 82 | + |
| 83 | +```bash |
| 84 | +model.dynamic_context_parallel=false model.sequence_packing_scheduler=dp_balanced |
| 85 | +``` |
| 86 | + |
| 87 | +On the recipe's own data, about nine in ten CoderForge trajectories exceed the |
| 88 | +32K cap, so nearly every packed sequence needs the full CP4 group. There, |
| 89 | +dynamic CP runs within noise of static CP4: a median step of 88 s against |
| 90 | +83 s, at about 90 model TFLOP/s per GPU. The example's README lists the |
| 91 | +measurements, including the effect of loader workers and recompute choices. |
| 92 | +Expect gains on streams where a meaningful share of sequences is short enough |
| 93 | +to run on one GPU or a small group. |
| 94 | + |
| 95 | +## Models with Wide Attention Heads |
| 96 | + |
| 97 | +Transformer Engine does not train with cuDNN fused attention at head dimension |
| 98 | +256 on Blackwell, so Qwen3.5 attention runs on FlashAttention. FlashAttention |
| 99 | +rejects THD bins that contain padding between sequences. Such models need two |
| 100 | +changes: |
| 101 | + |
| 102 | +- `fold_alignment_padding=True` on the dataset config, which reports each |
| 103 | + sequence's alignment padding as part of that sequence. The padding stays |
| 104 | + loss-masked and follows the sequence's last token, so causal attention and |
| 105 | + recurrent layers compute the same outputs for real tokens. |
| 106 | +- A Megatron-Core that derives `pad_between_seqs` from the actual `cu_seqlens`, |
| 107 | + proposed in [Megatron-LM #7417](https://github.com/NVIDIA/Megatron-LM/pull/7417). |
| 108 | + |
| 109 | +Leave `fold_alignment_padding` off for head dimensions of 128 or less, where |
| 110 | +cuDNN fused attention handles the padding itself. |
| 111 | + |
| 112 | +## Practical Caveats |
| 113 | + |
| 114 | +- The first step that uses a new group size pays communicator and kernel setup, |
| 115 | + so expect a few slow early steps. |
| 116 | +- The scheduler may place short sequences on larger groups to keep every rank |
| 117 | + busy, so fewer sequences run on one GPU than the length histogram suggests. |
| 118 | +- Transformer Engine's FusedAdam draws parameter handles from a pool of roughly |
| 119 | + 20,000. Many experts on few expert-parallel ranks can exhaust it with the |
| 120 | + precision-aware optimizer. Use more expert-parallel ranks, as the recipe does |
| 121 | + with EP8, or offload part of the optimizer to the CPU. |
| 122 | +- The fused cross-entropy materializes an fp32 buffer of tokens per rank times |
| 123 | + vocabulary size. At a 248K vocabulary this buffer sets |
| 124 | + `max_seqlen_per_dp_cp_rank` and makes context parallelism necessary at long |
| 125 | + caps. |
| 126 | + |
| 127 | +## Related Docs |
| 128 | + |
| 129 | +- [Packed Sequences](packed-sequences.md) |
| 130 | +- [Hierarchical Context Parallel](hierarchical-context-parallel.md) |
| 131 | +- [Activation Recomputation](activation-recomputation.md) |
0 commit comments