feat(Automodel): add MiniMax-M3 GRPO training support

[jQizhang]

What does this PR do?

Add MiniMax-M3 GRPO training support with the AutoModel backend and BF16 vLLM generation.

• Select the model load dtype based on optimizer FP32 master-weight support, allowing TE FusedAdam to load BF16 parameters while retaining FP32 master weights.
• Backport the vLLM 0.25.1 MiniMax-M3 top-k buffer layout fix, updating the indexer and sparse-attention paths together with regression coverage.
• Add a 32-node, 8-GPU-per-node non-colocated recipe using CP8/EP128 for training and TP16/EP16 for generation.
• Document setup, supported configurations, limitations, and reference training curves in the model-family guide.

Issues

No linked NeMo-RL issue. The vLLM compatibility patch references upstream issue vllm-project/vllm#48603.

Usage

From a 32-node allocation with 8 GPUs per node:

uv run examples/run_grpo.py \
  --config examples/configs/recipes/llm/grpo-minimax-m3-32n8g-automodel-cp8ep128-noncolocated.yaml

Before your PR is "Ready for review"

• Read the contributor guidelines.
• Add regression tests for optimizer-dependent load dtype and the vLLM patch.
• Run unit and functional tests locally for this submission.
• Add model documentation and a reference recipe.

Additional Information

• git diff --check passed against the locally recorded main merge base.
• The committed guide reports a short CP8/EP128 GRPO validation run and includes curves through step 100. Training was not rerun when opening this PR.
• Long-run convergence, additional training parallelism configurations, and longer sequences remain unvalidated.
• The branch includes the recipe but does not yet register a corresponding nightly driver.

[copy-pr-bot]

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.

[jQizhang]

/ok to test 3535f01

[jQizhang]

/ok to test 440d4b7

[jQizhang]

/ok to test 115ed78

[jQizhang]

/ok to test 4637622

[jQizhang]

Team review: PR #4143 — MiniMax-M3 GRPO training support

Reviewed by 6 agents (rl-expert, test-agent, bug-finder, design-reviewer, comment-reviewer, devil's-advocate). 2 non-blocking suggestions, 2 informational/minor items below.

No existing PR comments needed a response. Linters (pre-commit run --all-files) pass clean.

Generated by Claude Code

[jQizhang]

nemo_rl/models/automodel/setup.py:755-767

1 action item.

TL;DR — this silently changes stored param dtype (fp32→bf16) for 5+ pre-existing production recipes, not just the new MiniMax recipe, with no nightly evidence it's safe.

PR-introduced, not a numerics-correctness bug. Model weight storage dtype used to be unconditionally torch.float32; now it's runtime_config.dtype (e.g. bf16) whenever _has_optimizer_fp32_master is true — true for grpo-qwen3.5-35ba3b-2n8g-automodel-ep16.yaml, grpo-nemotron3-super-120BA12B-16n8g-automodel-ep8.v2.yaml, grpo-glm47-flash-4n8g-automodel.yaml, vlm_grpo-gemma4-e4b-geo3k-1n8g-automodel.yaml, and dapo-nanov3.5-30BA3B-4n8g-automodel.yaml.

Compute precision itself is unaffected: FSDP2's MixedPrecisionPolicy(param_dtype=dtype, ...) in setup_distributed already casts the all-gathered shard to bf16 for forward/backward regardless of resident storage dtype, so this only changes stored-param dtype / memory footprint (previously fp32 storage made TE FusedAdam's separate fp32 master copy redundant; now it isn't). Undisclosed in the PR description; no nightly evidence for the affected pre-existing recipes.

AI-1

Action: confirm via nightly (a) at least one affected pre-existing recipe (e.g. grpo-qwen3.5-35ba3b-2n8g-automodel-ep16) still converges post-change, and (b) DCP checkpoint save/resume compatibility across the storage-dtype change (resuming an old fp32-saved checkpoint under the new bf16-storage code, or vice versa).

[jQizhang]

docs/guides/models/minimax/minimax-m3.md:7

1 action item.

PR-introduced. The "Status: Functionally Ready" label rests on a single 100-step run (~0.72 val accuracy), while this same doc's Known Limitations section says long-run convergence is unvalidated — and there's no nightly driver to re-validate it going forward (see the test-suite registration comment on the recipe YAML).

AI-1

Action: either downgrade the status label (e.g. "Experimental"/"Preview") until a nightly driver exists, or add the nightly registration — don't ship "Functionally Ready" without an automated re-validation path.

[jQizhang]

nemo_rl/models/generation/vllm/patches.py:829

No action needed — suggestion only, not blocking.

Pre-existing pattern, this PR continues it. _patch_vllm_* is now 8 functions, hand-registered across two call sites (ensure_vllm_source_compat calls 4, _apply_vllm_patches calls 7 — an intentional subset/superset, not full duplication). A 9th patch author could add their function's body but forget one of the two registration sites — e.g. a diagnostics-only path that calls only ensure_vllm_source_compat would silently run unpatched.

Follow-up

A small registry (list[VllmPatch(name, apply, early_import_safe)] + a single loop per call site) would remove the risk of a call-site desync. Worth tracking given the file's growth rate — no tracking issue filed yet.

[jQizhang]

tests/unit/models/automodel/test_automodel_setup.py:919

No action needed — minor suggestion.

PR-introduced, minor. The new parametrized cases for _has_optimizer_fp32_master never exercise the except ImportError: return False branch or the optimizer_cls is not FusedAdam branch — both short-circuit before those paths are reached.

A small direct unit test calling _has_optimizer_fp32_master in isolation (no sys.modules monkeypatching needed) for these two branches would close the gap more cheaply than routing through the full setup_model_and_optimizer integration test.