DeepSpeed
32110a95 - Preserve AutoEP score correction bias buffers (#8369)

Commit
8 days ago
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
Parents
Loading