Switch Qwen3-Next to use MambaModel - #2520
Conversation
…into philip/qwen-on-mamba
…into philip/qwen-on-mamba
…into philip/qwen-on-mamba
…into philip/qwen-on-mamba
|
/claude review |
QK norm is now handled in SelfAttention.__init__ via config-driven fallback (Megatron-LM PR #4067), so Bridge no longer needs to modify the mamba stack spec. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Verify the 2*N (attention) / 2*N+1 (MLP) physical layer mapping, hybrid_override_pattern generation, GDN vs standard attention placement, MTP key paths, and decoder.final_norm naming. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
|
/claude review |
Wrap hf_hub_download in try/except so models without model.safetensors.index.json (single-file safetensors, .bin format) gracefully return False instead of crashing during conversion. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
|
/claude review |
|
/ok to test a4a841d |
📝 WalkthroughWalkthroughThis PR adds Mamba-based model bridge support for Qwen3-Next 80B-A3B, including a Deep EP container build, training and conversion automation scripts, updated bridge configuration with layer mapping logic, recipe constants, and comprehensive test coverage for the new Mamba bridge implementation. Changes
Sequence DiagramsequenceDiagram
participant HF as HuggingFace<br/>Qwen3-Next Model
participant Bridge as Qwen3NextBridge<br/>(MambaModelProvider)
participant Detector as MTP Detector<br/>(safetensors-index)
participant Mapper as Layer Mapper<br/>(Hybrid Logic)
participant Mamba as Megatron<br/>MambaModel
HF->>Bridge: from_hf_pretrained(model_id)
Bridge->>Detector: Check for MTP in safetensors index
Detector-->>Bridge: MTP detected?
Bridge->>Mapper: Build hybrid_override_pattern<br/>from full_attention_interval
Mapper->>Mapper: Compute num_layers = HF_layers × 2
Mapper->>Mapper: Map HF layer N to:<br/>Physical Attention: 2N<br/>Physical MoE FFN: 2N+1
Mapper-->>Bridge: Mapping registry with<br/>GDN/standard attention splits
Bridge->>Mamba: Create MambaModel with<br/>hybrid config + MTP params
Mamba-->>Bridge: to_megatron_provider()
Estimated code review effort🎯 4 (Complex) | ⏱️ ~50 minutes 🚥 Pre-merge checks | ✅ 3 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (3 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
There was a problem hiding this comment.
Actionable comments posted: 3
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
tests/unit_tests/models/qwen/test_qwen3_next_bridge.py (1)
379-381:⚠️ Potential issue | 🟡 MinorUpdate line 381 to use
final_norminstead offinal_layernorm.The bridge maps
decoder.final_norm.weight(line 249 of qwen3_next_bridge.py), but this test asserts fordecoder.final_layernorm.weight, which contradicts the expectation at lines 724-725 and the actual bridge implementation. The assertion should be:assert "decoder.final_norm.weight" in megatron_params🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tests/unit_tests/models/qwen/test_qwen3_next_bridge.py` around lines 379 - 381, The test asserts the wrong target name for the final layer norm: update the assertion that checks megatron_params to expect "decoder.final_norm.weight" instead of "decoder.final_layernorm.weight" so it matches the mapping produced by qwen3_next_bridge.py; locate the assertion using the variables hf_params and megatron_params in tests/unit_tests/models/qwen/test_qwen3_next_bridge.py and replace the string accordingly.
🧹 Nitpick comments (5)
docker/Dockerfile.deepep (2)
65-67: Consider documenting the pinned DeepEP commit purpose.The commit
eb9cee7de5a24193bf09500668d3a619d3d3f3fbis pinned without explanation. Adding a comment about what this commit provides (e.g., specific fix, feature, or compatibility) would help maintainers understand when/if to update it.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@docker/Dockerfile.deepep` around lines 65 - 67, Add a short comment above the git clone/checkout lines documenting why the DeepEP commit eb9cee7de5a24193bf09500668d3d3f3fb is pinned: state the specific fix/feature or compatibility it provides (e.g., "pins fix for X bug" or "ensures compatibility with Y version"), include the commit hash and optionally a link or reference to the upstream PR/issue, and mention when it can be revisited; update the comment near the existing git clone --branch hybrid-ep and git checkout eb9cee7de5a24193bf09500668d3d3f3fb lines so future maintainers understand the rationale.
61-63: Hardcoded Python version path may break with base image updates.The path
/opt/venv/lib/python3.12/site-packages/assumes Python 3.12. If the base image is updated to a different Python version, this symlink creation will fail silently or point to the wrong location.♻️ Proposed fix using dynamic Python version
- pushd /opt/venv/lib/python3.12/site-packages/nvidia/nvshmem/lib/ + PYTHON_VERSION=$(python -c "import sys; print(f'{sys.version_info.major}.{sys.version_info.minor}')") + pushd /opt/venv/lib/python${PYTHON_VERSION}/site-packages/nvidia/nvshmem/lib/ ln -s libnvshmem_host.so.3 libnvshmem_host.so popd🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@docker/Dockerfile.deepep` around lines 61 - 63, The Dockerfile is using a hardcoded path /opt/venv/lib/python3.12/site-packages/ when creating the nvshmem symlink (the pushd / popd + ln -s block), which will break if the Python minor version changes; update this step to dynamically discover the virtualenv's site-packages directory at build time (e.g., call Python to emit the correct site-packages path and use that value for pushd) before creating the symlink, and add a guard to verify the target libnvshmem_host.so.3 exists and fail with an explanatory message if not found.src/megatron/bridge/models/qwen/qwen3_next_bridge.py (2)
147-150: Narrow the exception clause to specific exception types.Catching bare
Exceptionhides potential programming errors. Consider catching specific exceptions likeFileNotFoundError,json.JSONDecodeError, andHfHubHTTPError.♻️ Proposed fix
+ from huggingface_hub.utils import HfHubHTTPError + # Try local path first, then download from hub try: local_path = Path(str(model_id)) / "model.safetensors.index.json" if local_path.exists(): index_path = local_path else: index_path = Path(hf_hub_download(str(model_id), "model.safetensors.index.json")) with open(index_path) as f: weight_map = json.load(f).get("weight_map", {}) return any(k.startswith("mtp.") for k in weight_map) - except Exception: + except (FileNotFoundError, json.JSONDecodeError, HfHubHTTPError, OSError): return False🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@src/megatron/bridge/models/qwen/qwen3_next_bridge.py` around lines 147 - 150, The current broad except in the weight_map loading block should be narrowed: catch FileNotFoundError for missing files, json.JSONDecodeError for invalid JSON, and HfHubHTTPError for HF hub fetch errors (import HfHubHTTPError from huggingface_hub), then return False in those cases; leave other exceptions to surface. Update the try/except around json.load(f) and the subsequent any(k.startswith("mtp.") ...) to specifically handle these three exception types and keep the same return False behavior for them.
123-124: Add type hint forhf_pretrainedparameter.The method has a return type hint but lacks a parameter type hint, which reduces clarity.
♻️ Proposed fix
+ from megatron.bridge.models.hf_pretrained.causal_lm import PreTrainedCausalLM + `@staticmethod` - def _hf_model_has_mtp(hf_pretrained) -> bool: + def _hf_model_has_mtp(hf_pretrained: PreTrainedCausalLM) -> bool:Note: The import may need to be placed at the module level or use
TYPE_CHECKINGto avoid circular imports.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@src/megatron/bridge/models/qwen/qwen3_next_bridge.py` around lines 123 - 124, The static method _hf_model_has_mtp is missing a parameter type for hf_pretrained; add a type hint such as transformers.PreTrainedModel (or a more specific HF model type) for the hf_pretrained parameter and import that symbol at module scope (or inside an if TYPE_CHECKING: block to avoid circular imports), keeping the existing -> bool return hint and updating the function signature to def _hf_model_has_mtp(hf_pretrained: PreTrainedModel) -> bool and adjust imports accordingly.examples/models/qwen3_next/slurm_sft.sh (1)
66-67: Default parallelism config may not match SFT recipe recommendation.The default
PARALLELISM_CONFIGS=("1,4,8,1,False")uses PP=4, but the SFT recipe atsrc/megatron/bridge/recipes/qwen/qwen3_next.py:179defaults topipeline_model_parallel_size = 2. Consider aligning the default or documenting the difference.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@examples/models/qwen3_next/slurm_sft.sh` around lines 66 - 67, The default PARALLELISM_CONFIGS=("1,4,8,1,False") in the slurm SFT script doesn't match the SFT recipe's default pipeline_model_parallel_size (2); update PARALLELISM_CONFIGS to use PP=2 (e.g., set the second value to 2) so TP,PP,EP still multiply to total GPUs, or explicitly document in the script why PP=4 is chosen and reference the pipeline_model_parallel_size value in the recipe to avoid confusion (look for the PARALLELISM_CONFIGS variable in the script and pipeline_model_parallel_size in the qwen3_next recipe).
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@examples/models/qwen3_next/conversion.sh`:
- Around line 36-42: The round-trip validation command for
examples/conversion/hf_megatron_roundtrip_multi_gpu.py omits the expert-parallel
flag; update the uv run invocation to pass the expert parallel size (add the
--ep argument) — e.g., include --ep 8 (or use the EP env var) alongside the
existing --tp 1 --pp 4 so the script runs with TP=1, PP=4, EP=8 for correct MoE
conversion validation.
In `@examples/models/qwen3_next/slurm_pretrain.sh`:
- Line 140: Replace the deprecated configuration key train.eval_iters with
validation.eval_iters in the slurm_pretrain.sh invocation (currently setting
train.eval_iters=$EVAL_ITERS); update the argument to
validation.eval_iters=$EVAL_ITERS so the script uses the non-deprecated key
(matching the pattern used in qwen3_next.py), avoiding the deprecation warning.
In `@tests/unit_tests/models/qwen/test_qwen3_next_bridge.py`:
- Around line 691-694: The test_hf_layer_index_mapping assigns params via
self._all_megatron_params(registry) but never uses it, causing a lint error;
remove the unused assignment (or replace with a throwaway name like _params) so
the call isn't stored in params, e.g. delete or rename the variable in the
test_hf_layer_index_mapping method where bridge_with_config.mapping_registry()
and self._all_megatron_params(registry) are used.
---
Outside diff comments:
In `@tests/unit_tests/models/qwen/test_qwen3_next_bridge.py`:
- Around line 379-381: The test asserts the wrong target name for the final
layer norm: update the assertion that checks megatron_params to expect
"decoder.final_norm.weight" instead of "decoder.final_layernorm.weight" so it
matches the mapping produced by qwen3_next_bridge.py; locate the assertion using
the variables hf_params and megatron_params in
tests/unit_tests/models/qwen/test_qwen3_next_bridge.py and replace the string
accordingly.
---
Nitpick comments:
In `@docker/Dockerfile.deepep`:
- Around line 65-67: Add a short comment above the git clone/checkout lines
documenting why the DeepEP commit eb9cee7de5a24193bf09500668d3d3f3fb is pinned:
state the specific fix/feature or compatibility it provides (e.g., "pins fix for
X bug" or "ensures compatibility with Y version"), include the commit hash and
optionally a link or reference to the upstream PR/issue, and mention when it can
be revisited; update the comment near the existing git clone --branch hybrid-ep
and git checkout eb9cee7de5a24193bf09500668d3d3f3fb lines so future maintainers
understand the rationale.
- Around line 61-63: The Dockerfile is using a hardcoded path
/opt/venv/lib/python3.12/site-packages/ when creating the nvshmem symlink (the
pushd / popd + ln -s block), which will break if the Python minor version
changes; update this step to dynamically discover the virtualenv's site-packages
directory at build time (e.g., call Python to emit the correct site-packages
path and use that value for pushd) before creating the symlink, and add a guard
to verify the target libnvshmem_host.so.3 exists and fail with an explanatory
message if not found.
In `@examples/models/qwen3_next/slurm_sft.sh`:
- Around line 66-67: The default PARALLELISM_CONFIGS=("1,4,8,1,False") in the
slurm SFT script doesn't match the SFT recipe's default
pipeline_model_parallel_size (2); update PARALLELISM_CONFIGS to use PP=2 (e.g.,
set the second value to 2) so TP,PP,EP still multiply to total GPUs, or
explicitly document in the script why PP=4 is chosen and reference the
pipeline_model_parallel_size value in the recipe to avoid confusion (look for
the PARALLELISM_CONFIGS variable in the script and pipeline_model_parallel_size
in the qwen3_next recipe).
In `@src/megatron/bridge/models/qwen/qwen3_next_bridge.py`:
- Around line 147-150: The current broad except in the weight_map loading block
should be narrowed: catch FileNotFoundError for missing files,
json.JSONDecodeError for invalid JSON, and HfHubHTTPError for HF hub fetch
errors (import HfHubHTTPError from huggingface_hub), then return False in those
cases; leave other exceptions to surface. Update the try/except around
json.load(f) and the subsequent any(k.startswith("mtp.") ...) to specifically
handle these three exception types and keep the same return False behavior for
them.
- Around line 123-124: The static method _hf_model_has_mtp is missing a
parameter type for hf_pretrained; add a type hint such as
transformers.PreTrainedModel (or a more specific HF model type) for the
hf_pretrained parameter and import that symbol at module scope (or inside an if
TYPE_CHECKING: block to avoid circular imports), keeping the existing -> bool
return hint and updating the function signature to def
_hf_model_has_mtp(hf_pretrained: PreTrainedModel) -> bool and adjust imports
accordingly.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: 7900c1db-d2a9-4495-b1b8-cfeb80d53371
📒 Files selected for processing (7)
docker/Dockerfile.deepepexamples/models/qwen3_next/conversion.shexamples/models/qwen3_next/slurm_pretrain.shexamples/models/qwen3_next/slurm_sft.shsrc/megatron/bridge/models/qwen/qwen3_next_bridge.pysrc/megatron/bridge/recipes/qwen/qwen3_next.pytests/unit_tests/models/qwen/test_qwen3_next_bridge.py
| # Round-trip validation | ||
| uv run python -m torch.distributed.run --nproc_per_node=8 \ | ||
| examples/conversion/hf_megatron_roundtrip_multi_gpu.py \ | ||
| --hf-model-id $HF_MODEL_ID \ | ||
| --megatron-load-path ${WORKSPACE}/models/$MODEL_NAME/iter_0000000 \ | ||
| --tp 1 --pp 4 \ | ||
| --trust-remote-code |
There was a problem hiding this comment.
Missing --ep (expert parallel) parameter in round-trip validation.
The pretrain recipe recommends TP=1, PP=4, EP=8, but the round-trip validation only passes --tp 1 --pp 4. For proper validation of the MoE model conversion, the expert parallel size should also be specified.
🔧 Proposed fix
uv run python -m torch.distributed.run --nproc_per_node=8 \
examples/conversion/hf_megatron_roundtrip_multi_gpu.py \
--hf-model-id $HF_MODEL_ID \
--megatron-load-path ${WORKSPACE}/models/$MODEL_NAME/iter_0000000 \
- --tp 1 --pp 4 \
+ --tp 1 --pp 4 --ep 8 \
--trust-remote-code📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| # Round-trip validation | |
| uv run python -m torch.distributed.run --nproc_per_node=8 \ | |
| examples/conversion/hf_megatron_roundtrip_multi_gpu.py \ | |
| --hf-model-id $HF_MODEL_ID \ | |
| --megatron-load-path ${WORKSPACE}/models/$MODEL_NAME/iter_0000000 \ | |
| --tp 1 --pp 4 \ | |
| --trust-remote-code | |
| # Round-trip validation | |
| uv run python -m torch.distributed.run --nproc_per_node=8 \ | |
| examples/conversion/hf_megatron_roundtrip_multi_gpu.py \ | |
| --hf-model-id $HF_MODEL_ID \ | |
| --megatron-load-path ${WORKSPACE}/models/$MODEL_NAME/iter_0000000 \ | |
| --tp 1 --pp 4 --ep 8 \ | |
| --trust-remote-code |
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@examples/models/qwen3_next/conversion.sh` around lines 36 - 42, The
round-trip validation command for
examples/conversion/hf_megatron_roundtrip_multi_gpu.py omits the expert-parallel
flag; update the uv run invocation to pass the expert parallel size (add the
--ep argument) — e.g., include --ep 8 (or use the EP env var) alongside the
existing --tp 1 --pp 4 so the script runs with TP=1, PP=4, EP=8 for correct MoE
conversion validation.
| train.train_iters=$TRAIN_ITERS \ | ||
| train.global_batch_size=$GLOBAL_BATCH_SIZE \ | ||
| train.micro_batch_size=$MICRO_BATCH_SIZE \ | ||
| train.eval_iters=$EVAL_ITERS \ |
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Check which config path is used for eval_iters
rg -n "eval_iters" src/megatron/bridge/recipes/qwen/qwen3_next.py
rg -n "eval_iters" src/megatron/bridge/training/config.pyRepository: NVIDIA-NeMo/Megatron-Bridge
Length of output: 147
🏁 Script executed:
#!/bin/bash
# Find the slurm_pretrain.sh file
fd -t f "slurm_pretrain.sh" --exec cat -n {} \; 2>/dev/null | head -150
# Also check the full structure of validation and train configs
rg -A 5 "class.*Validation" src/megatron/bridge/training/config.py
rg -A 5 "class.*Train" src/megatron/bridge/training/config.py | head -30Repository: NVIDIA-NeMo/Megatron-Bridge
Length of output: 8091
🏁 Script executed:
#!/bin/bash
# Find qwen3_next slurm file
fd -t f "slurm_pretrain.sh" examples/models/qwen3_next/ 2>/dev/null
# Also check config structure for validation
rg -B 5 -A 10 "class ValidationConfig" src/megatron/bridge/training/config.pyRepository: NVIDIA-NeMo/Megatron-Bridge
Length of output: 115
🏁 Script executed:
#!/bin/bash
# Show the actual qwen3_next slurm_pretrain.sh file, especially around line 140
cat -n examples/models/qwen3_next/slurm_pretrain.sh | sed -n '130,150p'
# Check the validation config class structure
rg -B 2 -A 15 "class ValidationConfig" src/megatron/bridge/training/config.pyRepository: NVIDIA-NeMo/Megatron-Bridge
Length of output: 1374
🏁 Script executed:
#!/bin/bash
# Check the validation config class structure
rg -B 2 -A 10 "class ValidationConfig" src/megatron/bridge/training/config.py
# Also check if there's a Train config with eval_iters
rg -B 2 -A 15 "class TrainConfig" src/megatron/bridge/training/config.py | head -50Repository: NVIDIA-NeMo/Megatron-Bridge
Length of output: 53
🏁 Script executed:
#!/bin/bash
# Search for validation-related config classes more broadly
rg -i "validation|eval_iters" src/megatron/bridge/training/config.py -B 2 -A 5
# Also check the qwen3_next recipe file to see what config structure it uses
cat -n src/megatron/bridge/recipes/qwen/qwen3_next.py | grep -A 3 -B 3 "eval_iters"Repository: NVIDIA-NeMo/Megatron-Bridge
Length of output: 8246
Use validation.eval_iters instead of deprecated train.eval_iters.
Line 140 uses train.eval_iters, which is deprecated. While it will work via auto-migration to validation.eval_iters, this path triggers a deprecation warning. Use validation.eval_iters directly instead, matching the pattern in qwen3_next.py (line 201).
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@examples/models/qwen3_next/slurm_pretrain.sh` at line 140, Replace the
deprecated configuration key train.eval_iters with validation.eval_iters in the
slurm_pretrain.sh invocation (currently setting train.eval_iters=$EVAL_ITERS);
update the argument to validation.eval_iters=$EVAL_ITERS so the script uses the
non-deprecated key (matching the pattern used in qwen3_next.py), avoiding the
deprecation warning.
| def test_hf_layer_index_mapping(self, bridge_with_config): | ||
| """Test that HF layer indices are correctly mapped to physical indices.""" | ||
| registry = bridge_with_config.mapping_registry() | ||
| params = self._all_megatron_params(registry) |
There was a problem hiding this comment.
Remove unused variable params (pipeline failure).
The variable params is assigned but never used in this test method, causing the ruff F841 lint error and pipeline failure.
🔧 Proposed fix
def test_hf_layer_index_mapping(self, bridge_with_config):
"""Test that HF layer indices are correctly mapped to physical indices."""
registry = bridge_with_config.mapping_registry()
- params = self._all_megatron_params(registry)
# Build a map of HF -> physical from the auto mappings
for n in range(8):📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| def test_hf_layer_index_mapping(self, bridge_with_config): | |
| """Test that HF layer indices are correctly mapped to physical indices.""" | |
| registry = bridge_with_config.mapping_registry() | |
| params = self._all_megatron_params(registry) | |
| def test_hf_layer_index_mapping(self, bridge_with_config): | |
| """Test that HF layer indices are correctly mapped to physical indices.""" | |
| registry = bridge_with_config.mapping_registry() |
🧰 Tools
🪛 Flake8 (7.3.0)
[error] 694-694: local variable 'params' is assigned to but never used
(F841)
🪛 GitHub Actions: CICD NeMo
[error] 694-694: ruff (hook id: ruff) failed with F841: Local variable params is assigned to but never used. help: Remove assignment to unused variable params.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@tests/unit_tests/models/qwen/test_qwen3_next_bridge.py` around lines 691 -
694, The test_hf_layer_index_mapping assigns params via
self._all_megatron_params(registry) but never uses it, causing a lint error;
remove the unused assignment (or replace with a throwaway name like _params) so
the call isn't stored in params, e.g. delete or rename the variable in the
test_hf_layer_index_mapping method where bridge_with_config.mapping_registry()
and self._all_megatron_params(registry) are used.
|
Closing in favor of the newer HybridModel migration stack in #4836, which implements Qwen3-Next and Qwen3.5 on the current model architecture. |
What does this PR do ?
Switch Qwen3-Next recipe to use
MambaModelinstead ofGPTModel.This PR depends on the Megatron-LM PRs: NVIDIA/Megatron-LM#3535 and NVIDIA/Megatron-LM#4067.
Changelog
GitHub Actions CI
See the CI sectionin the Contributing doc for how to trigger the CI. A Nvidia developer will need to approve and trigger the CI for external contributors.
Before your PR is "Ready for review"
Pre checks:
If you haven't finished some of the above items you can still open "Draft" PR.
Additional Information
Summary by CodeRabbit
New Features
Tests