diff --git a/docs/fern/versions/_nav_order.yml b/docs/fern/versions/_nav_order.yml index 66da986331..9dd320831a 100644 --- a/docs/fern/versions/_nav_order.yml +++ b/docs/fern/versions/_nav_order.yml @@ -17,6 +17,7 @@ "./versions/nightly/pages/models/llama/llama2.mdx": 1 "./versions/nightly/pages/models/mimo/mimo.mdx": 1 "./versions/nightly/pages/models/minimax/minimax-m2.mdx": 1 +"./versions/nightly/pages/models/minimax/minimax-m3.mdx": 2 "./versions/nightly/pages/models/mistral/mistral.mdx": 1 "./versions/nightly/pages/models/moonlight/moonlight.mdx": 1 "./versions/nightly/pages/models/nemotron/llama-nemotron.mdx": 1 diff --git a/docs/fern/versions/nightly.yml b/docs/fern/versions/nightly.yml index 82eca8a5b5..e411c82066 100644 --- a/docs/fern/versions/nightly.yml +++ b/docs/fern/versions/nightly.yml @@ -90,6 +90,8 @@ navigation: - contents: - page: MiniMax-M2 path: ./nightly/pages/models/minimax/minimax-m2.mdx + - page: MiniMax-M3 + path: ./nightly/pages/models/minimax/minimax-m3.mdx section: MiniMax - contents: - page: Mistral diff --git a/docs/fern/versions/nightly/pages/models/minimax/index.mdx b/docs/fern/versions/nightly/pages/models/minimax/index.mdx index 4e131a7812..14952a2fe1 100644 --- a/docs/fern/versions/nightly/pages/models/minimax/index.mdx +++ b/docs/fern/versions/nightly/pages/models/minimax/index.mdx @@ -5,3 +5,4 @@ MiniMax model documentation is organized by model variant. | Variant | Guide | |---------|-------| | MiniMax-M2 / M2.5 / M2.7 | [minimax-m2.md](minimax-m2.md) | +| MiniMax-M3 | [minimax-m3.md](minimax-m3.md) | diff --git a/docs/fern/versions/nightly/pages/models/minimax/minimax-m3.mdx b/docs/fern/versions/nightly/pages/models/minimax/minimax-m3.mdx new file mode 100644 index 0000000000..328196110e --- /dev/null +++ b/docs/fern/versions/nightly/pages/models/minimax/minimax-m3.mdx @@ -0,0 +1,54 @@ +# MiniMax-M3 + +[MiniMax-M3](https://huggingface.co/MiniMaxAI/MiniMax-M3) is a natively multimodal sparse MoE model from MiniMaxAI (428B total, ~23B active parameters). Megatron Bridge supports the M3 *language model* through the `MiniMaxM3Bridge`: the text backbone of the `MiniMaxM3SparseForConditionalGeneration` checkpoint is converted to a Megatron-Core `GPTModel`. + +## Supported Variants + +| Variant | Hugging Face ID | Notes | +|---------|-----------------|-------| +| MiniMax-M3 | `MiniMaxAI/MiniMax-M3` | Language model only (bf16 weights) | + +## Architecture Notes + +- Mixed dense/MoE decoder: 60 layers, the first 3 dense, the rest with 128 routed experts (top-4) plus one shared expert. +- Sigmoid router with expert-bias correction and `routed_scaling_factor` applied to the normalized top-k weights (DeepSeek-V3-style routing); the checkpoint's FP32 router weights remain FP32 during import. +- SwiGLU-OAI activation in every MLP and expert: clamped gate/up projections with a `+1` linear offset, mapped to `activation_func_clamp_value` / `glu_linear_offset` (same mechanism as GPT-OSS). +- Gemma-style RMSNorm (`x * (1 + w)`) on every norm, mapped to `layernorm_zero_centered_gamma`. +- GQA attention (64 query heads, 4 KV heads) with per-head QK RMSNorm and partial RoPE (64 of 128 head channels rotated, theta 5e6). + +## Known Limitations + +- **Language model only.** The CLIP-style vision tower, multimodal projector, and patch-merge MLP are not mapped. +- **Full attention only.** The lightning-indexer block-sparse attention branch (`self_attn.index_*` weights) is not mapped; the Megatron model runs full causal attention on every layer. Block selection keeps `index_topk_blocks * index_block_size` (2048) key tokens per query, so full attention is mathematically identical up to that sequence length and an approximation beyond it. +- **MTP modules are not mapped.** The released checkpoint advertises `num_nextn_predict_layers` in its config but ships no `mtp.*` weights. +- **Auxiliary-loss scoring differs for training.** The recipes use MCore's token-global load-balancing loss over normalized sigmoid scores. Hugging Face leaves its optional router loss disabled by default and uses softmax scores when enabled. +- The MXFP8 variant (`MiniMaxAI/MiniMax-M3-MXFP8`) is not supported; use the bf16 checkpoint. + +## Conversion + +```python +from megatron.bridge import AutoBridge + +bridge = AutoBridge.from_hf_pretrained("MiniMaxAI/MiniMax-M3", trust_remote_code=True) +provider = bridge.to_megatron_provider() +``` + +The bridge imports the language backbone from the multimodal checkpoint. +Standalone Hugging Face checkpoint export is not supported because the vision, +projector, lightning-indexer, and MTP weights are intentionally not mapped; +native Megatron checkpoints and in-memory weight round-trip verification are +supported. + +## Examples + +For real-checkpoint Slurm conversion, inference, hardware requirements, and +validated parallelism settings, see the [MiniMax-M3 examples README](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/examples/models/minimax/minimax_m3/README.md). + +## Recipes + +Pretraining and packed-sequence (THD) SFT recipes are available under [`src/megatron/bridge/recipes/minimax`](https://github.com/NVIDIA-NeMo/Megatron-Bridge/tree/main/src/megatron/bridge/recipes/minimax) (`minimax_m3_pretrain_256gpu_h100_bf16_config`, `minimax_m3_sft_128gpu_h100_bf16_config`), using a TP=2 / PP=4 / EP=32 baseline layout. + +## Related Implementation + +- Bridge implementation: [`src/megatron/bridge/models/minimax_m3`](https://github.com/NVIDIA-NeMo/Megatron-Bridge/tree/main/src/megatron/bridge/models/minimax_m3) +- Examples: [`examples/models/minimax/minimax_m3`](https://github.com/NVIDIA-NeMo/Megatron-Bridge/tree/main/examples/models/minimax/minimax_m3) diff --git a/docs/models/README.md b/docs/models/README.md index 113d1369dc..6a6693e537 100644 --- a/docs/models/README.md +++ b/docs/models/README.md @@ -16,7 +16,7 @@ Megatron Bridge conversion, training recipe links, and model-specific notes. | **GPT-OSS** | [GPT OSS](gpt_oss/gpt-oss.md) | | **Kimi** | [Kimi K2](kimi/kimi-k2.md), [Kimi-K2.5-VL](kimi/kimi-k25-vl.md) | | **Llama** | [Llama 2](llama/llama2.md), [Llama 3](llama/llama3.md) | -| **MiniMax** | [MiniMax-M2 / M2.5 / M2.7](minimax/minimax-m2.md) | +| **MiniMax** | [MiniMax-M2 / M2.5 / M2.7](minimax/minimax-m2.md), [MiniMax-M3](minimax/minimax-m3.md) | | **Mistral** | [Mistral](mistral/mistral.md), [Ministral 3](mistral/ministral3.md) | | **Xiaomi-MiMo** | [Xiaomi-MiMo](mimo/mimo.md) | | **Moonlight** | [Moonlight](moonlight/moonlight.md) | diff --git a/docs/models/minimax/index.md b/docs/models/minimax/index.md index fec9392bfe..c56eb391bb 100644 --- a/docs/models/minimax/index.md +++ b/docs/models/minimax/index.md @@ -6,8 +6,10 @@ MiniMax model documentation is organized by model variant. :hidden: minimax-m2.md +minimax-m3.md ``` | Variant | Guide | |---------|-------| | MiniMax-M2 / M2.5 / M2.7 | [minimax-m2.md](minimax-m2.md) | +| MiniMax-M3 | [minimax-m3.md](minimax-m3.md) | diff --git a/docs/models/minimax/minimax-m3.md b/docs/models/minimax/minimax-m3.md new file mode 100644 index 0000000000..328196110e --- /dev/null +++ b/docs/models/minimax/minimax-m3.md @@ -0,0 +1,54 @@ +# MiniMax-M3 + +[MiniMax-M3](https://huggingface.co/MiniMaxAI/MiniMax-M3) is a natively multimodal sparse MoE model from MiniMaxAI (428B total, ~23B active parameters). Megatron Bridge supports the M3 *language model* through the `MiniMaxM3Bridge`: the text backbone of the `MiniMaxM3SparseForConditionalGeneration` checkpoint is converted to a Megatron-Core `GPTModel`. + +## Supported Variants + +| Variant | Hugging Face ID | Notes | +|---------|-----------------|-------| +| MiniMax-M3 | `MiniMaxAI/MiniMax-M3` | Language model only (bf16 weights) | + +## Architecture Notes + +- Mixed dense/MoE decoder: 60 layers, the first 3 dense, the rest with 128 routed experts (top-4) plus one shared expert. +- Sigmoid router with expert-bias correction and `routed_scaling_factor` applied to the normalized top-k weights (DeepSeek-V3-style routing); the checkpoint's FP32 router weights remain FP32 during import. +- SwiGLU-OAI activation in every MLP and expert: clamped gate/up projections with a `+1` linear offset, mapped to `activation_func_clamp_value` / `glu_linear_offset` (same mechanism as GPT-OSS). +- Gemma-style RMSNorm (`x * (1 + w)`) on every norm, mapped to `layernorm_zero_centered_gamma`. +- GQA attention (64 query heads, 4 KV heads) with per-head QK RMSNorm and partial RoPE (64 of 128 head channels rotated, theta 5e6). + +## Known Limitations + +- **Language model only.** The CLIP-style vision tower, multimodal projector, and patch-merge MLP are not mapped. +- **Full attention only.** The lightning-indexer block-sparse attention branch (`self_attn.index_*` weights) is not mapped; the Megatron model runs full causal attention on every layer. Block selection keeps `index_topk_blocks * index_block_size` (2048) key tokens per query, so full attention is mathematically identical up to that sequence length and an approximation beyond it. +- **MTP modules are not mapped.** The released checkpoint advertises `num_nextn_predict_layers` in its config but ships no `mtp.*` weights. +- **Auxiliary-loss scoring differs for training.** The recipes use MCore's token-global load-balancing loss over normalized sigmoid scores. Hugging Face leaves its optional router loss disabled by default and uses softmax scores when enabled. +- The MXFP8 variant (`MiniMaxAI/MiniMax-M3-MXFP8`) is not supported; use the bf16 checkpoint. + +## Conversion + +```python +from megatron.bridge import AutoBridge + +bridge = AutoBridge.from_hf_pretrained("MiniMaxAI/MiniMax-M3", trust_remote_code=True) +provider = bridge.to_megatron_provider() +``` + +The bridge imports the language backbone from the multimodal checkpoint. +Standalone Hugging Face checkpoint export is not supported because the vision, +projector, lightning-indexer, and MTP weights are intentionally not mapped; +native Megatron checkpoints and in-memory weight round-trip verification are +supported. + +## Examples + +For real-checkpoint Slurm conversion, inference, hardware requirements, and +validated parallelism settings, see the [MiniMax-M3 examples README](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/examples/models/minimax/minimax_m3/README.md). + +## Recipes + +Pretraining and packed-sequence (THD) SFT recipes are available under [`src/megatron/bridge/recipes/minimax`](https://github.com/NVIDIA-NeMo/Megatron-Bridge/tree/main/src/megatron/bridge/recipes/minimax) (`minimax_m3_pretrain_256gpu_h100_bf16_config`, `minimax_m3_sft_128gpu_h100_bf16_config`), using a TP=2 / PP=4 / EP=32 baseline layout. + +## Related Implementation + +- Bridge implementation: [`src/megatron/bridge/models/minimax_m3`](https://github.com/NVIDIA-NeMo/Megatron-Bridge/tree/main/src/megatron/bridge/models/minimax_m3) +- Examples: [`examples/models/minimax/minimax_m3`](https://github.com/NVIDIA-NeMo/Megatron-Bridge/tree/main/examples/models/minimax/minimax_m3) diff --git a/examples/conversion/hf_megatron_roundtrip_multi_gpu.py b/examples/conversion/hf_megatron_roundtrip_multi_gpu.py index 5d2c495df1..afa30f6840 100644 --- a/examples/conversion/hf_megatron_roundtrip_multi_gpu.py +++ b/examples/conversion/hf_megatron_roundtrip_multi_gpu.py @@ -1,4 +1,4 @@ -# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2025-2026, NVIDIA CORPORATION. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -93,6 +93,8 @@ def main( trust_remote_code: bool | None = None, strict: bool = False, skip_save: bool = False, + atol: float = 1e-1, + rtol: float = 1e-5, ) -> None: """Perform round-trip conversion between HuggingFace and Megatron-LM models on multiple GPUs.""" if os.environ.get("WORLD_SIZE") is None: @@ -209,11 +211,21 @@ def main( elif compare_param.dtype != compare_original.dtype or any(p in name for p in IGNORE_PRECISION_PARAMS): compare_param = param.float() compare_original = original_param.float() - match = torch.allclose(compare_param, compare_original.to(compare_param.device), atol=1e-1) + match = torch.allclose( + compare_param, + compare_original.to(compare_param.device), + atol=atol, + rtol=rtol, + ) # --- Case 3: regular param → direct allclose --- else: - match = torch.allclose(compare_param, compare_original.to(compare_param.device), atol=1e-1) + match = torch.allclose( + compare_param, + compare_original.to(compare_param.device), + atol=atol, + rtol=rtol, + ) all_match = all_match and match table.add_row( @@ -286,6 +298,8 @@ def main( parser.add_argument( "--skip-save", action="store_true", help="Skip saving the model after comparison (verification only)" ) + parser.add_argument("--atol", type=float, default=1e-1, help="Absolute tolerance for tensor comparison") + parser.add_argument("--rtol", type=float, default=1e-5, help="Relative tolerance for tensor comparison") args = parser.parse_args() main( args.hf_model_id, @@ -298,6 +312,8 @@ def main( args.megatron_load_path, args.trust_remote_code, skip_save=args.skip_save, + atol=args.atol, + rtol=args.rtol, ) if torch.distributed.is_initialized(): diff --git a/examples/conversion/hf_to_megatron_generate_text.py b/examples/conversion/hf_to_megatron_generate_text.py index b06e86e129..7d173192dc 100644 --- a/examples/conversion/hf_to_megatron_generate_text.py +++ b/examples/conversion/hf_to_megatron_generate_text.py @@ -1,4 +1,4 @@ -# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2025-2026, NVIDIA CORPORATION. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -88,6 +88,30 @@ def loss_func(x, **kwargs): return model(**forward_args), loss_func +def _tokenize_prompt(tokenizer, prompt: str, *, apply_chat_template: bool, thinking_mode: str) -> torch.Tensor: + """Tokenize a raw prompt, optionally formatting it as a user chat turn.""" + if not apply_chat_template: + return tokenizer.encode(prompt, return_tensors="pt") + + encoded = tokenizer.apply_chat_template( + [{"role": "user", "content": prompt}], + add_generation_prompt=True, + tokenize=True, + return_dict=True, + return_tensors="pt", + thinking_mode=thinking_mode, + ) + return encoded["input_ids"] + + +def _decode_completion(tokenizer, generated_ids: torch.Tensor, prompt_length: int) -> str: + """Decode generated tokens without echoing the prompt or special tokens.""" + return tokenizer.decode( + generated_ids[0, prompt_length:].tolist(), + skip_special_tokens=True, + ) + + def main(args) -> None: """Main function for text generation from HuggingFace or Megatron models. @@ -202,7 +226,13 @@ def main(args) -> None: # Tokenize the input prompt prompt = args.prompt - input_ids = tokenizer.encode(prompt, return_tensors="pt").cuda() + input_ids = _tokenize_prompt( + tokenizer, + prompt, + apply_chat_template=args.apply_chat_template, + thinking_mode=args.thinking_mode, + ).cuda() + prompt_length = input_ids.size(1) position_ids = ( torch.arange(input_ids.size(1), dtype=torch.long, device=input_ids.device).unsqueeze(0).expand_as(input_ids) ) @@ -267,8 +297,9 @@ def main(args) -> None: if next_token_ids.item() in stop_tokens: break - # Decode the generated sequence - generated_text = tokenizer.decode(list(generated_ids[0])) + # Decode only the completion. Passing CUDA tensor objects directly to the + # tokenizer can produce corrupt text with some remote-code tokenizers. + generated_text = _decode_completion(tokenizer, generated_ids, prompt_length) print_rank_0("======== GENERATED TEXT OUTPUT ========") print_rank_0(f"Prompt: {prompt}") print_rank_0(f"Generated: {generated_text}") @@ -295,6 +326,17 @@ def main(args) -> None: default=20, help="Maximum number of new tokens to generate.", ) + parser.add_argument( + "--apply-chat-template", + action="store_true", + help="Format the prompt as a user turn using the tokenizer's chat template.", + ) + parser.add_argument( + "--thinking-mode", + choices=("enabled", "adaptive", "disabled"), + default="adaptive", + help="Thinking mode passed to the chat template when --apply-chat-template is set.", + ) parser.add_argument("--tp", type=int, default=1, help="Tensor parallelism size") parser.add_argument("--pp", type=int, default=1, help="Pipeline parallelism size") parser.add_argument("--ep", type=int, default=1, help="Expert parallelism size") diff --git a/examples/models/minimax/minimax_m3/README.md b/examples/models/minimax/minimax_m3/README.md new file mode 100644 index 0000000000..4056884d7f --- /dev/null +++ b/examples/models/minimax/minimax_m3/README.md @@ -0,0 +1,76 @@ +# MiniMax-M3 Examples + +This directory contains real-checkpoint conversion and inference examples for +[MiniMax-M3](https://huggingface.co/MiniMaxAI/MiniMax-M3). The bridge imports +the language backbone from the multimodal checkpoint; the vision tower, +projector, lightning-indexer weights, and MTP modules are not converted. + +## Hardware requirements + +MiniMax-M3 has about 428B parameters stored in bf16 (the published checkpoint +is about 869 GB). The supplied Slurm jobs use 32 GPUs with `TP=1`, `PP=1`, and +`EP=32`, which is a conservative layout for 80 GB GPUs. Hardware with larger +GPU memory can reduce `EP` and the node count as long as `EP` divides 128. + +## Setup + +Set the container and mounts without placing credentials in the scripts: + +```bash +export CONTAINER_IMAGE=/path/to/megatron-bridge.sqsh +export CONTAINER_MOUNTS=/shared:/shared,/path/to/Megatron-Bridge:/opt/Megatron-Bridge +export HF_HOME=/shared/cache/huggingface +export UV_CACHE_DIR=/shared/cache/uv +export HF_TOKEN=your_token_if_required +export SLURM_ACCOUNT=your_slurm_account +``` + +The repository is mounted at `/opt/Megatron-Bridge` by default. Override +`WORKDIR` if your mount uses a different path. Fully populate the shared model +cache before starting either 45-minute compute job; the checkpoint download is +about 869 GB: + +```bash +hf download MiniMaxAI/MiniMax-M3 +``` + +## Conversion round-trip + +Submit [slurm_conversion.sh](slurm_conversion.sh) to import the real HF +checkpoint into a distributed Megatron model and export every bridged tensor +back in memory. The job compares those tensors with the original checkpoint +and skips writing a second 869 GB copy. + +```bash +mkdir -p logs +sbatch --account="${SLURM_ACCOUNT}" examples/models/minimax/minimax_m3/slurm_conversion.sh +``` + +Success is reported only when all bridged language-model parameters match the +original checkpoint exactly (`atol=0`, `rtol=0`). This is an in-memory +verification; standalone Hugging Face checkpoint export is not supported +because the bridge intentionally omits the multimodal modules. + +## Inference + +Submit [slurm_inference.sh](slurm_inference.sh) to convert the real checkpoint, +apply the checkpoint's chat template with thinking disabled, and greedily +generate a short response with Megatron-Core: + +```bash +mkdir -p logs +sbatch --account="${SLURM_ACCOUNT}" examples/models/minimax/minimax_m3/slurm_inference.sh +``` + +The run is successful when it completes without missing-weight or forward +errors and the generated answer is coherent for the prompt. + +## Validated configuration + +The real `MiniMaxAI/MiniMax-M3` checkpoint was validated on 32 H100 80 GB +GPUs with `TP=1`, `PP=1`, `EP=32`, and `ETP=1`. All 1,053 mapped parameter +tasks passed the exact in-memory HF → Megatron → HF round-trip check. With the chat +template enabled, Megatron generated a coherent continuation beginning: + +> The sky appears blue because of Rayleigh scattering, where sunlight +> interacts with Earth's atmosphere. diff --git a/examples/models/minimax/minimax_m3/slurm_conversion.sh b/examples/models/minimax/minimax_m3/slurm_conversion.sh new file mode 100755 index 0000000000..52a6fa204f --- /dev/null +++ b/examples/models/minimax/minimax_m3/slurm_conversion.sh @@ -0,0 +1,77 @@ +#!/usr/bin/env bash +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +#SBATCH --job-name=minimax-m3-roundtrip +#SBATCH --nodes=4 +#SBATCH --ntasks-per-node=8 +#SBATCH --gpus-per-node=8 +#SBATCH --time=00:45:00 +#SBATCH --partition=batch +#SBATCH --output=logs/minimax_m3_roundtrip_%j.log +#SBATCH --exclusive + +set -euo pipefail + +: "${CONTAINER_IMAGE:?Set CONTAINER_IMAGE to the Megatron-Bridge container}" +: "${CONTAINER_MOUNTS:?Mount shared storage and this repository}" +: "${HF_HOME:?Set HF_HOME to a shared Hugging Face cache}" +: "${UV_CACHE_DIR:?Set UV_CACHE_DIR to a shared uv cache}" + +WORKDIR=${WORKDIR:-/opt/Megatron-Bridge} +HF_MODEL_ID=${HF_MODEL_ID:-MiniMaxAI/MiniMax-M3} +TP=${TP:-1} +PP=${PP:-1} +EP=${EP:-32} +ETP=${ETP:-1} + +export HF_HUB_DISABLE_TELEMETRY=1 +export HF_HUB_DISABLE_PROGRESS_BARS=1 +export TOKENIZERS_PARALLELISM=false +export TORCH_NCCL_AVOID_RECORD_STREAMS=1 +export NCCL_NVLS_ENABLE=0 +export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True + +mkdir -p logs + +MASTER_ADDR=$(scontrol show hostnames "${SLURM_JOB_NODELIST}" | head -n 1) +MASTER_PORT=${MASTER_PORT:-29500} +export MASTER_ADDR MASTER_PORT + +SRUN=( + srun + --mpi=pmix + --container-image="${CONTAINER_IMAGE}" + --container-mounts="${CONTAINER_MOUNTS}" + --no-container-mount-home +) + +"${SRUN[@]}" --nodes=1 --ntasks=1 bash -lc 'cd "$1" && uv sync --extra te' minimax-m3-sync "${WORKDIR}" + +"${SRUN[@]}" bash -lc ' + export RANK="${SLURM_PROCID:?}" + export WORLD_SIZE="${SLURM_NTASKS:?}" + export LOCAL_RANK="${SLURM_LOCALID:?}" + cd "$1" + uv run --no-sync python examples/conversion/hf_megatron_roundtrip_multi_gpu.py \ + --hf-model-id "$2" \ + --tp "$3" \ + --pp "$4" \ + --ep "$5" \ + --etp "$6" \ + --trust-remote-code \ + --skip-save \ + --atol 0 \ + --rtol 0 +' minimax-m3-roundtrip "${WORKDIR}" "${HF_MODEL_ID}" "${TP}" "${PP}" "${EP}" "${ETP}" diff --git a/examples/models/minimax/minimax_m3/slurm_inference.sh b/examples/models/minimax/minimax_m3/slurm_inference.sh new file mode 100755 index 0000000000..aabef0ce27 --- /dev/null +++ b/examples/models/minimax/minimax_m3/slurm_inference.sh @@ -0,0 +1,80 @@ +#!/usr/bin/env bash +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +#SBATCH --job-name=minimax-m3-inference +#SBATCH --nodes=4 +#SBATCH --ntasks-per-node=8 +#SBATCH --gpus-per-node=8 +#SBATCH --time=00:45:00 +#SBATCH --partition=batch +#SBATCH --output=logs/minimax_m3_inference_%j.log +#SBATCH --exclusive + +set -euo pipefail + +: "${CONTAINER_IMAGE:?Set CONTAINER_IMAGE to the Megatron-Bridge container}" +: "${CONTAINER_MOUNTS:?Mount shared storage and this repository}" +: "${HF_HOME:?Set HF_HOME to a shared Hugging Face cache}" +: "${UV_CACHE_DIR:?Set UV_CACHE_DIR to a shared uv cache}" + +WORKDIR=${WORKDIR:-/opt/Megatron-Bridge} +HF_MODEL_ID=${HF_MODEL_ID:-MiniMaxAI/MiniMax-M3} +PROMPT=${PROMPT:-Explain why the sky appears blue in one concise paragraph.} +MAX_NEW_TOKENS=${MAX_NEW_TOKENS:-64} +TP=${TP:-1} +PP=${PP:-1} +EP=${EP:-32} +ETP=${ETP:-1} + +export HF_HUB_DISABLE_TELEMETRY=1 +export HF_HUB_DISABLE_PROGRESS_BARS=1 +export TOKENIZERS_PARALLELISM=false +export TORCH_NCCL_AVOID_RECORD_STREAMS=1 +export NCCL_NVLS_ENABLE=0 +export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True + +mkdir -p logs + +MASTER_ADDR=$(scontrol show hostnames "${SLURM_JOB_NODELIST}" | head -n 1) +MASTER_PORT=${MASTER_PORT:-29500} +export MASTER_ADDR MASTER_PORT + +SRUN=( + srun + --mpi=pmix + --container-image="${CONTAINER_IMAGE}" + --container-mounts="${CONTAINER_MOUNTS}" + --no-container-mount-home +) + +"${SRUN[@]}" --nodes=1 --ntasks=1 bash -lc 'cd "$1" && uv sync --extra te' minimax-m3-sync "${WORKDIR}" + +"${SRUN[@]}" bash -lc ' + export RANK="${SLURM_PROCID:?}" + export WORLD_SIZE="${SLURM_NTASKS:?}" + export LOCAL_RANK="${SLURM_LOCALID:?}" + cd "$1" + uv run --no-sync python examples/conversion/hf_to_megatron_generate_text.py \ + --hf_model_path "$2" \ + --prompt "$3" \ + --max_new_tokens "$4" \ + --apply-chat-template \ + --thinking-mode disabled \ + --tp "$5" \ + --pp "$6" \ + --ep "$7" \ + --etp "$8" \ + --trust-remote-code +' minimax-m3-inference "${WORKDIR}" "${HF_MODEL_ID}" "${PROMPT}" "${MAX_NEW_TOKENS}" "${TP}" "${PP}" "${EP}" "${ETP}" diff --git a/src/megatron/bridge/models/__init__.py b/src/megatron/bridge/models/__init__.py index 545f0a021a..23ae85c167 100644 --- a/src/megatron/bridge/models/__init__.py +++ b/src/megatron/bridge/models/__init__.py @@ -107,6 +107,9 @@ from megatron.bridge.models.minimax_m2 import ( MiniMaxM2Bridge, ) +from megatron.bridge.models.minimax_m3 import ( + MiniMaxM3Bridge, +) from megatron.bridge.models.ministral3 import ( Ministral3Bridge, Ministral3Model, @@ -233,6 +236,7 @@ "Ministral3Model", "Ministral3ModelProvider", "MiniMaxM2Bridge", + "MiniMaxM3Bridge", "OlMoEBridge", "OlMoEModelProvider", "NemotronHBridge", diff --git a/src/megatron/bridge/models/conversion/auto_bridge.py b/src/megatron/bridge/models/conversion/auto_bridge.py index f2868fa81f..2ac0f8f116 100644 --- a/src/megatron/bridge/models/conversion/auto_bridge.py +++ b/src/megatron/bridge/models/conversion/auto_bridge.py @@ -923,6 +923,20 @@ def save_hf_pretrained( saves the configuration files, while weight saving is coordinated across all ranks. """ + # Some bridges (e.g. the language-model-only bridge of a multimodal model) cannot + # produce a valid standalone Hugging Face checkpoint; gate the export on the + # resolved bridge's capability flag. Resolving the bridge needs a concrete, + # registered architecture: a config-only save from a bare ``PretrainedConfig`` has + # none, so treat an unresolvable bridge as "nothing to gate" and fall through to the + # normal config-only path instead of failing here. + try: + model_bridge = self._model_bridge + except (ValueError, NotImplementedError): + model_bridge = None + if model_bridge is not None and not model_bridge.SUPPORTS_HF_PRETRAINED_EXPORT: + raise NotImplementedError( + f"{type(model_bridge).__name__} does not support standalone Hugging Face checkpoint export." + ) if not isinstance(self.hf_pretrained, (PreTrainedCausalLM, PretrainedConfig)): raise ValueError("save_hf_pretrained requires a pretrained HuggingFace model or config.") is_config_only = isinstance(self.hf_pretrained, PretrainedConfig) diff --git a/src/megatron/bridge/models/conversion/model_bridge.py b/src/megatron/bridge/models/conversion/model_bridge.py index b241afc4e0..334daf32be 100644 --- a/src/megatron/bridge/models/conversion/model_bridge.py +++ b/src/megatron/bridge/models/conversion/model_bridge.py @@ -414,6 +414,8 @@ def mapping_registry(self) -> MegatronMappingRegistry: - MegatronModel: The Megatron model type """ + SUPPORTS_HF_PRETRAINED_EXPORT: ClassVar[bool] = True + # Provider class to instantiate in provider_bridge (set via @register_bridge decorator) # For MLA models, use DeepSeekModelProvider or similar; for standard GPT, use GPTModelProvider PROVIDER_CLASS = None # Set by @register_bridge(provider=...) or defaults to GPTModelProvider diff --git a/src/megatron/bridge/models/minimax_m3/__init__.py b/src/megatron/bridge/models/minimax_m3/__init__.py new file mode 100644 index 0000000000..041131a887 --- /dev/null +++ b/src/megatron/bridge/models/minimax_m3/__init__.py @@ -0,0 +1,20 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from megatron.bridge.models.minimax_m3.minimax_m3_bridge import MiniMaxM3Bridge # noqa: F401 + + +__all__ = [ + "MiniMaxM3Bridge", +] diff --git a/src/megatron/bridge/models/minimax_m3/minimax_m3_bridge.py b/src/megatron/bridge/models/minimax_m3/minimax_m3_bridge.py new file mode 100644 index 0000000000..d530e8ab5f --- /dev/null +++ b/src/megatron/bridge/models/minimax_m3/minimax_m3_bridge.py @@ -0,0 +1,338 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from dataclasses import dataclass +from functools import partial + +import torch +from megatron.core.models.gpt.gpt_layer_specs import get_gpt_decoder_block_spec +from megatron.core.models.gpt.gpt_model import GPTModel +from megatron.core.transformer.moe.router import TopKRouter + +from megatron.bridge.models.conversion.mapping_registry import MegatronMappingRegistry +from megatron.bridge.models.conversion.model_bridge import MegatronModelBridge +from megatron.bridge.models.conversion.param_mapping import ( + AutoMapping, + GatedMLPMapping, + QKVMapping, +) +from megatron.bridge.models.gpt_provider import GPTModelProvider +from megatron.bridge.models.hf_pretrained.causal_lm import PreTrainedCausalLM + + +try: + import transformer_engine # noqa: F401 + + HAVE_TE = True +except (ImportError, ModuleNotFoundError): + HAVE_TE = False + +try: + from megatron.core.fusions.fused_bias_geglu import quick_gelu +except ImportError: + # Fallback if fused_bias_geglu is not available + quick_gelu = torch.nn.functional.gelu + + +def _promote_router_weights_to_float32(model: list[torch.nn.Module]) -> list[torch.nn.Module]: + """Keep MiniMax-M3 router parameters in FP32 for every load path. + + Megatron initializes router parameters in ``params_dtype`` even when + ``moe_router_dtype="fp32"``. Promoting them immediately after construction + prevents truncation when loading either HF weights or a native Megatron + checkpoint. + """ + for model_chunk in model: + for module in model_chunk.modules(): + if isinstance(module, TopKRouter) and module.weight.dtype != torch.float32: + module.weight.data = module.weight.data.float() + if isinstance(module, TopKRouter): + module._keep_in_float32_parameter_names = ("weight",) + return model + + +@dataclass +class MiniMaxM3ModelProvider(GPTModelProvider): + """GPT provider that preserves MiniMax-M3's FP32 router parameters.""" + + def __post_init__(self) -> None: + """Install the router hook on fresh and deserialized providers.""" + super().__post_init__() + self.register_pre_wrap_hook(_promote_router_weights_to_float32, prepend=True) + + +@MegatronModelBridge.register_bridge( + source="MiniMaxM3SparseForConditionalGeneration", + target=GPTModel, + model_type="minimax_m3_vl", +) +class MiniMaxM3Bridge(MegatronModelBridge): + """ + Megatron Bridge for the MiniMax-M3 language model. + + MiniMax-M3 ships as a natively multimodal checkpoint + (``MiniMaxM3SparseForConditionalGeneration``): a CLIP-style vision tower + plus a sparse-MoE text backbone. This bridge converts the *language model* + (``language_model.*`` weights) to a Megatron-Core ``GPTModel``; the vision + tower and multimodal projector are not yet bridged. + + Text backbone architecture: + - Mixed dense/MoE decoder: the first layers are dense + (``dense_intermediate_size`` MLP), the rest use 128 routed experts + (top-4) plus one shared expert. + - Sigmoid router scoring with expert-bias correction and + ``routed_scaling_factor`` applied to the normalized top-k weights + (same routing math as DeepSeek-V3). + - SwiGLU-OAI expert/MLP activation: clamped gate/up projections with a + ``+1`` linear offset (same as GPT-OSS), expressed via + ``activation_func_clamp_value`` and ``glu_linear_offset``. + - Gemma-style RMSNorm (``x * (1 + w)``) on every norm, expressed via + ``layernorm_zero_centered_gamma``. + - GQA attention with per-head QK RMSNorm and partial RoPE + (``rotary_dim`` of ``head_dim`` channels rotated). + + Known limitations: + - The lightning-indexer block-sparse attention branch + (``self_attn.index_{q,k}_{proj,norm}``) is not mapped; the Megatron + model runs full causal attention on every layer. Selection happens at + ``index_block_size`` granularity with ``index_topk_blocks`` kept per + query, so full attention is mathematically identical for sequences up + to ``index_topk_blocks * index_block_size`` tokens (2048 for the + released checkpoint) and an approximation beyond that. + - The vision tower, multimodal projector, and patch-merge MLP are not + mapped (language model only). + - MTP (Multi-Token Prediction) modules are not mapped. The released + checkpoint advertises ``num_nextn_predict_layers`` in its config but + ships no ``mtp.*`` weights, so ``mtp_num_layers`` is forced to None. + - Standalone Hugging Face checkpoint export is not supported. The HF + checkpoint is multimodal, while this bridge maps only its + ``language_model.*`` tensors. In-memory HF-to-Megatron-to-HF weight + verification remains supported. + + Example: + >>> from megatron.bridge import AutoBridge + >>> bridge = AutoBridge.from_hf_pretrained("MiniMaxAI/MiniMax-M3", trust_remote_code=True) + >>> provider = bridge.to_megatron_provider() + """ + + SUPPORTS_HF_PRETRAINED_EXPORT = False + + @classmethod + def hf_to_megatron_activation(cls, hidden_act: str): + """Convert HF activation name to Megatron activation function. + + The released MiniMax-M3 checkpoint declares ``hidden_act="swigluoai"``, + which is not a standard ACT2FN key (transformers normalizes it to + ``silu`` and computes the gate inline from ``swiglu_alpha`` / + ``swiglu_limit``). Map it to ``quick_gelu`` — the SwiGLU-OAI gate is + ``gate * sigmoid(1.702 * gate)``, i.e. exactly quick-GELU; the clamp + and ``+1`` offset are carried by separate provider fields. + """ + if hidden_act == "swigluoai": + return quick_gelu + return super().hf_to_megatron_activation(hidden_act) + + def provider_bridge(self, hf_pretrained: PreTrainedCausalLM) -> MiniMaxM3ModelProvider: + """Convert the HuggingFace MiniMax-M3 config to a GPTModelProvider.""" + hf_config = hf_pretrained.config + text_config = getattr(hf_config, "text_config", hf_config) + + provider_kwargs = self.hf_config_to_provider_kwargs(text_config) + provider_kwargs.pop("_mla_rope_params", None) + valid_fields = MiniMaxM3ModelProvider.__dataclass_fields__ + provider = MiniMaxM3ModelProvider(**{k: v for k, v in provider_kwargs.items() if k in valid_fields}) + + # Use decoder block spec to properly handle moe_layer_freq (mixed dense/MoE layers) + provider.transformer_layer_spec = partial(get_gpt_decoder_block_spec, use_transformer_engine=HAVE_TE) + + # Gemma-style RMSNorm: weights stored zero-centered, applied as x * (1 + w) + provider.normalization = "RMSNorm" + provider.layernorm_zero_centered_gamma = bool(getattr(text_config, "use_gemma_norm", True)) + provider.qk_layernorm = bool(getattr(text_config, "use_qk_norm", True)) + + provider.position_embedding_type = "rope" + provider.gated_linear_unit = True + provider.add_bias_linear = False + provider.add_qkv_bias = False + provider.hidden_dropout = 0.0 + + # tie_word_embeddings lives on text_config for MiniMax-M3 and is False + # for the released checkpoint (distinct lm_head and embed_tokens tensors) + provider.share_embeddings_and_output_weights = bool( + getattr(text_config, "tie_word_embeddings", getattr(hf_config, "tie_word_embeddings", False)) + ) + + # SwiGLU-OAI activation (same as GPT-OSS, but non-interleaved weights): + # gate = clamp(gate, max=limit); up = clamp(up, +-limit) + # out = (up + 1) * gate * sigmoid(alpha * gate), alpha = 1.702 (quick-GELU) + provider.activation_func = quick_gelu + provider.activation_func_clamp_value = float(getattr(text_config, "swiglu_limit", 7.0)) + provider.glu_linear_offset = 1.0 + + # Partial RoPE: only rotary_dim of head_dim channels are rotated + rotary_dim = getattr(text_config, "rotary_dim", None) + head_dim = getattr(text_config, "head_dim", None) + if rotary_dim is not None and head_dim: + provider.rotary_percent = rotary_dim / head_dim + + # Dense layers use dense_intermediate_size; text_config.intermediate_size + # is the per-expert FFN size (CONFIG_MAPPING would put it in ffn_hidden_size) + provider.moe_ffn_hidden_size = text_config.intermediate_size + dense_ffn_hidden_size = getattr(text_config, "dense_intermediate_size", None) + if dense_ffn_hidden_size is not None: + provider.ffn_hidden_size = dense_ffn_hidden_size + + # MoE settings — sigmoid routing with expert bias correction and + # normalized top-k weights scaled by routed_scaling_factor (DeepSeek-V3 style) + provider.moe_grouped_gemm = True + provider.moe_token_dispatcher_type = "alltoall" + provider.moe_permute_fusion = True + provider.moe_router_pre_softmax = False + provider.moe_router_score_function = "sigmoid" + provider.moe_router_enable_expert_bias = True + provider.moe_router_dtype = "fp32" + # HF exposes a token-global Switch auxiliary loss. MCore's closest + # equivalent uses normalized sigmoid scores rather than HF's softmax + # scores, but preserves the global (rather than per-sequence) scope. + provider.moe_router_load_balancing_type = "aux_loss" + provider.moe_aux_loss_coeff = getattr(text_config, "router_aux_loss_coef", 1e-3) + provider.moe_router_topk_scaling_factor = getattr(text_config, "routed_scaling_factor", 1.0) + # The overlapped shared-expert path applies a generic GLU and does not + # honor activation_func_clamp_value or glu_linear_offset. Keep it off so + # the shared expert uses MiniMax-M3's clamped (up + 1) SwiGLU-OAI math. + provider.moe_shared_expert_overlap = False + + n_shared_experts = getattr(text_config, "n_shared_experts", 0) or 0 + shared_intermediate_size = getattr(text_config, "shared_intermediate_size", 0) or 0 + provider.moe_shared_expert_intermediate_size = (n_shared_experts * shared_intermediate_size) or None + + # Per-layer dense/MoE pattern. The checkpoint config carries a 0/1 + # moe_layer_freq list; the native transformers config converts it into + # mlp_layer_types ("dense"/"sparse") strings. + moe_layer_freq = getattr(text_config, "moe_layer_freq", None) + if moe_layer_freq is None: + mlp_layer_types = getattr(text_config, "mlp_layer_types", None) + if mlp_layer_types is not None: + moe_layer_freq = [1 if layer_type == "sparse" else 0 for layer_type in mlp_layer_types] + if moe_layer_freq is not None: + provider.moe_layer_freq = [int(f) for f in moe_layer_freq] + + # The released checkpoint advertises num_nextn_predict_layers in its + # config but ships no mtp.* weights — keep MTP disabled so conversion + # does not look for weights that do not exist. + provider.mtp_num_layers = None + + provider.persist_layer_norm = True + # The fused bias-activation path only supports quick-GELU when MoE + # routing probabilities are supplied. MiniMax-M3 also uses the same + # activation in dense layers and the shared expert, so use the + # unfused path that applies the clamp and linear offset in both cases. + provider.bias_activation_fusion = False + provider.bias_dropout_fusion = True + + # Released checkpoints are bf16; text_config carries no dtype of its own + provider.fp16 = False + provider.bf16 = True + provider.params_dtype = torch.bfloat16 + provider.autocast_dtype = torch.bfloat16 + + # max_position_embeddings is 1M; keep the provider default conservative + # (recipes override seq_length explicitly) + provider.seq_length = 4096 + + return provider + + @classmethod + def megatron_to_hf_config(cls, provider: GPTModelProvider) -> dict: + """Reject standalone HF export until the full multimodal contract is mapped.""" + raise NotImplementedError( + "MiniMax-M3 standalone Hugging Face export is not supported: the source checkpoint is multimodal, " + "but this bridge maps only language_model.* tensors. Use HF import, native Megatron checkpoints, " + "or in-memory round-trip verification." + ) + + def mapping_registry(self) -> MegatronMappingRegistry: + """Return the parameter mappings for the MiniMax-M3 language model. + + All HF weights live under the ``language_model.`` prefix of the + multimodal checkpoint. MoE weights use the legacy ``block_sparse_moe`` + layout with per-expert ``w1`` (gate), ``w3`` (up), and ``w2`` (down) + tensors, the same on-disk format as MiniMax-M2. + """ + param_mappings = { + # Global weights + "embedding.word_embeddings.weight": "language_model.model.embed_tokens.weight", + "output_layer.weight": "language_model.lm_head.weight", + "decoder.final_layernorm.weight": "language_model.model.norm.weight", + # Input layernorm (fused into linear_qkv for the TE backend) + "decoder.layers.*.input_layernorm.weight": "language_model.model.layers.*.input_layernorm.weight", + "decoder.layers.*.self_attention.linear_qkv.layer_norm_weight": "language_model.model.layers.*.input_layernorm.weight", + # Post-attention layernorm: pre_mlp_layernorm on MoE layers, + # fused into linear_fc1 on dense layers + "decoder.layers.*.pre_mlp_layernorm.weight": "language_model.model.layers.*.post_attention_layernorm.weight", + "decoder.layers.*.mlp.linear_fc1.layer_norm_weight": "language_model.model.layers.*.post_attention_layernorm.weight", + # Attention + "decoder.layers.*.self_attention.linear_proj.weight": "language_model.model.layers.*.self_attn.o_proj.weight", + # Per-head QK RMSNorm (weight shape = head_dim) + "decoder.layers.*.self_attention.q_layernorm.weight": "language_model.model.layers.*.self_attn.q_norm.weight", + "decoder.layers.*.self_attention.k_layernorm.weight": "language_model.model.layers.*.self_attn.k_norm.weight", + # Dense-layer MLP down projection + "decoder.layers.*.mlp.linear_fc2.weight": "language_model.model.layers.*.mlp.down_proj.weight", + # MoE router and expert bias — on-disk uses the block_sparse_moe prefix + "decoder.layers.*.mlp.router.weight": "language_model.model.layers.*.block_sparse_moe.gate.weight", + "decoder.layers.*.mlp.router.expert_bias": "language_model.model.layers.*.block_sparse_moe.e_score_correction_bias", + # Shared expert down projection + "decoder.layers.*.mlp.shared_experts.linear_fc2.weight": "language_model.model.layers.*.block_sparse_moe.shared_experts.down_proj.weight", + } + + mapping_list = [ + AutoMapping(megatron_param=megatron_param, hf_param=hf_param) + for megatron_param, hf_param in param_mappings.items() + ] + + mapping_list.extend( + [ + # QKV + QKVMapping( + megatron_param="decoder.layers.*.self_attention.linear_qkv.weight", + q="language_model.model.layers.*.self_attn.q_proj.weight", + k="language_model.model.layers.*.self_attn.k_proj.weight", + v="language_model.model.layers.*.self_attn.v_proj.weight", + ), + # Dense-layer gated MLP + GatedMLPMapping( + megatron_param="decoder.layers.*.mlp.linear_fc1.weight", + gate="language_model.model.layers.*.mlp.gate_proj.weight", + up="language_model.model.layers.*.mlp.up_proj.weight", + ), + # Shared expert gated MLP + GatedMLPMapping( + megatron_param="decoder.layers.*.mlp.shared_experts.linear_fc1.weight", + gate="language_model.model.layers.*.block_sparse_moe.shared_experts.gate_proj.weight", + up="language_model.model.layers.*.block_sparse_moe.shared_experts.up_proj.weight", + ), + # Routed experts — on-disk layout: per-expert w1 (gate), w3 (up), w2 (down) + GatedMLPMapping( + megatron_param="decoder.layers.*.mlp.experts.linear_fc1.weight*", + gate="language_model.model.layers.*.block_sparse_moe.experts.*.w1.weight", + up="language_model.model.layers.*.block_sparse_moe.experts.*.w3.weight", + ), + AutoMapping( + megatron_param="decoder.layers.*.mlp.experts.linear_fc2.weight*", + hf_param="language_model.model.layers.*.block_sparse_moe.experts.*.w2.weight", + ), + ] + ) + + return MegatronMappingRegistry(*mapping_list) diff --git a/src/megatron/bridge/models/model_provider.py b/src/megatron/bridge/models/model_provider.py index b9514c506d..33dd5d9031 100644 --- a/src/megatron/bridge/models/model_provider.py +++ b/src/megatron/bridge/models/model_provider.py @@ -60,6 +60,42 @@ ModelT = TypeVar("ModelT", bound=MegatronModule) +def _apply_mixed_precision_wrapper( + model: list[MegatronModule], + model_config: Any, + mixed_precision_wrapper: Callable[[Any, MegatronModule], MegatronModule], +) -> list[MegatronModule]: + """Wrap a model while preserving parameters explicitly marked for FP32.""" + keep_in_fp32: list[tuple[torch.nn.Module, str, torch.Tensor]] = [] + for model_module in model: + for submodule in model_module.modules(): + # Preserve the existing MCore expert-bias contract. + if hasattr(submodule, "_maintain_float32_expert_bias"): + expert_bias = getattr(submodule, "expert_bias", None) + if expert_bias is not None: + keep_in_fp32.append((submodule, "expert_bias", expert_bias.data.clone())) + + # Model-specific modules can mark direct parameters that must not be + # truncated by Float16Module's recursive half()/bfloat16() cast. + parameter_names = vars(submodule).get("_keep_in_float32_parameter_names", ()) + if not isinstance(parameter_names, (list, tuple)): + raise TypeError("_keep_in_float32_parameter_names must be a list or tuple") + for parameter_name in parameter_names: + parameter = getattr(submodule, parameter_name, None) + if not isinstance(parameter, torch.nn.Parameter): + raise TypeError( + f"{type(submodule).__name__}.{parameter_name} must be a Parameter to remain in FP32" + ) + keep_in_fp32.append((submodule, parameter_name, parameter.data.clone())) + + wrapped_model = [mixed_precision_wrapper(model_config, model_module) for model_module in model] + + for submodule, parameter_name, fp32_data in keep_in_fp32: + getattr(submodule, parameter_name).data = fp32_data + + return wrapped_model + + class ModelProviderMixin(abc.ABC, Generic[ModelT]): """A mixin that implements the ModelProvider pattern for Megatron Bridge. @@ -632,20 +668,7 @@ def get_model( model_module.cuda(torch.cuda.current_device()) if (model_config.fp16 or model_config.bf16) and mixed_precision_wrapper is not None: - # Save expert bias in float32 to avoid precision loss during conversion - keep_in_fp32 = [] - for model_module in model: - for submodule in model_module.modules(): - if hasattr(submodule, "_maintain_float32_expert_bias"): - expert_bias = getattr(submodule, "expert_bias", None) - if expert_bias is not None: - keep_in_fp32.append((submodule, expert_bias.data.clone())) - - model = [mixed_precision_wrapper(model_config, model_module) for model_module in model] - - # Restore expert bias to float32 - for submodule, fp32_data in keep_in_fp32: - submodule.expert_bias.data = fp32_data + model = _apply_mixed_precision_wrapper(model, model_config, mixed_precision_wrapper) if correct_amax_history_if_needed is not None: correct_amax_history_if_needed(model) diff --git a/src/megatron/bridge/recipes/__init__.py b/src/megatron/bridge/recipes/__init__.py index fe00f9fb0e..613f3fc13f 100644 --- a/src/megatron/bridge/recipes/__init__.py +++ b/src/megatron/bridge/recipes/__init__.py @@ -42,6 +42,8 @@ from megatron.bridge.recipes.kimi_vl.h100 import * from megatron.bridge.recipes.llama import * from megatron.bridge.recipes.llama.h100 import * +from megatron.bridge.recipes.minimax import * +from megatron.bridge.recipes.minimax.h100 import * from megatron.bridge.recipes.ministral3 import * from megatron.bridge.recipes.ministral3.h100 import * from megatron.bridge.recipes.moonlight import * diff --git a/src/megatron/bridge/recipes/minimax/__init__.py b/src/megatron/bridge/recipes/minimax/__init__.py new file mode 100644 index 0000000000..3a291bde7e --- /dev/null +++ b/src/megatron/bridge/recipes/minimax/__init__.py @@ -0,0 +1,24 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from megatron.bridge.recipes.minimax.minimax_m3 import ( + minimax_m3_pretrain_config, + minimax_m3_sft_config, +) + + +__all__ = [ + "minimax_m3_pretrain_config", + "minimax_m3_sft_config", +] diff --git a/src/megatron/bridge/recipes/minimax/h100/__init__.py b/src/megatron/bridge/recipes/minimax/h100/__init__.py new file mode 100644 index 0000000000..2c5812a4a4 --- /dev/null +++ b/src/megatron/bridge/recipes/minimax/h100/__init__.py @@ -0,0 +1,21 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from megatron.bridge.recipes.minimax.h100.minimax_m3 import * # noqa: F403 + + +__all__ = [ + "minimax_m3_pretrain_256gpu_h100_bf16_config", + "minimax_m3_sft_128gpu_h100_bf16_config", +] diff --git a/src/megatron/bridge/recipes/minimax/h100/minimax_m3.py b/src/megatron/bridge/recipes/minimax/h100/minimax_m3.py new file mode 100644 index 0000000000..15ba7358dd --- /dev/null +++ b/src/megatron/bridge/recipes/minimax/h100/minimax_m3.py @@ -0,0 +1,197 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import torch + +from megatron.bridge import AutoBridge +from megatron.bridge.recipes.common import _pretrain_common, _sft_common +from megatron.bridge.recipes.utils.finetune_utils import default_squad_config +from megatron.bridge.recipes.utils.tokenizer_utils import DEFAULT_NULL_TOKENIZER_VOCAB_SIZE +from megatron.bridge.training.comm_overlap import CommOverlapConfig +from megatron.bridge.training.config import ConfigContainer +from megatron.bridge.training.mixed_precision import MixedPrecisionConfig + + +MINIMAX_M3_HF_PATH = "MiniMaxAI/MiniMax-M3" + + +def minimax_m3_pretrain_256gpu_h100_bf16_config() -> ConfigContainer: + """Return a pre-training config for MiniMax-M3 (428B total, ~23B active). + + MiniMax-M3 has 60 decoder layers (first 3 dense), 128 routed experts with + top-4 routing plus one shared expert, and 4 KV heads. Recommended + parallelism: TP=2, PP=4 (15 layers per stage), EP=32 (128 GPUs minimum; + 256 GPUs gives DP=2). + """ + cfg = _pretrain_common() + + cfg.model = AutoBridge.from_hf_pretrained(MINIMAX_M3_HF_PATH, trust_remote_code=True).to_megatron_provider( + load_weights=False + ) + + # Parallelism + cfg.model.tensor_model_parallel_size = 2 + cfg.model.pipeline_model_parallel_size = 4 + cfg.model.pipeline_dtype = torch.bfloat16 + cfg.model.virtual_pipeline_model_parallel_size = None + cfg.model.context_parallel_size = 1 + cfg.model.expert_model_parallel_size = 32 + cfg.model.expert_tensor_parallel_size = 1 + cfg.model.sequence_parallel = True + cfg.model.seq_length = 4096 + cfg.model.params_dtype = torch.bfloat16 + + # 60 layers split evenly across 4 stages + cfg.model.account_for_embedding_in_pipeline_split = False + cfg.model.account_for_loss_in_pipeline_split = False + cfg.model.num_layers_in_first_pipeline_stage = None + cfg.model.num_layers_in_last_pipeline_stage = None + + cfg.model.transformer_impl = "transformer_engine" + cfg.model.attention_backend = None + + cfg.model.moe_token_dispatcher_type = "alltoall" + cfg.model.moe_grouped_gemm = True + cfg.model.moe_permute_fusion = True + cfg.model.moe_router_force_load_balancing = False + cfg.model.cross_entropy_loss_fusion = True + cfg.model.cross_entropy_fusion_impl = "te" + + # Memory saving: recompute the MoE activation function only + cfg.model.recompute_granularity = "selective" + cfg.model.recompute_modules = ["moe_act"] + cfg.model.recompute_method = None + cfg.model.recompute_num_layers = None + cfg.model.cuda_graph_impl = "none" + + # Tokenizer - uses NullTokenizer by default (no HF tokenizer download needed) + cfg.tokenizer.tokenizer_type = "NullTokenizer" + cfg.tokenizer.tokenizer_model = None + cfg.tokenizer.vocab_size = DEFAULT_NULL_TOKENIZER_VOCAB_SIZE + + # Dataset config - mock data by default + cfg.dataset.blend = None + cfg.dataset.seq_length = 4096 + cfg.dataset.num_workers = 8 + + cfg.train.train_iters = 1_000_000 + cfg.train.global_batch_size = 2048 + cfg.train.micro_batch_size = 1 + cfg.train.manual_gc = True + cfg.train.manual_gc_interval = 5 + cfg.train.manual_gc_eval = 5 + cfg.validation.eval_interval = 2000 + + cfg.scheduler.lr_warmup_iters = 2000 + + cfg.logger.log_interval = 10 + cfg.checkpoint.save_interval = 2000 + cfg.checkpoint.async_save = False + + cfg.mixed_precision = MixedPrecisionConfig( + bf16=True, + params_dtype=torch.bfloat16, + pipeline_dtype=torch.bfloat16, + autocast_enabled=False, + grad_reduce_in_fp32=False, + ) + + cfg.optimizer.use_precision_aware_optimizer = True + cfg.optimizer.main_params_dtype = torch.float32 + cfg.optimizer.main_grads_dtype = torch.float32 + cfg.optimizer.exp_avg_dtype = torch.bfloat16 + cfg.optimizer.exp_avg_sq_dtype = torch.bfloat16 + + cfg.comm_overlap = CommOverlapConfig(tp_comm_overlap=False) + cfg.comm_overlap.delay_wgrad_compute = False + cfg.comm_overlap.overlap_moe_expert_parallel_comm = False + + cfg.ddp.overlap_grad_reduce = True + cfg.ddp.overlap_param_gather = True + cfg.ddp.check_for_nan_in_grad = True + cfg.ddp.use_distributed_optimizer = True + cfg.ddp.use_megatron_fsdp = False + cfg.ddp.grad_reduce_in_fp32 = False + cfg.ddp.data_parallel_sharding_strategy = "no_shard" + + return cfg + + +def minimax_m3_sft_128gpu_h100_bf16_config() -> ConfigContainer: + """MiniMax-M3 full SFT on packed (THD) sequences with Adam/bf16. + + Same TP=2 / PP=4 / EP=32 layout as the pretrain config (128 GPUs minimum). + """ + cfg = _sft_common() + + cfg.model = AutoBridge.from_hf_pretrained(MINIMAX_M3_HF_PATH, trust_remote_code=True).to_megatron_provider( + load_weights=False + ) + + # Parallelism + cfg.model.tensor_model_parallel_size = 2 + cfg.model.pipeline_model_parallel_size = 4 + cfg.model.pipeline_dtype = torch.bfloat16 + cfg.model.virtual_pipeline_model_parallel_size = None + cfg.model.context_parallel_size = 1 + cfg.model.expert_model_parallel_size = 32 + cfg.model.expert_tensor_parallel_size = 1 + cfg.model.sequence_parallel = True + cfg.model.seq_length = 4096 + cfg.model.params_dtype = torch.bfloat16 + + # 60 layers split evenly across 4 stages + cfg.model.account_for_embedding_in_pipeline_split = False + cfg.model.account_for_loss_in_pipeline_split = False + cfg.model.num_layers_in_first_pipeline_stage = None + cfg.model.num_layers_in_last_pipeline_stage = None + + cfg.model.transformer_impl = "transformer_engine" + cfg.model.attention_backend = None + + cfg.model.moe_token_dispatcher_type = "alltoall" + cfg.model.moe_grouped_gemm = True + cfg.model.moe_permute_fusion = True + cfg.model.moe_router_force_load_balancing = False + cfg.model.cross_entropy_loss_fusion = True + cfg.model.cross_entropy_fusion_impl = "te" + + # Memory saving: recompute the MoE activation function only + cfg.model.recompute_granularity = "selective" + cfg.model.recompute_modules = ["moe_act"] + cfg.model.recompute_method = None + cfg.model.recompute_num_layers = None + cfg.model.cuda_graph_impl = "none" + + # Tokenizer / dataset (real HF tokenizer; packed / THD) + cfg.tokenizer.tokenizer_model = MINIMAX_M3_HF_PATH + cfg.dataset = default_squad_config(seq_length=4096, packed_sequence=True) + + cfg.train.global_batch_size = 128 + cfg.train.micro_batch_size = 1 + + # Robustness defaults + cfg.comm_overlap = CommOverlapConfig(tp_comm_overlap=False) + cfg.comm_overlap.delay_wgrad_compute = False + cfg.comm_overlap.overlap_moe_expert_parallel_comm = False + cfg.ddp.check_for_nan_in_grad = True + cfg.ddp.use_megatron_fsdp = False + + return cfg + + +__all__ = [ + "minimax_m3_pretrain_256gpu_h100_bf16_config", + "minimax_m3_sft_128gpu_h100_bf16_config", +] diff --git a/src/megatron/bridge/recipes/minimax/minimax_m3.py b/src/megatron/bridge/recipes/minimax/minimax_m3.py new file mode 100644 index 0000000000..5ba3465056 --- /dev/null +++ b/src/megatron/bridge/recipes/minimax/minimax_m3.py @@ -0,0 +1,31 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ruff: noqa: F401 +"""Compatibility aliases for legacy recipe names.""" + +from __future__ import annotations + +from megatron.bridge.recipes.minimax.h100.minimax_m3 import MINIMAX_M3_HF_PATH +from megatron.bridge.recipes.minimax.h100.minimax_m3 import ( + minimax_m3_pretrain_256gpu_h100_bf16_config as minimax_m3_pretrain_config, +) +from megatron.bridge.recipes.minimax.h100.minimax_m3 import ( + minimax_m3_sft_128gpu_h100_bf16_config as minimax_m3_sft_config, +) + + +__all__ = [ + "minimax_m3_pretrain_config", + "minimax_m3_sft_config", +] diff --git a/src/megatron/bridge/training/checkpointing.py b/src/megatron/bridge/training/checkpointing.py index 22f372f429..5365cba4c0 100644 --- a/src/megatron/bridge/training/checkpointing.py +++ b/src/megatron/bridge/training/checkpointing.py @@ -879,6 +879,11 @@ def _save_hf_weights( hf_source = _resolve_hf_source(cfg) bridge = _build_auto_bridge_for_save(cfg, hf_source=hf_source) + model_bridge = bridge._model_bridge + if not model_bridge.SUPPORTS_HF_PRETRAINED_EXPORT: + raise NotImplementedError( + f"{type(model_bridge).__name__} does not support standalone Hugging Face checkpoint export." + ) distributed_save = bool(getattr(ckpt_cfg, "hf_distributed_save", False)) save_every_n_ranks = int(getattr(ckpt_cfg, "hf_save_every_n_ranks", 1)) diff --git a/tests/unit_tests/examples/test_hf_to_megatron_generate_text.py b/tests/unit_tests/examples/test_hf_to_megatron_generate_text.py new file mode 100644 index 0000000000..26dd29defd --- /dev/null +++ b/tests/unit_tests/examples/test_hf_to_megatron_generate_text.py @@ -0,0 +1,68 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import runpy +from pathlib import Path +from unittest.mock import MagicMock + +import pytest +import torch + + +_SCRIPT = Path(__file__).parents[3] / "examples" / "conversion" / "hf_to_megatron_generate_text.py" +_SCRIPT_GLOBALS = runpy.run_path(_SCRIPT) +_decode_completion = _SCRIPT_GLOBALS["_decode_completion"] +_tokenize_prompt = _SCRIPT_GLOBALS["_tokenize_prompt"] + + +@pytest.mark.unit +def test_tokenize_raw_prompt() -> None: + tokenizer = MagicMock() + tokenizer.encode.return_value = torch.tensor([[1, 2]]) + + input_ids = _tokenize_prompt(tokenizer, "hello", apply_chat_template=False, thinking_mode="adaptive") + + torch.testing.assert_close(input_ids, torch.tensor([[1, 2]])) + tokenizer.encode.assert_called_once_with("hello", return_tensors="pt") + tokenizer.apply_chat_template.assert_not_called() + + +@pytest.mark.unit +def test_tokenize_chat_prompt() -> None: + tokenizer = MagicMock() + tokenizer.apply_chat_template.return_value = {"input_ids": torch.tensor([[3, 4]])} + + input_ids = _tokenize_prompt(tokenizer, "hello", apply_chat_template=True, thinking_mode="disabled") + + torch.testing.assert_close(input_ids, torch.tensor([[3, 4]])) + tokenizer.apply_chat_template.assert_called_once_with( + [{"role": "user", "content": "hello"}], + add_generation_prompt=True, + tokenize=True, + return_dict=True, + return_tensors="pt", + thinking_mode="disabled", + ) + tokenizer.encode.assert_not_called() + + +@pytest.mark.unit +def test_decode_completion_excludes_prompt_and_special_tokens() -> None: + tokenizer = MagicMock() + tokenizer.decode.return_value = "The sky appears blue." + + text = _decode_completion(tokenizer, torch.tensor([[10, 11, 20, 21]]), prompt_length=2) + + assert text == "The sky appears blue." + tokenizer.decode.assert_called_once_with([20, 21], skip_special_tokens=True) diff --git a/tests/unit_tests/models/minimax_m3/__init__.py b/tests/unit_tests/models/minimax_m3/__init__.py new file mode 100644 index 0000000000..4fc25d0d3c --- /dev/null +++ b/tests/unit_tests/models/minimax_m3/__init__.py @@ -0,0 +1,13 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/tests/unit_tests/models/minimax_m3/test_minimax_m3_bridge.py b/tests/unit_tests/models/minimax_m3/test_minimax_m3_bridge.py new file mode 100644 index 0000000000..9d986c125a --- /dev/null +++ b/tests/unit_tests/models/minimax_m3/test_minimax_m3_bridge.py @@ -0,0 +1,343 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Unit tests for the MiniMax-M3 bridge. +""" + +from unittest.mock import Mock + +import pytest +import torch +from transformers import GenerationConfig, PretrainedConfig + +from megatron.bridge.models.conversion.auto_bridge import AutoBridge +from megatron.bridge.models.conversion.model_bridge import MegatronModelBridge +from megatron.bridge.models.hf_pretrained.causal_lm import PreTrainedCausalLM +from megatron.bridge.models.minimax_m3.minimax_m3_bridge import ( + MiniMaxM3Bridge, + MiniMaxM3ModelProvider, + TopKRouter, + _promote_router_weights_to_float32, + quick_gelu, +) +from megatron.bridge.models.model_provider import _apply_mixed_precision_wrapper + + +# Toy text-backbone config (mirrors the shape of MiniMaxAI/MiniMax-M3 text_config) +_MINIMAX_M3_TEXT_CONFIG = { + "architectures": ["MiniMaxM3SparseForCausalLM"], + "hidden_size": 64, + "intermediate_size": 32, + "dense_intermediate_size": 128, + "shared_intermediate_size": 32, + "n_shared_experts": 1, + "num_hidden_layers": 4, + "num_attention_heads": 8, + "num_key_value_heads": 4, + "head_dim": 16, + "hidden_act": "swigluoai", + "max_position_embeddings": 4096, + "rms_norm_eps": 1e-06, + "rope_theta": 5000000.0, + "rotary_dim": 8, + "partial_rotary_factor": 0.5, + "vocab_size": 1024, + "tie_word_embeddings": False, + "attention_dropout": 0.0, + "num_local_experts": 8, + "num_experts_per_tok": 2, + "scoring_func": "sigmoid", + "use_routing_bias": True, + "use_qk_norm": True, + "use_gemma_norm": True, + "qk_norm_type": "per_head", + "routed_scaling_factor": 2.0, + "router_aux_loss_coef": 0.001, + "moe_layer_freq": [0, 1, 1, 1], + "num_nextn_predict_layers": 1, + "swiglu_alpha": 1.702, + "swiglu_limit": 7.0, + "torch_dtype": "bfloat16", +} + +_MINIMAX_M3_VL_CONFIG = { + "architectures": ["MiniMaxM3SparseForConditionalGeneration"], + "model_type": "minimax_m3_vl", + "tie_word_embeddings": False, + "torch_dtype": "bfloat16", +} + + +_DELETE = object() + + +def _make_text_config(overrides: dict | None = None) -> Mock: + values = dict(_MINIMAX_M3_TEXT_CONFIG) + if overrides: + for key, value in overrides.items(): + if value is _DELETE: + values.pop(key, None) + else: + values[key] = value + cfg = Mock(spec=list(values.keys())) + for k, v in values.items(): + setattr(cfg, k, v) + return cfg + + +def _make_pretrained(text_overrides: dict | None = None) -> Mock: + text_cfg = _make_text_config(text_overrides) + + outer_keys = list(_MINIMAX_M3_VL_CONFIG.keys()) + ["text_config"] + outer_cfg = Mock(spec=outer_keys) + for k, v in _MINIMAX_M3_VL_CONFIG.items(): + setattr(outer_cfg, k, v) + outer_cfg.text_config = text_cfg + + m = Mock(spec=PreTrainedCausalLM) + m.config = outer_cfg + m.generation_config = Mock(spec=GenerationConfig) + return m + + +class TestMiniMaxM3Bridge: + """Unit tests for MiniMaxM3Bridge config mapping and mapping registry.""" + + @pytest.fixture + def mock_pretrained(self): + return _make_pretrained() + + def test_registration(self): + assert issubclass(MiniMaxM3Bridge, MegatronModelBridge) + + def test_provider_bridge_maps_core_config(self, mock_pretrained): + bridge = MiniMaxM3Bridge() + provider = bridge.provider_bridge(mock_pretrained) + + text_config = mock_pretrained.config.text_config + assert provider.hidden_size == text_config.hidden_size + assert provider.num_layers == text_config.num_hidden_layers + assert provider.num_attention_heads == text_config.num_attention_heads + assert provider.num_query_groups == text_config.num_key_value_heads + assert provider.kv_channels == text_config.head_dim + assert provider.vocab_size == text_config.vocab_size + assert provider.layernorm_epsilon == text_config.rms_norm_eps + assert provider.rotary_base == text_config.rope_theta + assert provider.num_moe_experts == text_config.num_local_experts + assert provider.moe_router_topk == text_config.num_experts_per_tok + assert provider.share_embeddings_and_output_weights is False + + def test_provider_bridge_splits_dense_and_expert_ffn_sizes(self, mock_pretrained): + """intermediate_size is the per-expert size; dense layers use dense_intermediate_size.""" + bridge = MiniMaxM3Bridge() + provider = bridge.provider_bridge(mock_pretrained) + + text_config = mock_pretrained.config.text_config + assert provider.ffn_hidden_size == text_config.dense_intermediate_size + assert provider.moe_ffn_hidden_size == text_config.intermediate_size + + def test_provider_bridge_sets_moe_sigmoid_routing(self, mock_pretrained): + bridge = MiniMaxM3Bridge() + provider = bridge.provider_bridge(mock_pretrained) + + assert provider.moe_grouped_gemm is True + assert provider.moe_router_pre_softmax is False + assert provider.moe_router_score_function == "sigmoid" + assert provider.moe_router_enable_expert_bias is True + assert provider.moe_token_dispatcher_type == "alltoall" + assert provider.moe_router_load_balancing_type == "aux_loss" + assert provider.moe_router_topk_scaling_factor == 2.0 + assert provider.moe_shared_expert_intermediate_size == 32 + assert provider.moe_shared_expert_overlap is False + assert provider.moe_layer_freq == [0, 1, 1, 1] + + def test_provider_bridge_derives_moe_layer_freq_from_mlp_layer_types(self): + """The native transformers config exposes mlp_layer_types instead of moe_layer_freq.""" + mock_pretrained = _make_pretrained( + {"moe_layer_freq": _DELETE, "mlp_layer_types": ["dense", "sparse", "sparse", "sparse"]} + ) + bridge = MiniMaxM3Bridge() + provider = bridge.provider_bridge(mock_pretrained) + + assert provider.moe_layer_freq == [0, 1, 1, 1] + + def test_provider_bridge_disables_mtp(self, mock_pretrained): + """The checkpoint config advertises MTP layers but ships no mtp.* weights.""" + bridge = MiniMaxM3Bridge() + provider = bridge.provider_bridge(mock_pretrained) + + assert provider.mtp_num_layers is None + + def test_provider_bridge_sets_gemma_style_norm(self, mock_pretrained): + bridge = MiniMaxM3Bridge() + provider = bridge.provider_bridge(mock_pretrained) + + assert provider.normalization == "RMSNorm" + assert provider.layernorm_zero_centered_gamma is True + assert provider.qk_layernorm is True + + def test_provider_bridge_sets_swigluoai_activation(self, mock_pretrained): + bridge = MiniMaxM3Bridge() + provider = bridge.provider_bridge(mock_pretrained) + + assert provider.gated_linear_unit is True + assert provider.activation_func is quick_gelu + assert provider.activation_func_clamp_value == 7.0 + assert provider.glu_linear_offset == 1.0 + assert provider.bias_activation_fusion is False + + def test_hf_to_megatron_activation_swigluoai(self): + assert MiniMaxM3Bridge.hf_to_megatron_activation("swigluoai") is quick_gelu + + def test_hf_to_megatron_activation_unknown_raises(self): + with pytest.raises(ValueError): + MiniMaxM3Bridge.hf_to_megatron_activation("not_a_real_activation") + + def test_provider_bridge_calculates_rotary_percent(self, mock_pretrained): + bridge = MiniMaxM3Bridge() + provider = bridge.provider_bridge(mock_pretrained) + + text_config = mock_pretrained.config.text_config + expected = text_config.rotary_dim / text_config.head_dim + assert abs(provider.rotary_percent - expected) < 1e-6 + + def test_provider_bridge_rotary_percent_missing_fields(self): + """When rotary_dim is absent, no AttributeError is raised.""" + mock_pretrained = _make_pretrained({"rotary_dim": None}) + bridge = MiniMaxM3Bridge() + provider = bridge.provider_bridge(mock_pretrained) + assert hasattr(provider, "rotary_percent") + + def test_provider_bridge_dtype_bfloat16(self, mock_pretrained): + bridge = MiniMaxM3Bridge() + provider = bridge.provider_bridge(mock_pretrained) + + assert provider.bf16 is True + assert provider.fp16 is False + assert provider.params_dtype == torch.bfloat16 + assert provider.autocast_dtype == torch.bfloat16 + + def test_provider_bridge_caps_seq_length(self, mock_pretrained): + bridge = MiniMaxM3Bridge() + provider = bridge.provider_bridge(mock_pretrained) + + assert provider.seq_length == 4096 + + def test_provider_bridge_flat_text_config(self): + """A flat (non-nested) text config is accepted for text-only checkpoints.""" + text_cfg = _make_text_config() + m = Mock(spec=PreTrainedCausalLM) + m.config = text_cfg + m.generation_config = Mock(spec=GenerationConfig) + + bridge = MiniMaxM3Bridge() + provider = bridge.provider_bridge(m) + assert provider.hidden_size == text_cfg.hidden_size + + def test_mapping_registry_contains_critical_weights(self): + bridge = MiniMaxM3Bridge() + registry = bridge.mapping_registry() + + megatron_params = [str(m.megatron_param) for m in registry] + assert any("word_embeddings" in p for p in megatron_params), "Embedding mapping missing" + assert any("output_layer" in p for p in megatron_params), "LM head mapping missing" + assert any("linear_qkv.weight" in p for p in megatron_params), "QKV mapping missing" + assert any("linear_proj" in p for p in megatron_params), "o_proj mapping missing" + assert any("q_layernorm" in p for p in megatron_params), "Q norm mapping missing" + assert any("k_layernorm" in p for p in megatron_params), "K norm mapping missing" + assert any("mlp.router.weight" in p for p in megatron_params), "MoE router mapping missing" + assert any("mlp.router.expert_bias" in p for p in megatron_params), "Expert bias mapping missing" + assert any("mlp.experts.linear_fc1" in p for p in megatron_params), "Expert gate/up mapping missing" + assert any("mlp.experts.linear_fc2" in p for p in megatron_params), "Expert down mapping missing" + assert any("mlp.shared_experts.linear_fc1" in p for p in megatron_params), "Shared expert mapping missing" + assert any("mlp.linear_fc1.weight" in p for p in megatron_params), "Dense MLP mapping missing" + + def test_router_hook_preserves_float32_during_native_state_reload(self): + router = TopKRouter.__new__(TopKRouter) + torch.nn.Module.__init__(router) + router.weight = torch.nn.Parameter(torch.empty(8, 64, dtype=torch.bfloat16)) + model = torch.nn.Module() + model.router = router + source_weight = torch.randn(8, 64, dtype=torch.float32) + + _promote_router_weights_to_float32([model]) + model.load_state_dict({"router.weight": source_weight}) + + assert router.weight.dtype == torch.float32 + assert torch.equal(router.weight, source_weight) + + def test_router_weight_survives_mixed_precision_wrapper(self): + router = TopKRouter.__new__(TopKRouter) + torch.nn.Module.__init__(router) + router.weight = torch.nn.Parameter(torch.empty(8, 64, dtype=torch.bfloat16)) + model = torch.nn.Module() + model.router = router + source_weight = torch.randn(8, 64, dtype=torch.float32) + + _promote_router_weights_to_float32([model]) + model.load_state_dict({"router.weight": source_weight}) + wrapped = _apply_mixed_precision_wrapper([model], object(), lambda _config, module: module.bfloat16()) + + assert wrapped == [model] + assert router.weight.dtype == torch.float32 + assert torch.equal(router.weight, source_weight) + + def test_provider_registers_router_dtype_hook_first(self, mock_pretrained): + provider = MiniMaxM3Bridge().provider_bridge(mock_pretrained) + + assert isinstance(provider, MiniMaxM3ModelProvider) + assert provider._pre_wrap_hooks[0] is _promote_router_weights_to_float32 + + def test_mapping_registry_uses_language_model_prefix(self): + """All HF-side params must live under the multimodal language_model. prefix.""" + bridge = MiniMaxM3Bridge() + registry = bridge.mapping_registry() + + for mapping in registry: + hf_param = mapping.hf_param + hf_params = hf_param.values() if isinstance(hf_param, dict) else [hf_param] + for p in hf_params: + assert str(p).startswith("language_model."), f"HF param missing language_model. prefix: {p}" + + def test_mapping_registry_does_not_map_indexer_or_vision(self): + """Lightning-indexer and vision-tower weights are intentionally unmapped.""" + bridge = MiniMaxM3Bridge() + registry = bridge.mapping_registry() + + for mapping in registry: + hf_param = mapping.hf_param + hf_params = hf_param.values() if isinstance(hf_param, dict) else [hf_param] + for p in hf_params: + assert "index_" not in str(p) + assert "vision_tower" not in str(p) + + def test_megatron_to_hf_config_rejects_incomplete_multimodal_export(self, mock_pretrained): + bridge = MiniMaxM3Bridge() + provider = bridge.provider_bridge(mock_pretrained) + + with pytest.raises(NotImplementedError, match="multimodal"): + MiniMaxM3Bridge.megatron_to_hf_config(provider) + + def test_public_hf_save_rejects_before_creating_output(self, tmp_path): + hf_config = PretrainedConfig() + hf_config.architectures = ["MiniMaxM3SparseForConditionalGeneration"] + hf_config.model_type = "minimax_m3_vl" + auto_bridge = AutoBridge(hf_config) + output_path = tmp_path / "incomplete-hf-export" + + with pytest.raises(NotImplementedError, match="standalone Hugging Face"): + auto_bridge.save_hf_pretrained([], output_path) + + assert not output_path.exists() diff --git a/tests/unit_tests/models/test_autobridge_registration_matrix.py b/tests/unit_tests/models/test_autobridge_registration_matrix.py index f1889ebe20..05458010c7 100644 --- a/tests/unit_tests/models/test_autobridge_registration_matrix.py +++ b/tests/unit_tests/models/test_autobridge_registration_matrix.py @@ -57,6 +57,7 @@ "MiMoForCausalLM": "megatron.bridge.models.mimo.mimo_bridge.MimoBridge", "MiMoV2FlashForCausalLM": ("megatron.bridge.models.mimo_v2_flash.mimo_v2_flash_bridge.MiMoV2FlashBridge"), "MiniMaxM2ForCausalLM": "megatron.bridge.models.minimax_m2.minimax_m2_bridge.MiniMaxM2Bridge", + "MiniMaxM3SparseForConditionalGeneration": ("megatron.bridge.models.minimax_m3.minimax_m3_bridge.MiniMaxM3Bridge"), "Mistral3ForConditionalGeneration": ("megatron.bridge.models.ministral3.ministral3_bridge.Ministral3Bridge"), "MistralForCausalLM": "megatron.bridge.models.mistral.mistral_bridge.MistralBridge", "NemotronForCausalLM": "megatron.bridge.models.nemotron.nemotron_bridge.NemotronBridge", @@ -110,6 +111,7 @@ "MiMoForCausalLM", "MiMoV2FlashForCausalLM", "MiniMaxM2ForCausalLM", + "MiniMaxM3SparseForConditionalGeneration", "NemotronHForCausalLM", "NemotronH_Nano_Omni_Reasoning_V3", "NemotronH_Nano_VL_V2", diff --git a/tests/unit_tests/training/test_checkpointing.py b/tests/unit_tests/training/test_checkpointing.py index ceb6a4ce5f..4054d7966c 100644 --- a/tests/unit_tests/training/test_checkpointing.py +++ b/tests/unit_tests/training/test_checkpointing.py @@ -39,6 +39,7 @@ _load_hf_pretrained_checkpoint, _load_model_state_dict, _save_hf_adapter_weights, + _save_hf_weights, checkpoint_exists, cleanup_old_non_persistent_checkpoint, create_checkpoint_manager, @@ -150,6 +151,22 @@ def test_get_checkpoint_tracker_filename(self): expected = "/checkpoints/latest_checkpointed_iteration.txt" assert result == expected + @patch("megatron.bridge.training.checkpointing.ensure_directory_exists") + @patch("megatron.bridge.training.checkpointing._build_auto_bridge_for_save") + def test_save_hf_weights_rejects_unsupported_bridge_before_writing(self, mock_build_bridge, mock_ensure_dir): + state = Mock() + state.cfg.peft = None + state.cfg.checkpoint.hf_source_path = "/hf/source" + bridge = mock_build_bridge.return_value + bridge._model_bridge.SUPPORTS_HF_PRETRAINED_EXPORT = False + + with pytest.raises(NotImplementedError, match="standalone Hugging Face"): + _save_hf_weights(state, [], "/checkpoints/iter_0000001/hf") + + mock_ensure_dir.assert_not_called() + bridge.export_hf_weights.assert_not_called() + bridge.hf_pretrained.save_artifacts.assert_not_called() + @patch("torch.distributed.is_initialized") @patch("torch.distributed.get_rank") @patch("torch.distributed.all_reduce")