[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.