fix(qwen-vl): pad tokenizer-derived vocabularies

Signed-off-by: Chen Cui <chcui@nvidia.com>

Author: cuichenx

src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/model.py

@@ -52,6 +52,7 @@
 )
 from megatron.bridge.models.qwen_vl.modelling_qwen3_vl.vision_model import Qwen3VLVisionModel
 from megatron.bridge.training.utils.packed_seq_utils import get_packed_seq_cp_partition_indices
+from megatron.bridge.utils.vocab_utils import calculate_padded_vocab_size
 
 
 def _is_mrope_position_ids(position_ids: torch.Tensor | None) -> bool:
@@ -204,10 +205,22 @@ def __init__(
                 pg_collection=pg_collection,
             )
         if self.add_decoder:
+            assert language_transformer_config.vocab_size is not None, (
+                "vocab_size must be configured before constructing the Qwen3-VL language model"
+            )
+            if language_transformer_config.should_pad_vocab:
+                language_model_vocab_size = calculate_padded_vocab_size(
+                    language_transformer_config.vocab_size,
+                    language_transformer_config.make_vocab_size_divisible_by,
+                    language_transformer_config.tensor_model_parallel_size,
+                )
+            else:
+                language_model_vocab_size = language_transformer_config.vocab_size
+
             self.language_model = Qwen3VLGPTModel(
                 config=language_transformer_config,
                 transformer_layer_spec=language_transformer_layer_spec,
-                vocab_size=language_transformer_config.vocab_size,
+                vocab_size=language_model_vocab_size,
                 max_sequence_length=language_transformer_config.language_max_sequence_length,
                 parallel_output=parallel_output,
                 position_embedding_type="mrope",

src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/transformer_config.py

@@ -28,6 +28,8 @@ class Qwen3VLTransformerConfig(TransformerConfig):
     """Configuration for Qwen3-VL transformer with vision and language components."""
 
     vocab_size: int = 64000
+    make_vocab_size_divisible_by: int = 128
+    should_pad_vocab: bool = False
     language_max_sequence_length: int = 4096
 
     patch_size: int = 16

src/megatron/bridge/perf_recipes/qwen_vl/common.py

@@ -39,6 +39,9 @@ def _use_model_vocab_null_tokenizer(cfg: ConfigContainer) -> None:
     cfg.tokenizer.tokenizer_type = "NullTokenizer"
     cfg.tokenizer.tokenizer_model = None
     cfg.tokenizer.vocab_size = cfg.model.vocab_size
+    # The synthetic tokenizer mirrors the fixed benchmark model shape; it does
+    # not define a new tokenizer-derived vocabulary for from-scratch training.
+    cfg.tokenizer.use_tokenizer_vocab_size = False
 
 
 def _qwen35_vl_common(cfg: ConfigContainer) -> None:

tests/unit_tests/models/qwen_vl/modelling_qwen3_vl/test_model.py

@@ -23,6 +23,7 @@
 import os
 from dataclasses import replace
 from types import SimpleNamespace
+from unittest.mock import Mock
 
 import pytest
 import torch
@@ -326,6 +327,53 @@ def test_shared_embedding_or_output_weight(self, hf_config):
         weight_no_decoder = model_no_decoder.shared_embedding_or_output_weight()
         assert weight_no_decoder is None
 
+    @pytest.mark.parametrize(
+        ("vocab_size", "should_pad_vocab", "expected_vocab_size"),
+        [
+            (151669, True, 152064),
+            (248077, True, 248320),
+            (151936, False, 151936),
+        ],
+    )
+    def test_language_model_honors_vocab_padding_policy(
+        self,
+        hf_config,
+        monkeypatch,
+        vocab_size,
+        should_pad_vocab,
+        expected_vocab_size,
+    ):
+        """Apply tokenizer-derived padding before constructing the Qwen language model."""
+        self._setup_parallel_state(tp_size=1, ep_size=1, pp_size=1)
+        pg_collection = ProcessGroupCollection.use_mpu_process_groups()
+        language_transformer_config = self.get_language_transformer_config(hf_config)
+        language_transformer_config.vocab_size = vocab_size
+        language_transformer_config.should_pad_vocab = should_pad_vocab
+        language_transformer_config.make_vocab_size_divisible_by = 128
+        language_transformer_config.tensor_model_parallel_size = 4
+
+        language_model = Mock()
+        language_model.config.cuda_graph_impl = "none"
+        language_model.share_embeddings_and_output_weights = False
+        language_model_constructor = Mock(return_value=language_model)
+        monkeypatch.setattr(
+            "megatron.bridge.models.qwen_vl.modelling_qwen3_vl.model.Qwen3VLGPTModel",
+            language_model_constructor,
+        )
+
+        Qwen3VLModel(
+            vision_transformer_config=self.get_vision_transformer_config(hf_config),
+            language_transformer_config=language_transformer_config,
+            language_transformer_layer_spec=self.get_language_model_layer_spec(),
+            pre_process=False,
+            post_process=True,
+            add_encoder=False,
+            add_decoder=True,
+            pg_collection=pg_collection,
+        )
+
+        assert language_model_constructor.call_args.kwargs["vocab_size"] == expected_vocab_size
+
     @pytest.mark.timeout(50)
     def test_set_input_tensor(self, hf_config):
         """Test set_input_tensor method."""

tests/unit_tests/recipes/test_all_perf_recipe_factories.py

@@ -88,3 +88,4 @@ def test_perf_recipe_factory_builds_config(recipe_factory: Callable[..., object]
     if "pretrain" in recipe_factory.__name__ and isinstance(cfg.dataset, MockVLMSFTDatasetConfig):
         assert cfg.tokenizer.tokenizer_type == "NullTokenizer"
         assert cfg.tokenizer.vocab_size == cfg.model.vocab_size
+        assert cfg.tokenizer.use_tokenizer_vocab_size is False