Preserve AutoEP score correction bias buffers (#8369)

## Summary

Fix AutoEP replacement silently dropping `e_score_correction_bias` when
Hugging Face models register it as a buffer or place it somewhere other
than the router module expected by the preset.

The parser now records the bias owner's actual path inside each MoE
layer. Replacement then copies the bias while preserving whether it is
an `nn.Parameter` or registered buffer, including the buffer's
persistence setting.

The implementation supports locations such as:

- `gate`
- the MoE block itself
- `router`
- nested modules such as `gate.moe_statics`

Ambiguous layers containing more than one `e_score_correction_bias` are
rejected rather than selecting one silently. Unsupported
non-parameter/non-buffer values produce a warning.

## Validation

- Full AutoEP unit test file: `103 passed`
- All pre-commit hooks passed
- Tested with a real Transformers 5.15.1 `DeepseekV3MoE`
- GPU validation environment:
  - NVIDIA GeForce RTX 5090
  - PyTorch 2.8.0+cu128
  - Transformers 5.15.1
  - DeepSpeed 0.19.6
- The replacement preserved the CUDA registered buffer and its value
- Native and EP2 replacement routing selected the same experts
- Single-GPU full MoE forward comparison had a maximum absolute
difference of `0.0`

The available validation machine has one GPU, so a multi-GPU
collective/ZeRO-3 end-to-end run was not performed.

Fixes #8358

Signed-off-by: poorpaper <43514747+poorpaper@users.noreply.github.com>

Author: poorpaper

deepspeed/module_inject/auto_ep.py

@@ -82,6 +82,19 @@ def _raise_if_duplicate_moe_specs(specs: list[MoELayerSpec]) -> None:
                      "AutoEP patterns so each MoE module matches exactly one preset.")
 
 
+def _resolve_e_score_correction_bias_path(moe_module: nn.Module, module_name: str) -> str | None:
+    matches = [
+        path for path, candidate in moe_module.named_modules()
+        if getattr(candidate, "e_score_correction_bias", None) is not None
+    ]
+    if len(matches) > 1:
+        locations = ", ".join(repr(path or "<root>") for path in matches)
+        raise ValueError(
+            f"AutoEP found e_score_correction_bias in multiple locations in layer '{module_name}': {locations}. "
+            "Keep exactly one score correction bias in each MoE layer so AutoEP can preserve it unambiguously.")
+    return matches[0] if matches else None
+
+
 def _source_param_shape(param: torch.Tensor | nn.Parameter) -> torch.Size:
     if is_zero_param(param):
         return torch.Size(param.ds_shape)
@@ -413,6 +426,7 @@ def ep_parser(self) -> list[MoELayerSpec]:
                     gate_bias = getattr(router_child, 'bias', None) is not None
 
                 forward_contract = adapter.adjust_forward_contract(_detect_forward_contract(module, router_child))
+                e_score_correction_bias_path = _resolve_e_score_correction_bias_path(module, module_name)
 
                 # Check shared experts
                 has_shared = False
@@ -472,6 +486,7 @@ def ep_parser(self) -> list[MoELayerSpec]:
                     preset_adapter=preset.preset_adapter,
                     router_logits_capture_mode=forward_contract.router_logits_capture_mode,
                     moe_output_shape=forward_contract.moe_output_shape,
