Skip to content

fix(models): use HybridEP for MiniMax-M3 expert parallelism - #5448

Merged
yaoyu-33 merged 7 commits into
mainfrom
fix/minimax-m3-fp32-combine
Aug 10, 2026
Merged

fix(models): use HybridEP for MiniMax-M3 expert parallelism#5448
yaoyu-33 merged 7 commits into
mainfrom
fix/minimax-m3-fp32-combine

Conversation

@yaoyu-33

@yaoyu-33 yaoyu-33 commented Aug 8, 2026

Copy link
Copy Markdown
Contributor

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.onlyWriteImm state 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

  • Configure the MiniMax-M3 Bridge mapping to use flex / hybridep, 16 dispatcher SMs, with HybridEP permute fusion disabled.
  • Apply the same settings to the 256-H100 pretrain and 128-H100 packed-SFT recipes, including an 8-rank NVLink domain for H100 nodes.
  • Remove the MiniMax-specific inference-only BF16-to-FP32 all-to-all dispatch/combine workaround and return to the standard MCore MoE layer.
  • Keep MiniMax-M3's FP32 router projection and parameter behavior unchanged.
  • Keep the generated-suffix-only VLM decoding fix from the earlier revision of this PR.
  • Add provider/recipe assertions and update the MiniMax-M3 verification card. Training entries remain explicitly unverified.

Bridge-created MiniMax-M3 providers use the new configuration. Existing native checkpoints may still contain a serialized alltoall setting, so inference callers loading those checkpoints must also override the deserialized model config.

Reproduction and controlled validation

The public reproduction shape is:

./scripts/inference/infer.sh --nodes 4 --gpus-per-node 8 \
  --task vlm-generation \
  --hf_model_path MiniMaxAI/MiniMax-M3 \
  --hf-revision 50942730318c7943fe83db7ec8e9f9177ecb1cf8 \
  --megatron_model_path <verified-imported-checkpoint> \
  --tp 1 --pp 1 --ep 32 --etp 1 \
  --image_path docs/images/tp1.png \
  --prompt "Read the green caption near the top of the image. Answer immediately with only the exact quoted term, including capitalization and hyphen; do not explain or reason." \
  --max_new_tokens 64 \
  --trust_remote_code

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

  • Model verification card schema: passed; its focused validator suite passed (18 tests).
  • Focused MiniMax-M3 provider/recipe tests passed (37 tests).
  • Focused VLM generation utility tests passed (13 tests).
  • Full repository pre-commit hooks passed.
  • The local GPU driver is older than the repository container runtime, so these local tests ran on CPU and do not substitute for GPU training validation.
  • The earlier revision's generated-suffix decoding tests and full CI passed, but that prior CI result does not validate this new commit.

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

  • Root cause isolated below MCore.
  • Deterministic real-model HybridEP inference validated on the target topology.
  • Bridge mapping, H100 training recipes, focused assertions, and verification card updated.
  • No MCore source, dependency, CI, or public API changed.
  • Pin an MCore main revision with padding-safe HybridEP routing for packed SFT.
  • Pass updated Bridge CI.
  • Run bounded finite-loss pretrain/SFT validation before marking training verified.

Signed-off-by: yaoyu-33 <yaoyu.094@gmail.com>
Signed-off-by: yaoyu-33 <yaoyu.094@gmail.com>
@copy-pr-bot

copy-pr-bot Bot commented Aug 8, 2026

Copy link
Copy Markdown

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

@yaoyu-33 yaoyu-33 added area:model Model implementations and HF bridge logic bug Something isn't working model-minimax labels Aug 8, 2026
@yaoyu-33

yaoyu-33 commented Aug 8, 2026

Copy link
Copy Markdown
Contributor Author

/ok to test fa3b371

@yaoyu-33
yaoyu-33 marked this pull request as ready for review August 8, 2026 21:36
@yaoyu-33 yaoyu-33 added the needs-review PR is ready for code review and waiting on a reviewer label Aug 8, 2026
@claude

claude Bot commented Aug 8, 2026

Copy link
Copy Markdown
Contributor

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:

  • The FP32-widening comment is accurate: BF16 is an exact subset of FP32, so the round trip is lossless.
  • probs is intentionally left in its original dtype in dispatch(); the test asserts dispatched_probs is router_probs.
  • Card expected_result is now internally consistent with the summary (bounded 64-token completion, no OCR-target/EOS claim).

Suggested test cases:

  • No perf tests impacted.
  • tests/unit_tests/models/minimax_m3/test_minimax_m3_bridge.py::TestMiniMaxM3Bridge::test_moe_layer_uses_fp32_ep_transport_during_distributed_inference
  • tests/unit_tests/models/minimax_m3/test_minimax_m3_bridge.py::TestMiniMaxM3Bridge::test_moe_layer_preserves_bf16_ep_transport_during_training
  • tests/unit_tests/utils/test_vlm_generate_utils.py::test_decode_generated_tokens_excludes_prompt_and_uses_integer_ids

