fix(models): use HybridEP for MiniMax-M3 expert parallelism - #5448
Conversation
Signed-off-by: yaoyu-33 <yaoyu.094@gmail.com>
Signed-off-by: yaoyu-33 <yaoyu.094@gmail.com>
|
/ok to test fa3b371 |
|
Light review - LGTM with one question. The change is well-scoped: it adds MiniMaxM3MoELayer to widen BF16 EP all-to-all payloads to FP32 during inference (correctness workaround for silent corruption), routes the block spec through it, adds decode_generated_tokens to strip the prompt/placeholders before decoding, and rewrites the verification card to honestly record the bounded (non-EOS) completion. Unit tests cover the inference and training dtype paths and the decode helper. No 3rdparty/Megatron-LM/ files are touched. Question (not blocking): dispatch() delegates to super().dispatch(hidden_states.float(), probs), but combine() bypasses super().combine() and calls self.token_dispatcher.token_combine(output.float()) directly (minimax_m3_bridge.py:82). Is that asymmetry intentional? If the base MoELayer.combine does more than token_combine (post-process/bias handling), the direct call would silently skip it. The inference-only guard makes it likely safe today, but super().combine(output.float()) would be more future-proof. Observations:
Suggested test cases:
|
Signed-off-by: yaoyu-33 <yaoyu.094@gmail.com>
|
Thanks—the asymmetry was not intentional. In both pinned MCore revisions, |
|
/ok to test 8ed1059 |
Signed-off-by: yaoyu-33 <yaoyu.094@gmail.com>
Signed-off-by: yaoyu-33 <yaoyu.094@gmail.com>
Signed-off-by: yaoyu-33 <yaoyu.094@gmail.com>
|
/ok to test 6399f33 |
|
Light review — LGTM. Switches MiniMax-M3 from alltoall to the flex/hybridep dispatcher in the bridge provider and the H100 pretrain/SFT recipes, fixes VLM decode to strip the prompt, and updates the verification card to match. Verified:
Observation (non-blocking): the recipe/bridge set the flex fields directly rather than going through apply_flex_dispatcher_backend, so on unsupported hardware the run raises in validate_flex_dispatcher_backend instead of silently falling back to alltoall (as kimi/exaone recipes do). Reasonable fail-closed choice; just flagging the behavioral difference. Suggested test cases:
|
Signed-off-by: yaoyu-33 <yaoyu.094@gmail.com>
|
/ok to test ed69466 |
What does this PR do?
Uses HybridEP as the MiniMax-M3 model-specific expert-parallel dispatcher default for inference and the public H100 pretrain/SFT recipes. This avoids the affected BF16 grouped all-to-all transport path and replaces the earlier inference-only FP32 transport workaround.
Root cause
On the affected release runtime, finite MiniMax-M3 BF16 tensors become corrupted at variable-split EP all-to-all dispatch/combine and produce step-0 NaN logits. Isolation below MCore found stale or uninitialized
ncclIbRequest.send.onlyWriteImmstate in the SPCX grouped all-to-all path; requests can be misclassified, losing or corrupting payload data. The same model, checkpoint, prompt, and EP32 topology remain finite with NCCL internal IB or HybridEP.This is not a prompt-formatting, checkpoint-conversion, router-math, or weight-roundtrip issue. It is also separate from the older cuBLAS router crash.
Changes
flex/hybridep, 16 dispatcher SMs, with HybridEP permute fusion disabled.Bridge-created MiniMax-M3 providers use the new configuration. Existing native checkpoints may still contain a serialized
alltoallsetting, so inference callers loading those checkpoints must also override the deserialized model config.Reproduction and controlled validation
The public reproduction shape is:
With the affected all-to-all path, step 0 produces NaN logits and empty/token-0 degeneration. With the four HybridEP model settings above and an 8-rank H100 NVLink domain, a matched 32-H100 TP1/PP1/EP32/ETP1 run observed 50/50 finite full-vocabulary logit vectors, no NaNs or infinities, no token 0, natural EOS, and the exact historical token sequence ending in
Pre-LayerNorm.This was a controlled configuration-level validation on clean Bridge/MCore source, not a training run and not a performance comparison.
Validation and known dependency
No backward pass, optimizer step, convergence run, or training performance run is claimed. The public recipes require 128 or 256 H100s, beyond the 32-GPU investigation ceiling.
Packed-SFT verification also requires the padding-safe dropless HybridEP routing changes merged to MCore dev in NVIDIA/Megatron-LM#5542. The current Bridge MCore main pin does not yet contain that complete change, so this PR should not merge until the applicable MCore main/backport pin is available and the bounded training checks pass.
Before this PR is ready for review