[bug] Every shipped MoE recipe defaults to `ep_size > 1` + `gradient_checkpoint full`, which fails at step 0

[NancyFyong]

Problem

All flagship MoE recipes ship ep_size > 1 together with
--model.gradient_checkpoint full. That combination crashes on the first
backward pass with a torch.utils.checkpoint.CheckpointError, so these 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 root cause is an AutoModel defect in the HybridEP dispatch layer, filed
upstream as NVIDIA-NeMo/Automodel#3325: forward and recompute disagree on the
permuted token count after expert dispatch, so grouped-GEMM activations change
shape between the two passes. Reproduced on four models and two architectures,
with and without LoRA, at ep_size 2 and 4.

This issue tracks the two molt-side consequences.

1. ignore_router_for_ac=True does not prevent it, but the comment says it does

molt/trainer/fsdp/strategy.py:242-246 reads:

# MoE parallelization config, required when ep_size > 1.
# ignore_router_for_ac=True → selective AC that saves the router projection so
# the topk routing is NOT recomputed in backward; otherwise a near-tie token
# re-routes on recompute → per-expert counts shift → grouped-GEMM shapes drift
# ±1 → CheckpointError.

Instrumentation shows the mitigation does exactly what it claims and the error
still occurs. The router projection is matched and marked MUST_SAVE, and
routing is byte-identical between forward and recompute — same layer instance,
same token count, index sum, mask, 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]

The drift originates below the router, in the dispatch/all-to-all layer, which
ignore_router_for_ac cannot reach. As written the comment tells the next reader
this class of failure is already handled, which sent this investigation down the
wrong path for a while.

2. No guard on a combination that always fails

molt accepts --fsdp.ep_size 4 --model.gradient_checkpoint full and surfaces a
raw CheckpointError mid-backward. The message names tensor positions and shapes
with no mention of expert parallel, activation checkpointing, or the dispatcher,
so the actionable workaround (gradient_checkpoint none) is not discoverable
from it.

MOLT_MOE_DISPATCHER=torch does not hit this, so the guard has to key on the
dispatcher, not on ep_size alone.

Minimal repro

# 4x H20. Any MoE checkpoint; a 4-layer / 8-expert config also reproduces,
# so no large download is needed.
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.
# Same command with --model.gradient_checkpoint none: 2 epochs complete,
# loss 1.59 -> 0.003 on Qwen3-30B-A3B-Base.

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

Expected behavior

The shipped recipes run as written. Until Automodel#3325 is fixed, either:

• the MoE recipes default to gradient_checkpoint none on the hybridep /
  deepep dispatchers, or
• molt fails fast when a real EP dispatcher is combined with an AC mode that
  recomputes dispatch, naming the workaround and linking Automodel#3325.

Either way strategy.py:242-246 should stop claiming ignore_router_for_ac=True
prevents this CheckpointError.

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 == the 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

Separately, HybridEP's preprocessing kernel defaults to a grid of 108 blocks and
does a grid-wide scan requiring all blocks co-resident, so it deadlocks on GPUs
with fewer SMs (H20 has 78). The grid has to be clamped to the device SM count
before the run reaches the bug above. Not a molt defect, noted for reproduction.

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]: 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 widths are 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.