src/megatron/bridge/training/vlm_step.py-124-125 (1)

124-125: ⚠️ Potential issue | 🟡 Minor

Hardcoded is_first = True / is_last = True bypasses pipeline-stage detection.

This forces all ranks to load labels/loss_mask and visual inputs, regardless of their actual PP stage. If this is intentional for the multi-module VLM distributed training path, please add a comment explaining why pipeline-stage gating was removed (and whether the original is_pp_first_stage/is_pp_last_stage imports on line 21 are now dead).

examples/recipes/decentralized_pg/pretrain_qwen3_vl_simple.py-41-43 (1)

41-43: ⚠️ Potential issue | 🟡 Minor

Misleading comment: "Qwen3 4B" vs actual 30B config.

Line 42 says "Get the standard Qwen3 4B pretrain config" but line 43 uses qwen3_vl_30b_a3b_pretrain_config.

Fix

-    # Get the standard Qwen3 4B pretrain config with overrides
+    # Get the standard Qwen3-VL 30B MoE pretrain config with overrides

src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/rope.py-265-307 (1)

265-307: ⚠️ Potential issue | 🟡 Minor

Silence unused cu_seqlens warning and make RoPE-fusion failure explicit.

• apply_rotary_pos_emb_thd_absolute(..., cu_seqlens, ...) doesn’t use cu_seqlens (ruff ARG001). Rename to _cu_seqlens (or add a targeted noqa) to keep the interface but avoid lint noise.
• assert not config.apply_rope_fusion will be stripped under python -O; prefer a real exception (NotImplementedError / ValueError) if this is a hard constraint.

As per coding guidelines, "When a feature is not supported (such as audio embeddings), raise an explicit error (e.g., NotImplementedError) instead of silently ignoring the input to fail fast with a clear message." Based on learnings: "when a feature is not supported ... raise an explicit error (e.g., NotImplementedError) instead of silently ignoring".

src/megatron/bridge/training/qwen3vl_step.py-41-47 (1)

41-47: ⚠️ Potential issue | 🟡 Minor

Silence unused-argument warnings and remove the unused variable.
Ruff flags ARG001/F841 in this function.

🛠️ Proposed fix

     batch = next(data_iterator)
+    _ = use_mtp
+    _ = is_first_pp_stage
+    _ = is_last_pp_stage
@@
-        max_seqlen_in_batch = seqlens_in_batch.max().item()

Also applies to: 224-224

src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/model.py-16-18 (1)

16-18: ⚠️ Potential issue | 🟡 Minor

Align type hints with repo conventions and drop the unused import.
Use built-in generics/union types and remove the unused mpu import to satisfy lint and style rules.

🛠️ Proposed fix

-from typing import List, Dict
+ 
@@
-from megatron.core import InferenceParams, mpu, tensor_parallel
+from megatron.core import InferenceParams, tensor_parallel
@@
-    def set_input_tensor(self, input_tensor: List[Dict[str, torch.Tensor]]):
+    def set_input_tensor(self, input_tensor: list[dict[str, torch.Tensor]]):
@@
-        cp_img_num: list[int] = None,
-        images_padded: list[bool] = None,
+        cp_img_num: list[int] | None = None,
+        images_padded: list[bool] | None = None,

As per coding guidelines, Use built-in generics (list, dict, tuple) instead of typing equivalents, and use 'T | None' for nullable types instead of 'Optional[T]'.

Also applies to: 195-200, 300-301

src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/attention.py-40-41 (1)

40-41: ⚠️ Potential issue | 🟡 Minor

Fix lint issues: unused parameter and f-string without placeholders.
Ruff flags the unused rotary_pos_cos_sin argument and the f prefix on a static string.

🛠️ Proposed fix

-        rotary_pos_cos_sin: Optional[Tensor] = None,
+        _rotary_pos_cos_sin: Optional[Tensor] = None,
@@
-            raise ValueError(f"CUDA graphs must use flash decode with static batching!")
+            raise ValueError("CUDA graphs must use flash decode with static batching!")

Also applies to: 126-127

src/megatron/bridge/training/train.py-39-40 (1)

39-40: ⚠️ Potential issue | 🟡 Minor

Remove unused MultiModulePipelineCommunicator import (F401).
Lint currently flags this as unused.

🛠️ Proposed fix

-from megatron.core.pipeline_parallel.multimodule_communicator import MultiModulePipelineCommunicator

src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/model.py-152-155 (1)

152-155: ⚠️ Potential issue | 🟡 Minor

Use print_rank_0 for model logging.
This avoids duplicate logs across ranks during distributed training.

🛠️ Proposed fix

