[AutoEP]Fix optimizer and replaced MOE parameter mismatch

[pengdurice]

Fix: AutoEP silently drops expert parameters from a caller-supplied optimizer

Problem

_configure_expert_parallel (engine.py) replaces every MoE module, and it runs before
_configure_optimizer. Nothing remaps optimizer param groups in between —
set_optimizer_flags only sets Muon flags.

torch.optim.Optimizer.__init__ materialises its argument eagerly (param_groups = list(params)),
so an optimizer the caller built from model.parameters() keeps hard references to the discarded
expert tensors, while the live GroupedExperts weights belong to no param group.

This is the path HF Trainer and Accelerate take whenever the DeepSpeed config declares no
optimizer block (transformers/integrations/deepspeed.py: optimizer = trainer.create_optimizer()
in the else branch).

Two symptoms, same cause:

setup	result
with zero.Init (zero3_init_flag: true)	silent — every expert and router is frozen. Loss still falls because attention, shared experts and norms train normally.
without zero.Init	AttributeError: 'Parameter' object has no attribute 'partition_numel' from _create_fp16_sub_groups (stage3.py), because the stale params were never ZeRO-converted.

Measured on a 2-layer / 4-expert model, ZeRO-3, autoep_size=2, one step at lr=0.1:

parameter                                kind   max |delta|  verdict
model.layers.0.self_attn.*               LIVE   1.003e-01    moved
model.layers.0.mlp.router.gate.weight    LIVE   0.000e+00    *** FROZEN ***
model.layers.0.mlp.experts.w1/w2/w3      LIVE   0.000e+00    *** FROZEN ***
model.layers.0.mlp.experts.gate_up_proj  GHOST  0.000e+00    frozen   (detached from the model)

The same run with the optimizer declared in ds_config moves every tensor by 1.005e-01.

Fix

_remap_client_optimizer_after_module_replacement, called immediately after
_configure_expert_parallel:

• no-op unless a caller supplied a real optimizer (DummyOptim / config-built paths return early);
• no-op unless something was actually replaced;
• drops params no longer in the module, adds the new ones to the group the stale ones came from,
  so per-group hyper-parameters (e.g. a weight-decay split) survive;
• purges stale optimizer.state entries, which Adam keys on the parameter object;
• raises rather than guesses if replacement invalidated params across more than one group — a wrong
  assignment would silently mis-apply weight decay.

Tests

New tests/unit/v1/moe/test_autoep_client_optimizer.py, world_size = 2:

• test_client_optimizer_covers_replacement_parameters — every trainable param is in a param group;
  replacement actually happened; no optimized param is detached from the module.
• test_client_optimizer_updates_expert_weights — after one real step, no parameter is
  bit-identical. This is the test that catches the silent freeze.
• test_client_optimizer_preserves_param_group_hyperparameters — with decay / no-decay groups,
  replacement params land in the right group and both weight_decay values survive.

3 passed with the fix; 3 failed with engine.py reverted and everything else identical.

The tests use bf16 rather than the shared mixed_precision_config() helper: fp16 carries a loss
scaler that skips the first optimizer step on overflow, which would leave every parameter
unchanged and make the update test vacuous.

Scope

ZeRO-3 only. A caller-supplied optimizer with MoE on ZeRO-1/2 additionally requires param groups
marked {"moe": True} (stage_1_and_2.py:780, bf16_optimizer.py:128) — a separate pre-existing
requirement, not addressed here.

Why this was not caught

Every existing AutoEP test supplies the optimizer through ds_config
("optimizer": {"type": "Adam"} in make_autoep_config), so DeepSpeed builds it after
replacement and the client-optimizer path is never exercised.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: fc3c948431

ℹ️ 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]

Add the required Signed-off-by trailer

This non-merge commit has no Signed-off-by trailer, so it violates the repository's mandatory commit policy and may fail DCO/CI validation. Recreate the commit with --signoff using the configured Git identity.

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

Useful? React with 👍 / 👎.

[chatgpt-codex-connector]

Reject partially optimized grouped expert sources

When module_list storage packs several local experts into one replacement tensor and the client optimizer contains only some of those source weights, this filtering produces one nonempty group and assigns the entire grouped tensor there. The previously excluded experts then receive gradients and optimizer updates, silently defeating selective-expert fine-tuning; detect sources missing from all groups and reject this partially represented mapping rather than treating it as unambiguous.

Useful? React with 👍 / 👎.

[chatgpt-codex-connector]

Preserve unrelated external optimizer parameters

When AutoEP performs any replacement under ZeRO-0 and a valid caller-supplied optimizer also owns a parameter outside model.parameters()—for example, an auxiliary trainable loss parameter—this filter removes that parameter merely because it is not in the module tree; the following state cleanup also discards its optimizer state. Restrict removal to parameters known to belong to discarded AutoEP modules instead of every non-model parameter.

Useful? React with 👍 / 👎.

[tohtana]

Hi @pengdurice, thank you for the fix!
The direction looks good, and I left a few comments about some cornor cases.

I opened a PR to clarify the issues and share suggested fix: pengdurice#1
Feel free to refer to it when you think it is useful.

[tohtana]

Thank you for accepting the PR. This looks good to me.