[bug] MoE + expert parallel + `gradient_checkpoint full` fails at step 0: CheckpointError from HybridEP dispatch token-count drift

Problem

Any MoE model trained with --fsdp.ep_size > 1 and the hybridep dispatcher
fails on the first backward pass with:

torch.utils.checkpoint.CheckpointError: torch.utils.checkpoint: Recomputed values
for the following tensors have different metadata than during the forward pass.

The mismatched dimension is the permuted token count after expert dispatch —
forward and recompute disagree by a few tokens, so the grouped-GEMM activations
change shape between the two passes.

This is the default configuration of every shipped flagship MoE recipe, so those
recipes cannot run as written:

recipe	ep_size	AC
examples/scripts/slurm/sft_qwen3_6_35b.sh:57,116	8	full
examples/scripts/quick_start/sft_qwen3_6_35b.sh:43,78	8	full
examples/scripts/slurm/sft_omni3_30b.sh:62,76	8	full
examples/scripts/slurm/rl_qwen3_6_35b.sh:42,47	8	full
examples/scripts/slurm/rl_glm5_2_753b.sh:42,45	256	full

The existing mitigation does not prevent it

molt/trainer/fsdp/strategy.py:242-262 sets ignore_router_for_ac=True
specifically to avoid this error, on the theory that recomputing the router
re-routes near-tie tokens and shifts per-expert counts.

Instrumentation shows that mitigation works as designed and is still not
sufficient:

1. The router matcher fires. _is_router_projection in
   nemo_automodel/components/moe/parallelizer.py:459-477 does match the router
   projection, which is therefore marked MUST_SAVE:

   [DBG-AC] want hidden=256 experts=8 | aten.mm.default args=((56, 256), (256, 8)) MATCHED=True
   

2. Routing is byte-identical across forward and recompute. Probing right after
   weights, indices, aux_loss = self.gate(...) in
   nemo_automodel/components/moe/layers.py, the same layer object (id=) yields
   identical token counts, index sums, masks, and per-expert bincounts:

   forward   id=...157750 ntok=58 isum=309 mask=57 counts=[0,55,0,1,49,11,0,0]
   recompute id=...157750 ntok=58 isum=309 mask=57 counts=[0,55,0,1,49,11,0,0]
   forward   id=...1566d0 ntok=58 isum=388 mask=57 counts=[48,0,6,0,0,0,58,4]
   recompute id=...1566d0 ntok=58 isum=388 mask=57 counts=[48,0,6,0,0,0,58,4]
   

So the routing decision is preserved, yet the post-dispatch token count still
drifts. The divergence originates below the router, in the HybridEP
dispatch/all-to-all layer.

Suspected cause

_HybridEPManager in
nemo_automodel/components/moe/megatron/token_dispatcher.py:340-445 carries
mutable per-instance state across the dispatch call — self.num_permuted_tokens,
self.pad_multiple, self.routing_map, self.token_probs, self.handle.
Under non-reentrant activation checkpointing the same manager instance is
re-entered during recompute, so stale state from the forward pass is visible.

Only one of those fields has been guarded, and its comment names this exact
mechanism (token_dispatcher.py:408-410):

# Reset num_permuted_tokens to None to avoid reusing cached state from a prior dispatch.
# This can happen in non-reentrant activation checkpointing mode.

The remaining fields appear to need the same treatment.

Minimal repro

# 4x H20 (sm_90). Any MoE checkpoint; smallest observed repro is a 4-layer /
# 8-expert config, so no large download is needed to reproduce.
torchrun --standalone --nproc_per_node=4 -m molt.cli.train_sft \
  --model.model_name_or_path /path/to/Qwen3.6-35B-A3B \
  --data.dataset sft.jsonl --data.input_key prompt --data.output_key response \
  --data.max_len 256 --data.max_samples 128 \
  --model.lora.rank 16 --model.lora.alpha 32 \
  --model.lora.target_modules '*_proj' '*experts' \
  --train.max_epochs 2 --train.batch_size 16 --train.micro_batch_size 1 \
  --fsdp.param_dtype bf16 --fsdp.attn_implementation sdpa \
  --fsdp.ep_size 4 \
  --model.gradient_checkpoint full \
  --model.aux_loss_coef 0.001 --adam.lr 2e-4