Signed-off-by: yaoyu-33 <yaoyu.094@gmail.com>
@yaoyu-33

yaoyu-33 commented Aug 8, 2026

Copy link
Copy Markdown
Contributor Author

Thanks—the asymmetry was not intentional. In both pinned MCore revisions, MoELayer.combine() currently delegates directly to token_combine(), so this was not a present correctness gap. I updated the FP32 path in 8ed1059be to call super().combine(output.float()) and preserve the base-layer contract. The focused inference transport test and all-file pre-commit checks pass.

@yaoyu-33

yaoyu-33 commented Aug 8, 2026

Copy link
Copy Markdown
Contributor Author

/ok to test 8ed1059

@yaoyu-33 yaoyu-33 added the needs-more-tests Requires additional L0 and L1 test coverage before merge label Aug 8, 2026
Signed-off-by: yaoyu-33 <yaoyu.094@gmail.com>
@yaoyu-33 yaoyu-33 changed the title fix(models): stabilize MiniMax-M3 EP inference fix(models): use HybridEP for MiniMax-M3 expert parallelism Aug 9, 2026
@yaoyu-33
yaoyu-33 marked this pull request as draft August 9, 2026 22:49
@yaoyu-33 yaoyu-33 removed the needs-review PR is ready for code review and waiting on a reviewer label Aug 9, 2026
Signed-off-by: yaoyu-33 <yaoyu.094@gmail.com>
Signed-off-by: yaoyu-33 <yaoyu.094@gmail.com>
@yaoyu-33
yaoyu-33 marked this pull request as ready for review August 10, 2026 03:34
@yaoyu-33

Copy link
Copy Markdown
Contributor Author

/ok to test 6399f33

@claude

claude Bot commented Aug 10, 2026

Copy link
Copy Markdown
Contributor

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:

  • Bridge (minimax_m3_bridge.py:349-352) and recipes set moe_token_dispatcher_type=flex, moe_flex_dispatcher_backend=hybridep, moe_flex_dispatcher_num_sms=16, moe_permute_fusion_into_hybridep=False. These are consumed by transformer_config.py (HybridEP padding) and validated at config time via validate_flex_dispatcher_backend. HybridEP supports Hopper (major 9), so H100 is valid.
  • NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN=8 matches the H100 NVL8 domain and the dispatcher-selection guidance. As an int it is coerced via str(value) in apply_env_vars and fits the COMMON_RECIPE_ENV_VARS scalar type.
  • decode_generated_tokens slices generated_ids[0, prompt_length:] (prompt_length captured before the loop, before per-step TP padding), so only generated tokens are decoded. Good unit test asserting assert_called_once_with([20, 21]).
  • Card enabled_features, bridge_commit, last_verified, and expected-result text are updated consistently to record the 64-token-bound limitation.

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:

  • No perf tests impacted (no scripts/performance/configs entries for MiniMax-M3).
  • Unit: tests/unit_tests/models/minimax_m3/test_minimax_m3_bridge.py::TestMiniMaxM3Bridge::test_provider_bridge_sets_moe_sigmoid_routing
  • Unit: tests/unit_tests/models/minimax_m3/test_minimax_m3_bridge.py::TestMiniMaxM3Bridge::test_vlm_provider_can_recover_checkpoint_compatible_text_provider
  • Unit: tests/unit_tests/recipes/minimax/test_minimax_m3_recipes.py::test_text_recipe_uses_text_only_provider[minimax_m3_pretrain_256gpu_h100_bf16_config]
  • Unit: tests/unit_tests/recipes/minimax/test_minimax_m3_recipes.py::test_text_recipe_uses_text_only_provider[minimax_m3_sft_128gpu_h100_bf16_config]
  • Unit: tests/unit_tests/utils/test_vlm_generate_utils.py::test_decode_generated_tokens_excludes_prompt_and_uses_integer_ids

Signed-off-by: yaoyu-33 <yaoyu.094@gmail.com>
@yaoyu-33

Copy link
Copy Markdown
Contributor Author

/ok to test ed69466

@yaoyu-33 yaoyu-33 added blocked Work cannot move forward until an external dependency is cleared r0.6.0 Auto-cherrypick to release branch. Apply before merge; cherrypick happens after merge. and removed blocked Work cannot move forward until an external dependency is cleared labels Aug 10, 2026
@yaoyu-33
yaoyu-33 merged commit 4868740 into main Aug 10, 2026
125 checks passed
@yaoyu-33
yaoyu-33 deleted the fix/minimax-m3-fp32-combine branch August 10, 2026 16:30
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

area:model Model implementations and HF bridge logic bug Something isn't working model-minimax needs-more-tests Requires additional L0 and L1 test coverage before merge r0.6.0 Auto-cherrypick to release branch. Apply before merge; cherrypick happens after merge.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant