perf(dllm): use the DeepEP dispatcher in the DiffusionGemma recipes

[akoumpa]

Recipe-only change. Related shared-MoE work lives in #3684 (DeepEP layout reuse + HybridEP warning).

Targets main directly. The HybridEP recipe added by #3651 (diffusion_gemma_te_cp_100k.yaml) is intentionally not touched here — see the last section for why.

What

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. This sets dispatcher: deepep explicitly.

Verified in the CI container that DeepEP is genuinely usable, not just importable:

import deep_ep: OK -> /usr/local/lib/python3.12/dist-packages/deep_ep/__init__.py
HAVE_DEEP_EP: True
DEFAULT BackendConfig -> dispatcher = deepep

…and end-to-end: the patched recipe runs 8 clean steps at 10.47 s/step vs 15.7 s before, TORCHRUN_EXIT=0, 59.17 GiB.

Why

At ep_size=8 the fallback path in moe/experts.py 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.

Profiled on 8×H100. Per step, per rank:

kernel launches	59,482
NCCL collectives	1,184
GPU kernel time in NCCL	97.4 %
exposed (non-overlapped) communication	99.9 %
blocking device→host copies	2,175
compute-kernel time	528 ms
…of which GEMM	73 ms

About 85 % of the collectives are launched from experts.py, and 1,091 of the 1,184 are issued on the same CUDA stream as compute, so overlap is not merely missed — it is impossible.

Measured

eos, 8×H100, 8 steps, shipped recipe with only this change:

dispatcher	step time	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.47× 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.

Why deepep and not hybridep

hybridep was benchmarked as the alternative and is not a drop-in for this recipe. It fails with:

CheckpointError: Recomputed values for the following tensors have different metadata
  saved metadata:      {'shape': torch.Size([481, 2816]),  'dtype': torch.bfloat16}
  recomputed metadata: {'shape': torch.Size([1087, 2816]), 'dtype': torch.bfloat16}

The tokens routed to a rank differ between forward and activation-checkpoint recompute, even though ignore_router_for_ac already defaults to True — so the nondeterminism is beyond what that guard covers. Worth a separate look, since hybridep is what would be needed to scale this recipe past one node (DeepEP's own internode path faults separately at internode.cu:346).

Why diffusion_gemma_te_cp_100k.yaml is untouched

It is deliberately left on dispatcher: torch. It runs ep_size: 1 across two nodes, and the dispatcher branch in moe/layers.py:770-778 keys off get_world_size_safe() rather than ep_size — so changing it would construct GroupedExpertsDeepEP on a 16-GPU internode configuration I have not tested.

Not in this PR

• attn: sdpa and linear: torch are also downgrades from the library default (te); diffusion_gemma_te_cp_100k.yaml already uses attn: te. I have not verified TE attention against the block-diffusion additive mask.
• Raising local_batch_size is worth more (4× the batch costs only ~4 % more step time — the step is dominated by fixed per-collective latency). deepep + local_batch_size: 2 measured 3.08×. It changes the effective global batch, so it belongs in a separate, deliberate change.

---

Update: making both dispatchers usable

Follow-up investigation into why hybridep could not simply be swapped in.

Working configurations (8xH100, measured)

dispatcher	activation ckpt	step time	result
torch (before)	on	15.7 s	works
deepep	on	10.6 s	works — what this PR ships
deepep	off	5.83 s	works
hybridep	off	5.46 s	works — fastest measured
hybridep	on	—	fails: CheckpointError

Why hybridep + activation checkpointing fails

It dies in the first backward with:

CheckpointError: Recomputed values ... have different metadata
  saved:      torch.Size([2791, 2816])
  recomputed: torch.Size([2701, 2816])

That message points at the router — a shape change like this is documented as non-deterministic
re-routing, which ignore_router_for_ac exists to prevent. That is not the cause here.

Instrumenting Gemma4Gate (DiffusionGemma swaps in its own gate, so the standard Gate is never
called) shows routing is reproduced exactly. Encoder pass vs its recompute, top-k index checksum:

rank 0: 282623 == 282623     rank 4: 425381 == 425381
rank 1: 275769 == 275769     rank 5: 279545 == 279545
rank 2: 276947 == 276947     rank 6: 278018 == 278018
rank 3: 273901 == 273901     rank 7: 262048 == 262048

All 30 decoder layers match too (0/30 mismatches on every rank). With bit-identical routing on all
8 ranks, hybrid_ep_dispatch returned 2791 rows in the forward and 2701 on replay. The
non-determinism is inside the dispatch, so no router-side mitigation can prevent it.

Also ruled out along the way:

• Shared dispatch state — dispatcher_share_token_dispatcher=False (per-layer managers) still fails.
• pad_multiple — set once at construction, never reassigned.
• Adding the EP collectives to the AC save-list — the torch.ops.deepep.dispatch /
  hybridep.dispatch entries already in activation_checkpointing.py are inert: dispatch goes
  through FusedDispatch.apply, an autograd.Function, which a __torch_dispatch__ policy never
  observes. Those four entries have never matched anything. Worth fixing separately.

A warning for this combination, and the DeepEP dispatch-layout reuse
found while profiling this, are both in #3684. This PR is only the two recipe files.

[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.