+                    e_score_correction_bias_path=e_score_correction_bias_path,
                 )
                 specs.append(spec)
                 logger.debug(f"Detected MoE layer: {module_name} (family={preset_name}, "

deepspeed/module_inject/auto_ep_layer.py

@@ -85,6 +85,33 @@ def _copy_parameter_data(target: nn.Parameter, source: torch.Tensor) -> None:
         target.data.copy_(source_data)
 
 
+def _copy_e_score_correction_bias(
+    target_router: nn.Module,
+    source_owner: nn.Module,
+    source_bias,
+    source_path: str,
+) -> None:
+    if isinstance(source_bias, nn.Parameter):
+        target_router.e_score_correction_bias = nn.Parameter(source_bias.data.clone(),
+                                                             requires_grad=source_bias.requires_grad)
+    elif (torch.is_tensor(source_bias) and source_owner._buffers.get("e_score_correction_bias") is source_bias):
+        copied_bias = source_bias.detach().clone()
+        copied_bias.requires_grad_(source_bias.requires_grad)
+        if hasattr(target_router, "e_score_correction_bias"):
+            delattr(target_router, "e_score_correction_bias")
+        persistent = "e_score_correction_bias" not in source_owner._non_persistent_buffers_set
+        target_router.register_buffer("e_score_correction_bias", copied_bias, persistent=persistent)
+    else:
+        logger.warning(
+            "AutoEP: cannot copy e_score_correction_bias from source module path '%s': expected "
+            "an nn.Parameter or registered buffer, got %s.", source_path or "<root>",
+            type(source_bias).__name__)
+        return
+
+    logger.info("AutoEP: copied e_score_correction_bias from source module path '%s' (shape=%s)", source_path
+                or "<root>", source_bias.shape)
+
+
 def apply_scores_before_experts_if_enabled(
     routed_input: torch.Tensor,
     top_scores: torch.Tensor,
@@ -398,7 +425,14 @@ def __init__(
         # Router: copy gate weights from source
         source_gate = getattr(source_module, spec.router_name)
         source_gate_bias = getattr(source_gate, 'bias', None)
-        source_ecb = getattr(source_gate, 'e_score_correction_bias', None)
+        source_ecb_path = spec.e_score_correction_bias_path
+        if source_ecb_path is None:
+            source_ecb_owner = source_gate
+            source_ecb_path = spec.router_name
+        else:
+            source_ecb_owner = (source_module
+                                if source_ecb_path == "" else source_module.get_submodule(source_ecb_path))
+        source_ecb = getattr(source_ecb_owner, "e_score_correction_bias", None)
         unsupported_router_biases = [
             getattr(source_gate, bias_name, None) for bias_name in spec.unsupported_router_bias_names
         ]
@@ -434,11 +468,8 @@ def __init__(
                 self.router.gate.bias.requires_grad_(source_gate_bias.requires_grad)
 
             # Copy pre-trained score correction bias (DeepSeek-V3/Moonlight noaux_tc routing)
-            if source_ecb is not None and isinstance(source_ecb, nn.Parameter):
-                self.router.e_score_correction_bias = nn.Parameter(source_ecb.data.clone(),
-                                                                   requires_grad=source_ecb.requires_grad)
-                logger.info('AutoEP: copied e_score_correction_bias from source gate '
-                            '(shape=%s)', source_ecb.shape)
+            if source_ecb is not None:
+                _copy_e_score_correction_bias(self.router, source_ecb_owner, source_ecb, source_ecb_path)
 
         # Alias router under the name OutputRecorder expects (layer_name if provided),
         # but only when OutputRecorder captures from the router child and the alias is safe.

deepspeed/module_inject/auto_ep_presets/base.py

@@ -90,6 +90,7 @@ class MoELayerSpec:
     preset_adapter: str = "default"
     router_logits_capture_mode: Literal["raw", "post_score"] = "post_score"
     moe_output_shape: Literal["batched", "flat"] = "batched"
+    e_score_correction_bias_path: str | None = None
 
 
 @dataclass

tests/unit/v1/moe/test_autoep_unit.py

@@ -1308,7 +1308,9 @@ def test_deepseek_v3_detection_and_score_correction_bias_copy(self, monkeypatch)
         FakeGatheredParameters.calls = []
         monkeypatch.setattr(ep_repack, "GatheredParameters", FakeGatheredParameters)
         monkeypatch.setattr(get_preset_adapter("deepseek_v3"), "_installed_transformers_version", lambda: "5.0.0")
+
         model = MockDeepSeekV3Transformer(num_layers=1, num_experts=8)
+
         auto_ep = AutoEP(model, _runtime_config(enabled=True, autoep_size=2))
         specs = auto_ep.ep_parser()
 
@@ -1317,6 +1319,7 @@ def test_deepseek_v3_detection_and_score_correction_bias_copy(self, monkeypatch)
         assert specs[0].expert_storage == "module_list"
         assert specs[0].expert_w1_name == "gate_proj"
         assert specs[0].has_shared_experts is True
+        assert specs[0].e_score_correction_bias_path is None
 
         source_bias = torch.arange(8, dtype=torch.float32)
         model.model.layers[0].mlp.gate.e_score_correction_bias = nn.Parameter(source_bias.clone())
@@ -1333,6 +1336,59 @@ def test_deepseek_v3_detection_and_score_correction_bias_copy(self, monkeypatch)
         torch.testing.assert_close(replaced.router.e_score_correction_bias, source_bias)
         assert ["router.e_score_correction_bias"] in [call["names"] for call in FakeGatheredParameters.calls]
 
+    @pytest.mark.parametrize(
+        "owner_path,bias_kind,persistent",
+        [
+            ("gate", "buffer", True),
+            ("", "buffer", False),
+            ("router", "buffer", True),
+            ("gate.moe_statics", "parameter", True),
+        ],
+    )
+    def test_score_correction_bias_location_and_registration(self, monkeypatch, owner_path, bias_kind, persistent):
+        monkeypatch.setattr(get_preset_adapter("deepseek_v3"), "_installed_transformers_version", lambda: "5.0.0")
+        model = MockDeepSeekV3Transformer(num_layers=1, num_experts=8)
+        source = model.model.layers[0].mlp
+        owner = source
+        for part in owner_path.split(".") if owner_path else ():
+            if not hasattr(owner, part):
+                owner.add_module(part, nn.Module())
+            owner = getattr(owner, part)
+
+        source_bias = torch.arange(8, dtype=torch.float32)
+        if bias_kind == "parameter":
+            owner.e_score_correction_bias = nn.Parameter(source_bias.clone(), requires_grad=False)
+        else:
+            owner.register_buffer("e_score_correction_bias", source_bias.clone(), persistent=persistent)
+
+        auto_ep = AutoEP(model, _runtime_config(enabled=True, autoep_size=2))
+        spec = auto_ep.ep_parser()[0]
+        assert spec.e_score_correction_bias_path == owner_path
+
+        auto_ep.replace_moe_layer(spec, ep_size=2, ep_rank=0)
+
+        replaced_bias = model.model.layers[0].mlp.router.e_score_correction_bias
+        torch.testing.assert_close(replaced_bias, source_bias)
+        assert replaced_bias.requires_grad is False
+        if bias_kind == "parameter":
+            assert dict(
+                model.model.layers[0].mlp.router.named_parameters())["e_score_correction_bias"] is replaced_bias
+            assert "e_score_correction_bias" not in dict(model.model.layers[0].mlp.router.named_buffers())
+        else:
+            assert dict(model.model.layers[0].mlp.router.named_buffers())["e_score_correction_bias"] is replaced_bias
+            assert "e_score_correction_bias" not in dict(model.model.layers[0].mlp.router.named_parameters())
+            assert ("e_score_correction_bias" in model.model.layers[0].mlp.router.state_dict()) is persistent
+
+    def test_score_correction_bias_multiple_locations_are_rejected(self, monkeypatch):
+        monkeypatch.setattr(get_preset_adapter("deepseek_v3"), "_installed_transformers_version", lambda: "5.0.0")
+        model = MockDeepSeekV3Transformer(num_layers=1, num_experts=8)
+        source = model.model.layers[0].mlp
+        source.register_buffer("e_score_correction_bias", torch.zeros(8))
+        source.gate.register_buffer("e_score_correction_bias", torch.ones(8))
+
+        with pytest.raises(ValueError, match="e_score_correction_bias in multiple locations"):
+            AutoEP(model, _runtime_config(enabled=True, autoep_size=2)).ep_parser()
+
 
 def _eager_pep604_lines(module):
     """Line numbers where a module evaluates PEP 604 unions at import time."""