[ISEEKYAN]

Summary

Add opt-in token-wise ChunkedEP for Qwen3 MoE, using DeepEP to overlap communication and expert computation.

• Three primitives: forward, backward, and fused forward/backward. Operators do not inspect recomputation or optimizer settings.
• Normal training uses forward with saved context plus backward; full-layer recomputation composes graph-free forward with fused forward/backward.
• Caller-owned expert output/dgrad buffers and an event-guarded activation arena shared across matching layers. Capacity grows lazily; cleanup uses a generic model-offload callback.
• Native EP remains a parallel path. Reuse native dispatcher/expert setup; keep Qwen3 integration small.
• No new head-loss/linear CE implementation, environment variables, or experimental/profiling scaffolding. Existing linear CE remains independently configured.

Candidate: 45634901b5a80ed350a2b9d456cde7be2c0292a2, 15 files, +4000/-15 relative to dev bb5dfd08f09ce06c5925af453fef06b3129f199d.

Validation and performance

Actual optimizer updates: 48 layers, 16K local tokens

Qwen3 BF16 THD, 16 H100 GPUs across two nodes, 48 layers, 16384 tokens/GPU, 16 microbatches, EP8/top-k8, expert DP2, two chunks, full recomputation, random initial weights, identical existing linear CE, lazy activation capacity. Both arms execute real precision-aware distributed FusedAdam and scheduler updates. Three warmup and five measured steps; no profiler or private allocator-retention modification.

Metric	Native	ChunkedEP
Median distributed step time	70.340 s	61.933 s
Speedup	1.000x	1.136x
Worst-rank allocated peak reduction	—	6.15%
Worst-rank reserved peak reduction	—	1.45%

Step latency takes the maximum across ranks before taking the median across measured steps. Both arms completed all eight updates on all ranks, with zero allocator retries/OOMs. Results were measured at 83f7552e4; current production code differs only in RuntimeError messages. Subsequent changes consolidate test setup and documentation, without changing operator computation.

After timing, all ChunkedEP ranks completed model/gradient CPU offload, CUDA reload and a forward/backward recovery step with finite loss and zero allocator retries/OOMs. This does not qualify recovery-gradient equivalence or a recovery optimizer update.

Historical single-layer OP measurements

At a4f2e2b4a, 32768 local tokens, EP8/top-k8, two chunks, five warmup and thirty measured repeats:

OP	Native → ChunkedEP	Speedup	Allocated peak reduction
Forward	16.563 → 15.005 ms	1.104x	17.38%
Fused forward/backward	40.125 → 27.946 ms	1.436x	30.03%

These predate the router-gradient correction, exclude the model head, and are not whole-model or activation-only memory results.

Historical whole-model runs at the same older source used eight GPUs, full recomputation, identical existing linear CE, no optimizer updates, three warmup and ten measured steps. One layer/32768 tokens/one microbatch achieved 1.0263–1.0283x with allocated peak reduction 12.67–16.11%; 48 layers/16384 tokens/16 microbatches achieved 1.1302x with allocated peak reduction 5.48–9.70%. Ranges cover ranks, not confidence intervals. These must not be conflated with the 16-GPU actual-update result above.

Numerical status — unresolved

This is a Draft for external evaluation, not a claim of accepted training equivalence.

Router weight gradients now accumulate across chunks in FP32 before a single dtype conversion. A separate single-layer 32K-token full-gradient check at 83f7552e4 compared 1,245,452,544 finalized optimizer-owned gradient elements: maximum absolute error 1.1921e-7; router relative L2 error 0.006774%, improved from the earlier 0.2871% diagnostic. It is not bitwise parity, and no non-bitwise tolerance has been approved.

In the 48-layer actual-update run, maximum loss absolute difference was 2.4224e-4. Gradient-norm relative difference reached 9.30% after eight updates, versus 0.00685% on the initial step. The later comparison is at independently updated weights; it is not a same-weight full-gradient error. The cause remains under investigation, including a native/native repeatability baseline. Small loss differences do not establish training equivalence.

Tests and limitations

• 89 related CPU contract tests passed; CPU test doubles do not prove GPU equivalence. Selected Ruff 0.9.10 and whitespace checks passed.
• Qwen3 only; MTP is rejected. Two chunks are the measured target. Final-source n=3/4 GPU qualification is incomplete.
• Normal backward is supported, but normal-mode speed/memory improvements are not qualified.
• CUDA graph capture/replay is not qualified. Lazy capacity and eager parking do not imply zero allocations or stable eager addresses.
• Full-recompute composition replaces the outer recompute list. Ordinary ChunkedEP rejects outer MoE/full checkpoint composition; caller-owned outputs do not support moe_act, and router auxiliary loss is disabled on this path.
• Full delivery review and numerical qualification remain in progress. No merge approval is implied.

Usage: experimental/lite/examples/bench/chunked_ep.md.
Fork review: ISEEKYAN#225.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.