-from transformers.models.qwen3_vl.configuration_qwen3_vl import Qwen3VLConfig as Qwen3VLConfigHF
+from transformers.models.qwen3_vl.configuration_qwen3_vl import Qwen3VLConfig as Qwen3VLConfigHF
+from megatron.bridge.utils.common_utils import print_rank_0
@@
-                print(f"rank {torch.distributed.get_rank()} use hf vision model")
+                print_rank_0(f"rank {torch.distributed.get_rank()} use hf vision model")

As per coding guidelines, Use 'print_rank_0' for logging in model bridge to avoid duplicate output across ranks.

src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/utils.py-242-242 (1)

242-242: ⚠️ Potential issue | 🟡 Minor

Typo: "flaaten" → "flatten".

-    assert input_ids.dim() == 1, "input_ids should be flaaten"
+    assert input_ids.dim() == 1, "input_ids should be flattened"

src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/vision_model.py-1-24 (1)

1-24: ⚠️ Potential issue | 🟡 Minor

Missing NVIDIA copyright header and unused imports.

This file is missing the required NVIDIA copyright header. Additionally, static analysis flags several unused imports: dataclass, Union, MultimodalProjector, MegatronModule, build_module, and get_tensor_model_parallel_group_if_none.

Proposed fix

+# Copyright (c) 2025, 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 typing import Optional, Union
+from __future__ import annotations
 
 import torch
 from megatron.core import InferenceParams
 from megatron.core.models.common.vision_module.vision_module import VisionModule
-from megatron.core.models.vision.multimodal_projector import MultimodalProjector
 from megatron.core.packed_seq_params import PackedSeqParams
 from megatron.core.process_groups_config import ProcessGroupCollection
 from megatron.core.transformer.enums import ModelType
-from megatron.core.transformer.module import MegatronModule
-from megatron.core.transformer.spec_utils import ModuleSpec, build_module
-from megatron.core.utils import get_tensor_model_parallel_group_if_none
+from megatron.core.transformer.spec_utils import ModuleSpec
 from torch import nn
 from torch.nn import functional as F

As per coding guidelines: "Add NVIDIA copyright header to all Python files" and "Use T | None for nullable types instead of Optional[T]".

src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/vision_model.py-197-245 (1)

197-245: ⚠️ Potential issue | 🟡 Minor

Docstring args/return mismatch with actual signature and return value.

• Docstring refers to parameter x (Line 208) but the actual parameter is hidden_states.
• Docstring mentions packed_seq_params (Line 210) which is not in the function signature (it's computed internally on Line 235).
• The return type annotation says torch.Tensor (Line 203) but the function returns a tuple (hidden_states, deepstack_feature_lists) on Line 245.

Proposed fix

     def forward(
         self,
         hidden_states: Optional[torch.Tensor],
         grid_thw: torch.Tensor,
         inference_params: Optional[InferenceParams] = None,
         extra_block_kwargs: dict = None,
-    ) -> torch.Tensor:
+    ) -> tuple[torch.Tensor, list]:
         """Forward function of the Qwen3 Vision Model. This function passes the input tensors
         through the embedding layer and then the transformer.
 
         Args:
-            x (torch.Tensor): input image/video data of shape [n_tokens, n_dims]
+            hidden_states (torch.Tensor): input image/video data of shape [n_tokens, n_dims]
             grid_thw (torch.Tensor): the size tensor indicates grid size of each image/frame
-            packed_seq_params (PackedSeqParams): parameters to build attention mask in the backend
+            inference_params (InferenceParams, optional): inference parameters
+            extra_block_kwargs (dict, optional): additional keyword arguments for the decoder block
 
         Returns:
-            x (torch.Tensor): output after final transformer block of shape [b, s, h].
+            tuple: (hidden_states, deepstack_feature_lists)
         """

src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/utils.py-666-679 (1)

666-679: ⚠️ Potential issue | 🟡 Minor

cu_seqlens_padded is passed undivided to PackedSeqParams, but this is currently a latent issue since the returned packed_seq_params is never used.

