[AutoEP] Make router weight persistent to avoid 0 tensor size crash

[pengdurice]

Fix: ZeRO-3 releases AutoEP's router, crashing step-2 backward on fine-grained MoE

Problem

Under ZeRO-3 + AutoEP, training dies in step-2 backward for any MoE model with
num_experts * hidden_size > stage3_param_persistence_threshold:

RuntimeError: size mismatch, got input (2048), mat (2048x128), vec (0)
  at deepspeed/runtime/zero/linear.py:126, in backward

mat (2048x128) is [tokens, num_experts] -- the router gate. vec (0) is its weight, zero
elements, the placeholder ZeRO-3 leaves after releasing a partitioned parameter.

mark_persistent_parameters (parameter_offload.py:318) pins a parameter only when
param.ds_numel <= param_threshold (default 100,000). AutoEP's router gate is
num_experts x hidden_size:

model	router numel	above 100k?
MockMoETransformer (the AutoEP tests)	512	no
Mixtral-8x22B	49,152	no
GLM-4.5-Air, Qwen3-235B-A22B	524,288	yes
MiniMax-M3	786,432	yes
DeepSeek-V3 / V3.2, GLM-5.2	1.6M-1.8M	yes

Neither configuration the existing tests exercise is above the threshold, which is why this shipped.

Step 1 succeeds and step 2 fails, consistent with ZeRO-3 recording its parameter trace on the first
iteration and switching to trace-driven prefetch/release from the second.

Evidence that the threshold is the trigger

Two runs, identical except stage3_param_persistence_threshold. GLM-4.5-Air @4 layers (128 experts,
hidden 4096 -> router 524,288), 8x H200, ZeRO-3, autoep_size=8, seq 2048, AC off. Same seed --
both report Step 1, Loss: 12.784355.

threshold	router vs threshold	result
100,000 (default)	above -> release candidate	FAILED, vec (0) after 1 step
1,000,000	below -> pinned	COMPLETED, 15/15 steps, 103,542 tok/s

Fix

An opt-in persistence marker a module-replacing transform can set:

# deepspeed/runtime/zero/parameter_offload.py
if param.ds_numel <= param_threshold or getattr(param, "ds_force_persist", False):

set by AutoEP on router parameters only:

# deepspeed/module_inject/auto_ep_layer.py
param.ds_force_persist = True

Cost is num_layers * num_experts * hidden * 2 bytes -- about 98 MB for a 94-layer/128-expert
model -- for a tensor needed on every forward and backward.

Verification

Same config as the failing run, with the fix:

run	fix	result
28722	none	FAILED -- vec (0), 1 step
28724	ds_force_persist	COMPLETED -- 15/15 steps, 104,717 tok/s

Also exercised at scale on this branch's lineage: Qwen3-235B-A22B (235B, 64 GPUs) and GLM-5.2
(745B, 64 GPUs, autoep_size=32, CPU optimizer offload), both training to completion.

Alternatives considered

Add AutoEPMoELayer to DEFAULT_LEAF_MODULE_CLASSES. Works -- verified independently at
108,833 tok/s with no leaf_module entry in the user config -- but escalates the whole MoE block to
a leaf, coarsening gather granularity for the experts too. The targeted marker leaves expert
gathering unchanged.

memory_efficient_linear: false. Not a fix. It moves the failure: step 1 then succeeds and step
2 fails elsewhere with The size of tensor a (0) must match the size of tensor b (4096). Worth
noting because it looks like a fix if you stop at the first green step.

Note on test coverage

The existing AutoEP tests cannot catch this -- their router is 512 elements, ~200x below the
threshold. A regression test needs both num_experts * hidden_size > 100_000 (e.g.
num_experts=64, hidden_size=2048 = 131,072) and more than one step, since step 1 passes. In our
testing a mock containing only MoE blocks did not reproduce it even above the threshold; the release
appears to need a full transformer execution trace, so a regression test may need a more complete
model than the current mock.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 3b39766d32

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

[chatgpt-codex-connector]

Handle forced routers before applying the model persistence cap

When stage3_model_persistence_threshold is smaller than a router weight, or earlier small parameters have consumed its budget, the preceding cap check skips this new ds_force_persist branch entirely. In that supported configuration the router remains releasable and the step-2 zero-length-weight crash this patch targets still occurs; forced parameters need to bypass or reserve space outside this cap, or AutoEP should reject an insufficient cap explicitly.

Useful? React with 👍 / 👎.

[chatgpt-codex-connector]

Add regression coverage and documentation for forced persistence

This introduces a new cross-module ds_force_persist contract, but the patch adds neither a test nor documentation for it. At minimum, cover an AutoEP router above stage3_param_persistence_threshold across the second training step and the interaction with stage3_model_persistence_threshold, so the marker cannot silently stop being propagated or honored.

AGENTS.md reference: AGENTS.md:L26-L26

Useful? React with 👍 / 👎.

[chatgpt-codex-connector]

Add the required Signed-off-by trailer

The reviewed commit is a non-merge commit, but its message contains no Signed-off-by trailer, so it does not satisfy the repository's commit requirements. Recreate the commit with --signoff using the configured author identity.

AGENTS.md reference: AGENTS.md:L8-L8

Useful? React with 👍 / 👎.