feat(dllm): add Transformer Engine context parallelism for DiffusionGemma - #3651
feat(dllm): add Transformer Engine context parallelism for DiffusionGemma#3651akoumpa wants to merge 7 commits into
Conversation
Signed-off-by: Alexandros Koumparoulis <akoumparouli@nvidia.com>
|
Multinode Transformer Engine loss parity completed on cw-dfw (Slurm
The small BF16 differences are consistent with the different distributed attention/reduction order; the loss is numerically equivalent.
|
Red-team loss validationStatus: the TE context-parallel loss and gradients pass an independent oracle; one pre-existing objective-reduction caveat remains. Independent checks:
The diffusion term matches the documented objective: plain CE over every supervised token in the selected canvas, including uncorrupted tokens, with no inverse-noise weighting. Since selected canvases are fixed-width after EOS fill, token reduction and per-example reduction coincide for this term. One caveat is unrelated to CP: the encoder AR term currently reduces as sum(CE) / sum(valid tokens), while Eq. 13 specifies a per-example sequence mean followed by a batch mean. These are identical for equal-length examples but weight longer examples more heavily in variable-length batches. An exaggerated 11/5/2-token oracle case differed by 7.67%. CP preserves the implemented token-weighted objective correctly, but we should not describe the encoder term as exact Eq. 13 unless that weighting is an explicit design choice. |
Signed-off-by: Alexandros Koumparoulis <akoumparouli@nvidia.com>
|
/ok to test 4c50835 |
Encoder-loss equivalence follow-upThe variable-length reduction caveat is resolved in The encoder objective is now:
Under CP, every rank divides its local numerator by the replicated full-sequence pair count. The recipe uses the DP-global nonempty-example count across gradient-accumulation microbatches, and the existing DP x CP backward scaling combines the disjoint token shards. Independent validation:
Matched W&B rerun (cw-dfw Slurm
|
100K-token CP training curveThe 100K validation now has a real multi-step W&B history rather than a one-point snapshot:
This uses the same one-layer tiny DiffusionGemma and real model/TE-CP/fused-CE path as the 100K correctness probe. It performs synchronized AdamW updates: CP-local loss-shard gradients are summed across all four ranks before every optimizer step. The plotted objective is the encoder AR loss; the diffusion-canvas loss is not included in this long-context curve. Remote log: |
Signed-off-by: Alexandros Koumparoulis <akoumparouli@nvidia.com>
|
Updated every dLLM diffusion training recipe to retain FP32 parameters ( Validation:
A focused regression test now prevents future diffusion recipes from reverting the master-parameter precision contract. |
|
/ok to test 747f423 |
diffusion_gemma_sft.yaml and diffusion_gemma_lora.yaml pinned `dispatcher: torch`. Both run ep_size=8 on a single node, so this forced the DTensor fallback even on DeepEP-capable images such as the CI container. At ep_size=8 that fallback runs, for every MoE layer on every forward pass: a token-count all_gather, an `[int(t.item()) ...]` host sync, and four variable-length all_gathers (activations, router weights, indices, token mask) that replicate every token to all 8 ranks, then an all-reduce to combine. DiffusionGemma runs the shared stack three times per step (causal encoder, no_grad self-conditioning decode, real decode) and activation-checkpoint recompute replays the dispatch again in backward. Measured with ntrace on 8xH100 (google/diffusiongemma-26B-A4B-it), per step per rank: 59,482 kernel launches, 1,184 NCCL collectives holding 97.4% of GPU kernel time, 99.9% of that communication exposed (1,091 of the 1,184 are issued on the same CUDA stream as compute, so overlap is impossible), 2,175 blocking device-to-host copies, and only 528 ms of compute-kernel time of which 73 ms is GEMM. Benchmarked on eos, 8xH100, 8 steps, shipped recipe with only this change: | dispatcher | step | throughput | peak mem | |-----------------|---------|-------------|----------| | torch (before) | 15.7 s | 263 tok/s | 59.0 GiB | | deepep (after) | 10.6 s | 388 tok/s | 59.2 GiB | 1.47x faster steps at equal memory. Loss trajectories match to within floating-point reassociation (step 0: 4.9946 vs 4.9583), as expected for a change that alters only how tokens are exchanged. Set explicitly rather than left unset: the value is visible in the recipe, and BackendConfig's default would otherwise resolve differently per image. "hybridep" was benchmarked as the alternative and is not a drop-in for this recipe -- it fails with CheckpointError because the tokens routed to a rank differ between forward and activation-checkpoint recompute (481 vs 1087 rows), even though ignore_router_for_ac already defaults to True. The diffusion_gemma_te_cp_100k.yaml recipe added by #3651 is deliberately left on `dispatcher: torch`: it runs ep_size=1 across two nodes, and the dispatcher branch in moe/layers.py keys off world size rather than ep_size, so switching it would activate the DeepEP path on an untested internode configuration. Signed-off-by: Alexandros Koumparoulis <akoumparouli@nvidia.com>
diffusion_gemma_sft.yaml and diffusion_gemma_lora.yaml pinned `dispatcher: torch`. Both run ep_size=8 on a single node, so this forced the DTensor fallback even on DeepEP-capable images such as the CI container. At ep_size=8 that fallback runs, for every MoE layer on every forward pass: a token-count all_gather, an `[int(t.item()) ...]` host sync, and four variable-length all_gathers (activations, router weights, indices, token mask) that replicate every token to all 8 ranks, then an all-reduce to combine. DiffusionGemma runs the shared stack three times per step (causal encoder, no_grad self-conditioning decode, real decode) and activation-checkpoint recompute replays the dispatch again in backward. Measured with GPU kernel tracing on 8xH100 (google/diffusiongemma-26B-A4B-it), per step per rank: 59,482 kernel launches, 1,184 NCCL collectives holding 97.4% of GPU kernel time, 99.9% of that communication exposed (1,091 of the 1,184 are issued on the same CUDA stream as compute, so overlap is impossible), 2,175 blocking device-to-host copies, and only 528 ms of compute-kernel time of which 73 ms is GEMM. Benchmarked on eos, 8xH100, 8 steps, shipped recipe with only this change: | dispatcher | step | throughput | peak mem | |-----------------|---------|-------------|----------| | torch (before) | 15.7 s | 263 tok/s | 59.0 GiB | | deepep (after) | 10.6 s | 388 tok/s | 59.2 GiB | 1.47x faster steps at equal memory. Loss trajectories match to within floating-point reassociation (step 0: 4.9946 vs 4.9583), as expected for a change that alters only how tokens are exchanged. Set explicitly rather than left unset: the value is visible in the recipe, and BackendConfig's default would otherwise resolve differently per image. "hybridep" was benchmarked as the alternative and is not a drop-in for this recipe -- it fails with CheckpointError because the tokens routed to a rank differ between forward and activation-checkpoint recompute (481 vs 1087 rows), even though ignore_router_for_ac already defaults to True. The diffusion_gemma_te_cp_100k.yaml recipe added by #3651 is deliberately left on `dispatcher: torch`: it runs ep_size=1 across two nodes, and the dispatcher branch in moe/layers.py keys off world size rather than ep_size, so switching it would activate the DeepEP path on an untested internode configuration. Signed-off-by: Alexandros Koumparoulis <akoumparouli@nvidia.com>
Signed-off-by: Alexandros Koumparoulis <akoumparouli@nvidia.com>
|
Fixed the CP=8 decoder NaN at exact head TE P2P context parallelism evaluates K/V blocks incrementally, so an intermediate block can be fully masked even when the complete attention row is valid. Using the dtype minimum for that block produced a non-finite partial softmax. The additive mask now uses Validation on cw-dfw:
I am now running the final matched 10-step CP=1/CP=8 W&B comparison at this head. |
|
/ok to test c4dbf8b |
|
Final exact-head cw-dfw validation completed for
Both runs are finished and each W&B history contains 10 loss rows. CP=1 vs CP=8 total-loss parity:
AR maximum relative delta is The runtime dtype audit also verifies FP32 master/resident parameters do not make the whole model run in FP32. In both runs:
So FP32 is retained for master parameter storage while the expensive linear, TE attention, and fused CE compute remains BF16. |
1 participant
Summary
Loss semantics
The diffusion decoder remains non-causal. Its exact block-attention mask is supplied as a post-scale attention bias while Transformer Engine performs context-parallel attention. Padding and dummy cache positions are excluded from the objective.
The encoder AR term first averages valid next-token CE within each example, then averages those per-example means over the global batch. Each CP rank contributes only its local token numerator while using the replicated full-sequence token count. The existing DP x CP backward scaling combines those disjoint contributions into the same objective as the non-CP path. Empty examples are excluded from the denominator.
Validation
16555164: two nodes, four H100s, Transformer Engine CP=4, completed successfully[17, 7]: loss 6.396856 versus 6.406250 and LM-head gradient cosine 0.99997294