While the code structure does create a mismatch when cp_size > 1 (buffer is sized as sum(seqlens_in_batch_padded_cpu) // cp_size but cu_seqlens_padded reflects full, undivided cumulative lengths), this is not a critical runtime bug in the current codebase. All calls to preprocess_packed_seqs() discard the returned packed_seq_params with _ (lines 437, 456, 467, 523 in model.py), and the parameter passed to the language model is always None (line 536).

The data itself is correctly divided during packing (indices are adjusted: start_idx = cu_seqlens_padded_cpu[i] // cp_size on line 650). However, if packed_seq_params were ever used downstream with CP enabled, the mismatch would cause incorrect indexing. Consider either: (1) adjust cu_seqlens_padded and max_seqlen_in_batch for the per-rank view when cp_size > 1, or (2) clarify whether this function is intended to support CP splitting and document the limitation.

scripts/performance/argument_parser.py (1)

148-148: Consider whether qwen3vl belongs as a --domain choice or should be handled via --model_family_name.

The existing domain values (llm, vlm) are generic categories, whereas qwen3vl is model-specific. This sets a precedent where each new VLM variant could require its own domain entry. If the Qwen3-VL training path diverges significantly enough from the general vlm path to justify a separate domain, this is fine — but if the differences are minor, routing through --model_family_name (or a sub-option) would keep the domain list stable.

examples/recipes/qwen_vl/finetune_qwen_vl.py (1)

93-93: Use built-in tuple instead of typing.Tuple for Python 3.10+.

Line 118 uses Tuple[argparse.Namespace, list[str]], mixing typing.Tuple with built-in list. Per coding guidelines, prefer built-in generics.

Proposed fix

-from typing import Tuple

And on line 118:

-def parse_cli_args() -> Tuple[argparse.Namespace, list[str]]:
+def parse_cli_args() -> tuple[argparse.Namespace, list[str]]:

As per coding guidelines: "Use built-in generics (list, dict, tuple) instead of typing equivalents".

src/megatron/bridge/training/utils/packed_seq_utils.py (2)

77-78: Missing type hint for use_fp8_padding parameter.

Fix

-def preprocess_packed_seqs(
-    input_ids: torch.Tensor, attention_mask: torch.Tensor, pre_process: bool = True, use_fp8_padding=False
-) -> tuple[torch.Tensor, PackedSeqParams]:
+def preprocess_packed_seqs(
+    input_ids: torch.Tensor, attention_mask: torch.Tensor | None, pre_process: bool = True, use_fp8_padding: bool = False
+) -> tuple[torch.Tensor, PackedSeqParams]:

As per coding guidelines: "Use type hints for function arguments and return types" and "Use 'T | None' for nullable types instead of 'Optional[T]'".

---

196-197: Minor: Use unpacking instead of list concatenation.

Per Ruff RUF005, prefer [batch_size, seq_len, *list(output.shape[2:])].

Fix

-    shape = [batch_size, seq_len] + list(output.shape[2:])  # 1,packed, dim -> batch_size, seq_len, dim
+    shape = [batch_size, seq_len, *output.shape[2:]]  # 1,packed, dim -> batch_size, seq_len, dim

src/megatron/bridge/models/qwen_vl/qwen3_vl_provider.py (2)

39-39: Remove unused import get_vision_model_config.

Flake8 confirms get_vision_model_config is imported but never used in this file.

Fix

-from megatron.bridge.models.qwen_vl.modelling_qwen3_vl.transformer_config import get_vision_model_config

---

139-157: Vision spec construction is duplicated between dense and MoE providers.

Lines 139-144 (dense provider) and lines 307-312 (MoE provider) construct identical vision_transformer_layer_spec and vision_patch_merger_spec objects. Consider extracting a shared helper method (e.g., on a common base class or as a module-level function) to keep these in sync.

src/megatron/bridge/training/initialize.py (1)

589-648: Dist-train PG split: cache rank-membership checks + validate world-size partitioning.

This block calls is_rank_in_pg(...) multiple times for the same collections; please cache results to avoid repeated get_process_group_ranks() calls during init.

Also recommend adding a sanity check that model_config.vision_world_size + model_config.language_world_size == torch.distributed.get_world_size() (or get_world_size_safe()), otherwise you can end up with “neither in vision nor language” ranks that pass the current logic until a later failure.

As per coding guidelines, "Follow the existing code style and conventions as documented in CODING_GUIDELINES.md".

src/megatron/bridge/training/config.py (1)

1318-1413: Add invariant checks for use_dist_train so dp-size computation can’t silently go wrong.

When use_dist_train=True, dp-size is derived from model.language_world_size. That’s fine if model.tensor_model_parallel_size, pipeline_model_parallel_size, and context_parallel_size are also the language-side values at this point. If those fields still reflect a “combined” or vision-side topology, dp-size will be incorrect.

Suggestion: add a small validation in set_data_parallel_size() (or model.finalize()) that language_world_size % (tp*pp*cp) == 0 with the exact values used here, and raise a clear error if not.

As per coding guidelines, "Be explicit about required vs optional fields in configuration objects; do not add arbitrary defaults".

src/megatron/bridge/models/gpt_provider.py (1)

200-210: Use T | None for the new nullable config fields.
Align new fields with the repo’s nullable type-hint convention.

🛠️ Proposed fix

-    vision_model_type: Optional[str] = None
+    vision_model_type: str | None = None
@@
-    dist_train_vision_chunk_size: Optional[int] = 1
-    vision_world_size: Optional[int] = None
-    language_world_size: Optional[int] = None
-    vision_tensor_model_parallel_size: Optional[int] = None
-    vision_pipeline_model_parallel_size: Optional[int] = None
-    vision_context_parallel_size: Optional[int] = None
-    vision_expert_tensor_parallel_size: Optional[int] = None
-    vision_expert_model_parallel_size: Optional[int] = None
+    dist_train_vision_chunk_size: int | None = 1
+    vision_world_size: int | None = None
+    language_world_size: int | None = None
+    vision_tensor_model_parallel_size: int | None = None
+    vision_pipeline_model_parallel_size: int | None = None
+    vision_context_parallel_size: int | None = None
+    vision_expert_tensor_parallel_size: int | None = None
+    vision_expert_model_parallel_size: int | None = None

As per coding guidelines, Use 'T | None' for nullable types instead of 'Optional[T]'.

src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/utils.py (4)

16-17: Use modern union type syntax per coding guidelines.

The coding guidelines require T | None instead of Optional[T] and X | Y instead of Union[X, Y] (Python 3.10+). These typing imports are used throughout the file (lines 92–94, 103, etc.).

Suggested fix

-from typing import Optional, Union
+from __future__ import annotations

Then replace all Optional[X] with X | None and Union[X, Y] with X | Y throughout the file.

---

569-596: AllGatherVisionEmbeddings.backward — integer start_idx when cp_rank == 0 may cause slicing issues.

On Line 593, when cp_rank == 0, torch.cat(seqlens_on_cp_ranks[:0]) produces an empty tensor, and .sum() returns a tensor (scalar 0), not a Python int. Then on Line 595, grad_output[start_idx:end_idx] uses tensor indices. While this works in practice, mixing tensor and int index types is fragile. More importantly, the ctx.save_for_backward stores all seqlens_on_cp_ranks tensors. If these are large lists, this is fine, but verify the tensors are detached and on the correct device.

Safer start_idx computation

-        start_idx = torch.cat(seqlens_on_cp_ranks[:cp_rank]).sum() if cp_rank != 0 else 0
+        start_idx = int(torch.cat(seqlens_on_cp_ranks[:cp_rank]).sum().item()) if cp_rank != 0 else 0

---

700-701: postprocess_packed_seqs: redundant .cpu() call may trigger an extra D2H sync.

attention_mask.sum(dim=1, dtype=torch.int32).cpu().tolist() on Line 701 explicitly moves to CPU. If attention_mask is already on GPU, .tolist() alone will perform the D2H transfer. The extra .cpu() call is harmless but redundant. More importantly, this is a synchronization point — consider batching it with other D2H transfers above if performance matters here.

---

290-292: Magic default token IDs should reference the config, not be hardcoded.

image_token_id: int = 151655 and video_token_id: int = 151656 are hardcoded defaults in reorganize_inputs. These IDs are model-specific and already defined in Qwen3VLTransformerConfig. Hardcoding them here risks silent mismatch if the config changes.

src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/vision_model.py (2)

134-195: fast_pos_embed_interpolate: potential performance concern with .tolist() in inner loop.

Lines 170–171 call .tolist() and .extend() inside the per-image loop, converting GPU tensors to Python lists element by element. For a large number of images/grids, this creates many small D2H synchronizations and Python list operations. Consider accumulating the index and weight tensors on GPU and converting once after the loop.

Also, the unused loop variable t on Line 140 (flagged by Ruff B007) can be replaced with _ to signal intent.

---

239-244: torch.split followed by torch.cat is a no-op.

Lines 242–244 split hidden_states into chunks and immediately concatenate them back. Unless there's a side effect or a planned insertion between split and cat, this can be removed.

-        split_sizes = (grid_thw.prod(-1) // self.spatial_merge_size**2).tolist()
-        hidden_states = torch.split(hidden_states, split_sizes)
-        hidden_states = torch.cat(hidden_states, dim=0)

If this is a placeholder for future logic, add a comment explaining it.

src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/transformer_block.py (2)

64-64: Remove or document commented-out code.

Lines 64 and 76 contain commented-out model_comm_pgs parameters without any explanation. As per coding guidelines, commented-out code should include a comment describing its purpose and why it's commented out, or be removed before merging.

-        # model_comm_pgs: ModelCommProcessGroups = None,
...
-            # model_comm_pgs=model_comm_pgs,

Also applies to: 76-76

---

336-339: O(n) index lookup self.deepstack_visual_indexes.index(l_no) on every layer forward pass.

Both in the checkpointed path (Line 134) and the non-checkpointed path (Line 337), .index(l_no) performs a linear scan of deepstack_visual_indexes for every matching layer. Consider converting to a dict mapping layer_no → deepstack_idx in __init__ for O(1) lookup.

Proposed optimization in __init__

         self.deepstack_visual_indexes = config.deepstack_visual_indexes
+        self._deepstack_index_map = {
+            l_no: idx for idx, l_no in enumerate(config.deepstack_visual_indexes)
+        }

Then use self._deepstack_index_map[l_no] instead of self.deepstack_visual_indexes.index(l_no).