Preserve AutoEP score correction bias buffers

[poorpaper]

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

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 83a4f6c5f4

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

Preserve gate lookup when the parsed path is absent

When a caller constructs MoELayerSpec directly (the new field defaults to None) or attaches e_score_correction_bias between ep_parser() and replace_moe_layer(), this leaves source_ecb unset even if the source gate has a valid bias. The previous implementation discovered the bias from source_gate during replacement, and the pre-change unit test explicitly exercised the post-parser attachment case; the revised test now moves that attachment before parsing. This regression silently drops the correction bias and changes expert selection, so replacement should fall back to inspecting source_gate when no path was recorded.

Useful? React with 👍 / 👎.

[poorpaper]

fixed