Skip to content

Switch Qwen3-Next to use MambaModel - #2520

Closed
Phlip79 wants to merge 22 commits into
mainfrom
philip/qwen-on-mamba
Closed

Switch Qwen3-Next to use MambaModel#2520
Phlip79 wants to merge 22 commits into
mainfrom
philip/qwen-on-mamba

Conversation

@Phlip79

@Phlip79 Phlip79 commented Feb 25, 2026

Copy link
Copy Markdown
Member

What does this PR do ?

Switch Qwen3-Next recipe to use MambaModel instead of GPTModel.

This PR depends on the Megatron-LM PRs: NVIDIA/Megatron-LM#3535 and NVIDIA/Megatron-LM#4067.

Changelog

  • Add specific line by line info of high level changes in this PR.

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:

  • Make sure you read and followed Contributor guidelines
  • Did you write any new necessary tests?
  • Did you add or update any necessary documentation?
  • Does the PR affect components that are optional to install? (Ex: Numba, Pynini, Apex etc)
    • Reviewer: Does the PR have correct import guards for all optional libraries?

If you haven't finished some of the above items you can still open "Draft" PR.

Additional Information

  • Related to # (issue)

Summary by CodeRabbit

  • New Features

    • Added Docker container support for DeepEP integration workflows
    • Added example scripts for Qwen3-Next model conversion, pretraining, and supervised fine-tuning
    • Enhanced Qwen3-Next model bridge with Mamba architecture and Multi-Token Prediction support
  • Tests

    • Added comprehensive tests for Qwen3 bridge layer indexing, mapping validation, and architecture behavior

@copy-pr-bot

copy-pr-bot Bot commented Feb 25, 2026

Copy link
Copy Markdown

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

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

@Phlip79

Phlip79 commented Mar 31, 2026

Copy link
Copy Markdown
Member Author

/claude review

Comment thread src/megatron/bridge/models/qwen/qwen3_next_bridge.py
Comment thread src/megatron/bridge/models/qwen/qwen3_next_bridge.py Outdated
Phlip79 and others added 4 commits March 31, 2026 02:02
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>
@Phlip79

Phlip79 commented Mar 31, 2026

Copy link
Copy Markdown
Member Author

/claude review

Comment thread src/megatron/bridge/models/qwen/qwen3_next_bridge.py
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>
@Phlip79

Phlip79 commented Mar 31, 2026

Copy link
Copy Markdown
Member Author

/claude review

@claude claude Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM

@Phlip79

Phlip79 commented Mar 31, 2026

Copy link
Copy Markdown
Member Author

/ok to test a4a841d

@Phlip79
Phlip79 marked this pull request as ready for review April 1, 2026 03:37
@Phlip79
Phlip79 requested a review from a team as a code owner April 1, 2026 03:37
@Phlip79
Phlip79 requested review from yaoyu-33 and removed request for yaoyu-33 April 1, 2026 03:37
@Phlip79
Phlip79 marked this pull request as draft April 1, 2026 03:38
@coderabbitai

coderabbitai Bot commented Apr 1, 2026

Copy link
Copy Markdown
Contributor
📝 Walkthrough

Walkthrough

This 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

