Support discard-output recompute for MoE shared experts under A2A overlap

[wplf]

What

Switches shared_experts selective recompute from a standard checkpoint (which keeps the output activation) to CheckpointWithoutOutput: the shared-expert output activation is discarded in the forward and regenerated in backward from a grad hook. Supported on both the A2A-overlap (fine-grained callables) and the single-call (non-overlap) MoE paths.

• tensor_parallel/random.py: split CheckpointWithoutOutput.discard_output_and_register_recompute into discard_output() + register_recompute_hook() so the output can be freed and the recompute hook registered at different points.
• moe/moe_layer.py: shared_experts_compute uses CheckpointWithoutOutput when discard-output recompute is enabled (gated only off MoE cudagraph partial capture, where the shared-expert output is a graph output that must keep its storage).
• models/gpt/fine_grained_callables.py (A2A-overlap path): register the recompute on the moe node's expert_output (right after the pre_mlp_norm recompute, same tensor) and free the output in the combine node after postprocess() consumes it.
• transformer/transformer_layer.py (non-overlap path): in _forward_post_mlp, free the output and register the recompute on mlp_output_with_bias[0] right after the pre_mlp_norm recompute.

Ordering invariant (both paths)

The shared expert's input is pre_mlp_layernorm_output, which the layernorm recompute discards and only restores in its own backward hook. So the shared-expert recompute must fire after the pre_mlp_norm recompute (input restored) and before the shared expert's backward. Both paths register the hook on the same tensor as — and immediately after — pre_mlp_norm, guaranteeing the order. A hook placed too early (e.g. on the combine node output, which detaches between schedule nodes) hits a freed pre_mlp_layernorm_output → cuBLAS error. The two paths are mutually exclusive (overlap bypasses _forward_post_mlp; non-overlap bypasses the fine-grained callables), so there is no double registration.

Memory savings

Discard-output frees one extra activation that a standard checkpoint retains: the shared-expert output, shape [seq, mbs, hidden] (~0.5 GB/MoE-layer at seq=4096, hidden=4096, bf16, mbs=4). The realized peak saving is smaller than the per-layer × num-layers sum because only the few layers whose shared-expert output is co-resident at the instantaneous peak (bwd → optimizer) are freed.

Controlled measurement on the Qwen3.5-VL 397B proxy (8×GB200, EP=8, mbs=4, full stack: layernorm + gdn_norm_out + moe_act + shared_experts recompute + offload + precision-aware optimizer + A2A overlap) — same recipe, the only difference is the shared_experts checkpoint mechanism:

shared_experts recompute	peak max_allocated
standard checkpoint (keeps output)	85.94 GB
discard-output (this PR)	84.75 GB
saving	−1.19 GB

Notes:

• This is an incremental saving on top of an already-recompute+offload stack (standard checkpoint already recomputes the shared-expert intermediates; this PR additionally frees the output). It is not a large single-point reduction.
• The saving scales ~linearly with micro_batch_size (output ∝ mbs) and with how many layers' shared-expert outputs are co-resident at the peak.
• Its main value is bringing the shared expert into the same discard-output recompute scheme as layernorm / gdn_norm_out, rather than a large standalone win.

Validation (Qwen3.5-VL 397B proxy, 8×GB200, EP=8, mbs=4, multimodal entry)

Full stack, 20 iters (10 for non-overlap), stable loss, no NaN:

path	result
A2A overlap (+ offload + paopt)	20 iters clean, 84.75 GB max_alloc
non-overlap	10 iters clean (functional; no offload/paopt in this run)
isolated (recompute_modules=[shared_experts], no layernorm nesting)	20 iters clean

Numerical parity (3-way, same seed=1234, A2A overlap on, mbs=4)

REF (shared_experts NOT recomputed) vs TEST (discard-output) vs REF-dup (REF re-run, FP-noise floor), per-iter lm loss:

iter	REF	TEST	REF-dup	|REF−TEST|	|REF−dup|
1	13.24673	13.24672	13.24665	1e-5	8e-5
5	13.23938	13.23925	13.23932	1.3e-4	6e-5
6	13.24373	13.24389	13.24388	1.6e-4	1.5e-4
10	13.24379	13.24394	13.24366	1.5e-4	1.3e-4

|REF − TEST| is the same magnitude as the |REF − REF-dup| non-determinism floor with no systematic drift; grad norms match to 3 digits. So discard-output recompute is numerically equivalent to not recomputing, within run-to-run FP noise. (Bitwise-deterministic comparison isn't possible: SM100 attention backward with head_dim=256 has no deterministic kernel, so --deterministic-mode asserts out for this model.)

🤖 Generated with Claude Code

[Wohox]

Please add unit test for shared expert recompute + ep overlap, you can refer to tests under unit_test/a2a_overlap/.

[Wohox]

LGTM