# Fails at "Train step of epoch 0: 0%| | 0/32" on all ranks, rc=1.
# Workaround: --model.gradient_checkpoint none  ->  2 epochs complete cleanly.

Reproduced across four models, two architectures, with and without LoRA

model	architecture	LoRA	drift
4-layer / 8-expert MoE	Qwen3MoeForCausalLM	no (full FT)	151 → 150
Qwen3-30B-A3B-Base	Qwen3MoeForCausalLM	yes	271 → 268
Qwen3.5-35B-A3B	Qwen3_5MoeForConditionalGeneration	yes	218 → 220
Qwen3.6-35B-A3B	Qwen3_5MoeForConditionalGeneration	yes	246 → 248

The drift is not one-directional (both shrink and grow occur), is independent of
LoRA, model scale, and architecture, and reproduces at ep_size 2 and 4.

MOLT_MOE_DISPATCHER=torch does not hit this — the DTensor all-gather path
produces fixed shapes. The failure is specific to the ragged permuted buffer used
by hybridep (and, by inspection of the shared code path, deepep).

Expected behavior

gradient_checkpoint full composes with ep_size > 1, as the shipped recipes
assume. Recompute should observe the same permuted token count as the forward
pass, so backward succeeds.

Failing that, molt should reject the combination up front with an actionable
message instead of surfacing a raw CheckpointError at step 0.

Environment

Branch / commit:      main @ 64b6e44
Container:            not the project image; local venv matching dockerfile/Dockerfile
                      pins (torch 2.11.0+cu130, DeepEP 42144303 incl. DeepEP #638,
                      AutoModel a3aa09bcc == requirements.txt pin). Python 3.11
                      vs the image's 3.12.
GPU / machine:        4x NVIDIA H20 (sm_90, 78 SMs, 97871 MiB), full-mesh NV18 NVLink
Python / CUDA / torch: 3.11 / 13.0 / 2.11.0+cu130
Dispatcher:           hybridep (default), GroupedExpertsDeepEPLoRA on all 40 layers

Note on H20: HybridEP's preprocessing kernel defaults to a grid of 108 blocks,
which exceeds this GPU's 78 SMs. Since the kernel does a grid-wide scan requiring
all blocks to be co-resident, it deadlocks before reaching the bug above; the
grid must be clamped to the device SM count first. That is a separate issue.

Logs

[LoRA] rank=16 alpha=32 dropout=0.0: trainable 934.1M / 36041.2M params (2.59%)
Train step of epoch 0:   0%|          | 0/32

[rank2]: Traceback (most recent call last):
[rank2]:     raise CheckpointError(
[rank2]: torch.utils.checkpoint.CheckpointError: torch.utils.checkpoint: Recomputed
[rank2]: values for the following tensors have different metadata than during the
[rank2]: forward pass.
[rank2]: tensor at position 91:
[rank2]: saved metadata:      {'shape': torch.Size([246, 2048]), 'dtype': torch.bfloat16, ...}
[rank2]: recomputed metadata: {'shape': torch.Size([248, 2048]), 'dtype': torch.bfloat16, ...}
[rank2]: tensor at position 94:
[rank2]: saved metadata:      {'shape': torch.Size([246, 16]), 'dtype': torch.bfloat16, ...}
[rank2]: recomputed metadata: {'shape': torch.Size([248, 16]), 'dtype': torch.bfloat16, ...}
[rank2]: tensor at position 95:
[rank2]: saved metadata:      {'shape': torch.Size([246, 1024]), 'dtype': torch.bfloat16, ...}
[rank2]: recomputed metadata: {'shape': torch.Size([248, 1024]), 'dtype': torch.bfloat16, ...}
rc=1

The three shapes are the same tensor at three widths: hidden (2048), the LoRA
rank-16 intermediate, and the fused gate/up projection (2 x moe_intermediate 512).
Positions 91/94/95 are identical across models; only the token count differs.