Cohort / File(s) Summary
Deep EP Container
docker/Dockerfile.deepep
Creates NVIDIA PyTorch-based container with uv package manager, installs project dependencies, patches and builds DeepEP repository from branch hybrid-ep with CUDA architecture configuration.
Qwen3-Next Training & Conversion Scripts
examples/models/qwen3_next/conversion.sh, examples/models/qwen3_next/slurm_pretrain.sh, examples/models/qwen3_next/slurm_sft.sh
Automation scripts for model checkpoint conversion (HF↔Megatron), SLURM-based pretraining with configurable parallelism configs (TP/PP/EP/CP/SP), and supervised fine-tuning with per-config overrides and error handling.
Qwen3-Next Bridge Implementation
src/megatron/bridge/models/qwen/qwen3_next_bridge.py
Switches bridge target from GPTModel to MambaModel with MambaModelProvider; adds hybrid attention pattern detection from HF config, MTP (Multi-Token Prediction) safetensors-based detection, and per-layer mapping logic for physical Mamba layers with conditional GDN/standard attention and MoE FFN routing.
Qwen3-Next Recipe Configuration
src/megatron/bridge/recipes/qwen/qwen3_next.py
Introduces _HF_MODEL_ID constant, replaces hardcoded model IDs, adds MTP parameters to SFT config (mtp_num_layers, mtp_loss_scaling_factor), and adjusts initialization/parallelism settings.
Qwen3-Next Bridge Tests
tests/unit_tests/models/qwen/test_qwen3_next_bridge.py
Adds TestQwen3NextMambaLayerIndexing test class validating hybrid override patterns, layer count calculations, MTP detection, physical-to-HF layer mapping rules, and MTP-specific parameter prefixing under mtp.layers paths.

Sequence Diagram

sequenceDiagram
    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()
Loading

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~50 minutes

🚥 Pre-merge checks | ✅ 3 | ❌ 1

❌ Failed checks (1 warning)

Check name Status Explanation Resolution
Test Results For Major Changes ⚠️ Warning PR makes major architectural change (GPTModel to MambaModel) but lacks documented test results validating convergence, numerical equivalence, or performance impact. Embed W&B report links with loss curves and metrics in PR description, or add integration tests demonstrating convergence behavior on new MambaModel architecture.
✅ Passed checks (3 passed)
Check name Status Explanation
Description Check ✅ Passed Check skipped - CodeRabbit’s high-level summary is enabled.
Title check ✅ Passed The title 'Switch Qwen3-Next to use MambaModel' directly and concisely describes the main change across the PR, which switches the Qwen3-Next bridge architecture from GPTModel to MambaModel.
Docstring Coverage ✅ Passed Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%.

✏️ Tip: You can configure your own custom pre-merge checks in the settings.

✨ Finishing Touches
📝 Generate docstrings
  • Create stacked PR
  • Commit on current branch
🧪 Generate unit tests (beta)
  • Create PR with unit tests
  • Commit unit tests in branch philip/qwen-on-mamba

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.

❤️ Share

Comment @coderabbitai help to get the list of available commands and usage tips.

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 | 🟡 Minor

Update line 381 to use final_norm instead of final_layernorm.

The bridge maps decoder.final_norm.weight (line 249 of qwen3_next_bridge.py), but this test asserts for decoder.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 eb9cee7de5a24193bf09500668d3a619d3d3f3fb is 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 Exception hides potential programming errors. Consider catching specific exceptions like FileNotFoundError, json.JSONDecodeError, and HfHubHTTPError.

♻️ 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 for hf_pretrained parameter.

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_CHECKING to 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 at src/megatron/bridge/recipes/qwen/qwen3_next.py:179 defaults to pipeline_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

📥 Commits

Reviewing files that changed from the base of the PR and between 23ac10e and a4a841d.

📒 Files selected for processing (7)
  • docker/Dockerfile.deepep
  • examples/models/qwen3_next/conversion.sh
  • examples/models/qwen3_next/slurm_pretrain.sh
  • examples/models/qwen3_next/slurm_sft.sh
  • src/megatron/bridge/models/qwen/qwen3_next_bridge.py
  • src/megatron/bridge/recipes/qwen/qwen3_next.py
  • tests/unit_tests/models/qwen/test_qwen3_next_bridge.py

Comment on lines +36 to +42
# 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

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟡 Minor

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.

Suggested change
# 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 \

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟡 Minor

🧩 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.py

Repository: 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 -30

Repository: 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.py

Repository: 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.py

Repository: 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 -50

Repository: 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.

Comment on lines +691 to +694
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)

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟡 Minor

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.

Suggested change
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.

@yaoyu-33

Copy link
Copy Markdown
Contributor

Closing in favor of the newer HybridModel migration stack in #4836, which implements Qwen3-Next and Qwen3.5 on the current model architecture.

@yaoyu-33 yaoyu-33 closed this Jul 24, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants