[Model] Update qwen3vl for current Megatron-LM-FL (flagos-ai#1235)

### PR Category
<!-- One of [ Train | Inference | Compress | Serve | RL | Core |
Hardware | CICD | Tools | Others ] -->
Train
### PR Types
<!-- One of [ User Experience | New Features | Bug Fixes | Improvements
| Performance | Breaking Change| Deprecations | Test Case | Docs |
Others ] -->
Others
### PR Description
<!-- Describe what you’ve done -->
Update qwen3vl for the current `Megatron-LM-FL` and add functional
tests.

Author: heavyrain-lzy

flagscale/models/megatron/qwen2_5_vl/language_module.py

@@ -36,9 +36,10 @@ def __init__(
         position_embedding_type: Literal['learned_absolute', 'rope', 'none'] = 'learned_absolute',
         num_tokentypes: int = 0,
         scatter_to_sequence_parallel: bool = False, # chage default to False
+        tp_group = None,
     ):
         assert scatter_to_sequence_parallel == False, "QwenVLLanguageModelEmbedding does not support scatter_to_sequence_parallel"
-        super().__init__(config, vocab_size, max_sequence_length, position_embedding_type, num_tokentypes, scatter_to_sequence_parallel)
+        super().__init__(config, vocab_size, max_sequence_length, position_embedding_type, num_tokentypes, scatter_to_sequence_parallel, tp_group)
 
 
     def forward(

flagscale/models/megatron/qwen2_5_vl/vit_model.py

@@ -68,7 +68,7 @@ def enable_linear(self):
         # and https://github.com/huggingface/transformers/pull/45041
         # By default, we use CUDNN's convolution ops with optimization.
         return self.kernel_size == self.stride and \
-                version.parse(torch.__version__) > version.parse('2.9.0')
+                version.parse(torch.__version__) > version.parse('2.9.0') and version.parse(torch.__version__) < version.parse('2.11.0')
 
     def _forward_matmul(self, hidden_states):
         target_dtype = self.proj.weight.dtype

flagscale/models/megatron/qwen3_vl/language_model.py

@@ -27,6 +27,7 @@
 from megatron.core.utils import WrappedTensor, deprecate_inference_params
 from megatron.core.models.gpt.gpt_model import GPTModel
 from megatron.core.process_groups_config import ProcessGroupCollection
+from megatron.core.transformer.multi_token_prediction import MultiTokenPredictionBlock
 
 from .language_transformer_block import LanguageTransformerBlock
 
@@ -100,7 +101,7 @@ def apply_interleaved_mrope(self, freqs, mrope_section):
             freqs_t[..., idx] = freqs[dim, ..., idx]
         return freqs_t
 
-    def forward(self, position_ids: torch.Tensor, mrope_section: List[int]) -> Tensor:
+    def forward(self, position_ids: torch.Tensor, mrope_section: List[int], cp_group: Optional[torch.distributed.ProcessGroup] = None) -> Tensor:
         """Forward pass of multimodal RoPE embedding.
 
         Args:
@@ -137,10 +138,12 @@ def forward(self, position_ids: torch.Tensor, mrope_section: List[int]) -> Tenso
 
         # shape (seq_length, bs, 1, 2 * dim)
         emb = emb[..., None, :].transpose(0, 1).contiguous()
-        if self.cp_group is not None and self.cp_group.size() > 1:
+        if cp_group is None:
+            cp_group = self.cp_group
+        if cp_group is not None and cp_group.size() > 1:
             # slice rotary_pos_emb along sequence dimension and select the parition of the current
             # CP rank
-            emb = get_pos_emb_on_this_cp_rank(emb, 0, self.cp_group)
+            emb = get_pos_emb_on_this_cp_rank(emb, 0, cp_group)
         return emb
 
 class Qwen3VLLanguageModule(GPTModel):
@@ -228,6 +231,7 @@ def __init__(
                 vocab_size=self.vocab_size,
                 max_sequence_length=self.max_sequence_length,
                 position_embedding_type=position_embedding_type,
+                tp_group=self.pg_collection.tp,
             )
         if self.position_embedding_type == 'mrope' and not self.config.multi_latent_attention:
             self.rotary_pos_emb = Qwen3VLLanguageRotaryEmbedding(
@@ -258,7 +262,7 @@ def __init__(
 
         if self.mtp_process:
             self.mtp = MultiTokenPredictionBlock(
-                config=self.config, spec=self.mtp_block_spec, vp_stage=vp_stage
+                config=self.config, spec=self.mtp_block_spec, vp_stage=vp_stage, pg_collection=self.pg_collection
             )
 
         # Output
@@ -314,18 +318,20 @@ def forward(self, input_ids, position_ids, attention_mask,
                 visual_pos_masks: Optional[torch.Tensor] = None,
                 deepstack_visual_embeds: Optional[list[torch.Tensor]] = None,
                 *, inference_params = None,
-                loss_mask = None):
+                loss_mask = None,
+                padding_mask: Optional[torch.Tensor] = None):
 
         inference_context = deprecate_inference_params(inference_context, inference_params)
 
-        decoder_input, rotary_pos_emb, rotary_pos_cos, rotary_pos_sin, sequence_len_offset = (
+        decoder_input, rotary_pos_emb, rotary_pos_cos, rotary_pos_sin, sequence_len_offset, padding_mask = (
             self._preprocess(
                 input_ids=input_ids,
                 position_ids=position_ids,
                 decoder_input=decoder_input,
                 inference_context=inference_context,
                 packed_seq_params=packed_seq_params,
-            )
+                padding_mask=padding_mask,
+            )[:6]
         )
 
         # Run decoder.
@@ -338,6 +344,7 @@ def forward(self, input_ids, position_ids, attention_mask,
             rotary_pos_sin=rotary_pos_sin,
             packed_seq_params=packed_seq_params,
             sequence_len_offset=sequence_len_offset,
+            padding_mask=padding_mask,
             visual_pos_masks = visual_pos_masks,
             deepstack_visual_embeds = deepstack_visual_embeds,
             **(extra_block_kwargs or {}),

flagscale/models/megatron/qwen3_vl/language_transformer_block.py

@@ -62,6 +62,21 @@
 
 class LanguageTransformerBlock(TransformerBlock):
 
+    def __init__(
+        self,
+        config,
+        spec,
+        post_layer_norm=True,
+        pre_process=True,
+        post_process=True,
+        pg_collection=None,
+        vp_stage=None,
+        dualpipev_stage=None,
+    ):
+        super().__init__(
+            config, spec, post_layer_norm, pre_process, post_process, pg_collection, vp_stage, dualpipev_stage
+        )
+
     def _checkpointed_forward(
         self,
         hidden_states: Tensor,
@@ -207,6 +222,7 @@ def forward(
         inference_context: Optional[BaseInferenceContext] = None,
         packed_seq_params: Optional[PackedSeqParams] = None,
         sequence_len_offset: Optional[Tensor] = None,
+        padding_mask: Optional[Tensor] = None,
         # args for deepstack
         visual_pos_masks: Optional[torch.Tensor] = None,
         deepstack_visual_embeds: Optional[list[torch.Tensor]] = None,
@@ -343,6 +359,7 @@ def forward(
                             inference_context=inference_context,
                             packed_seq_params=packed_seq_params,
                             sequence_len_offset=sequence_len_offset,
+                            padding_mask=padding_mask,
                         )
                     # Deepstack visual embedding addition
                     # NOTE: Assume that this is first pipeline stage that has at least three layers.

flagscale/models/megatron/qwen3_vl/model.py

@@ -72,6 +72,7 @@ def __init__(
         language_rotary_base: int = 10000,
         fp16_lm_cross_entropy: bool = False,
         language_share_embeddings_and_output_weights: bool=False,
+        pg_collection = None,
         vp_stage: int=None
     ) -> None:
         super().__init__(config=language_transformer_config)
@@ -121,6 +122,7 @@ def __init__(
             share_embeddings_and_output_weights=language_share_embeddings_and_output_weights,
             rope_scaling=False,
             mtp_block_spec=None,
+            pg_collection=pg_collection,
             vp_stage=vp_stage,
         )
         self.share_embeddings_and_output_weights = (

flagscale/models/megatron/qwen3_vl/vision_transformer_block.py

@@ -67,12 +67,12 @@ class VisionTransformerBlock(TransformerBlock):
     """
     def __init__(self, config, spec,
                  post_layer_norm = True, pre_process = True, post_process = True,
-                 pg_collection = None, vp_stage = None,
+                 pg_collection = None, vp_stage = None, dualpipev_stage = None,
                  projection_config = None,  # Note: DeepStack
                  projection_layer_spec = None,  # Note: DeepStack
                  projection_type = 'mlp',  # Note: DeepStack):
                 ):
-        super().__init__(config, spec, post_layer_norm, pre_process, post_process, pg_collection, vp_stage)
+        super().__init__(config, spec, post_layer_norm, pre_process, post_process, pg_collection, vp_stage, dualpipev_stage)
 
         if self.final_layernorm != None:
             # NOTE(lizhiyu): replace final layernorm with TENorm if using TE

flagscale/train/megatron/train_qwen3_vl.py

@@ -92,13 +92,17 @@
 #### especially for qwen2.5-vl ####
 IGNORE_IDX=-100
 def model_provider(
-    pre_process=True, post_process=True, add_encoder=True, add_decoder=True
+    pre_process=True, post_process=True, vp_stage=None, config=None, pg_collection=None
 ) -> Union[Qwen3VLModel]:
     args = get_args()
     print_rank_0("start building qwen3-vl model ...")
 
     # Config of vit, llm and projector
-    config = core_transformer_config_from_args(args, Qwen3VLTransformerConfig)
+    if config is None:
+        config = core_transformer_config_from_args(args, Qwen3VLTransformerConfig)
+    else:
+        # config passed from backend, use it directly
+        pass
     use_te = args.transformer_impl == "transformer_engine"
     if not use_te:
         raise NotImplementedError("The Qwen3-VL model is only implemented with TransformerEngine!")
@@ -144,12 +148,12 @@ def model_provider(
 
         pre_process=pre_process,
         post_process=post_process,
-        add_decoder=add_decoder,
-        add_encoder=add_encoder,
 
         fp16_lm_cross_entropy=args.fp16_lm_cross_entropy,
         parallel_output=True,
         language_share_embeddings_and_output_weights=not args.untie_embeddings_and_output_weights,
+        pg_collection=pg_collection,
+        vp_stage=vp_stage,
     )
 
     model.freeze(

tests/functional_tests/train/qwen3_vl/conf/tp2.yaml

@@ -0,0 +1,33 @@
+defaults:
+  - _self_
+  - train: tp2
+
+experiment:
+  exp_name: tp2
+  exp_dir: tests/functional_tests/train/qwen3_vl/test_results/tp2
+  task:
+    type: train
+    backend: megatron
+    entrypoint: flagscale/train/megatron/train_qwen3_vl.py
+  runner:
+    ssh_port: null
+  shell_cmds: null
+  envs:
+    HYDRA_FULL_ERROR: 1
+    CUDA_VISIBLE_DEVICES: "0,1,2,3,4,5,6,7"
+    CUDA_DEVICE_MAX_CONNECTIONS: 1
+    # CUBLAS_WORKSPACE_CONFIG: ":4096:8"
+    NCCL_ALGO: "Ring"
+    NVTE_APPLY_QK_LAYER_SCALING: 0
+    NVTE_ALLOW_NONDETERMINISTIC_ALGO: 0
+    NVTE_FLASH_ATTN: 1
+    NVTE_FUSED_ATTN: 0
+    CUDNN_BENCHMARK: "false"
+    CUDNN_DETERMINISTIC: "true"
+  cmds:
+    before_start: source /root/miniconda3/bin/activate flagscale-train && pip install git+https://github.com/NVIDIA/Megatron-Energon.git@ab40226100830f41de38d1f1204d7848b54b1f3e && pip install "multi-storage-client==0.42.0" && pip install transformers==4.57
+action: run
+
+hydra:
+  run:
+    dir: ${experiment.exp_dir}/hydra

tests/functional_tests/train/qwen3_vl/conf/train/tp2.yaml

@@ -0,0 +1,105 @@
+system:
+  vision_ration: 0.1
+  num_workers: 1
+  calculate_per_token_loss: true
+  tensor_model_parallel_size: 4
+  pipeline_model_parallel_size: 1
+  context_parallel_size: 1
+  use_flash_attn: True
+  use_distributed_optimizer: True
+  sequence_parallel: True
+  tp_comm_overlap: False
+  overlap_grad_reduce: False # if has text-only must be false
+  overlap_param_gather: False # if has text-only must be false
+  use_mcore_models: True
+  transformer_impl: transformer_engine
+  use_te: True
+  precision:
+    bf16: True
+    attention_softmax_in_fp32: True
+  logging:
+    log_interval: 1
+    tensorboard_log_interval: 1
+    log_throughput: True
+    wandb_project: ${experiment.exp_name}
+    wandb_exp_name: ${experiment.exp_name}
+    log_params_norm: True
+    log_num_zeros_in_grad: True
+  checkpoint:
+    save_interval: 1000
+    # pretrained_checkpoint: /nfs/lizhiyu/embody/data/Qwen3-VL-8B-Instruct-tp2/
+    dataloader_save: ${experiment.exp_dir}/checkpoints/dataloader
+    ckpt_format: torch
+    async_save: False
+
+model:
+  kv_channels: 128 # the weight out_size of prepare qkv
+  qk_layernorm: True
+  attention_backend: flash # don't use "auto(nvte_flash_attn)"
+  disable_bias_linear: True
+  num_layers: 36
+  hidden_size: 4096
+  ffn_hidden_size: 12288
+  num_attention_heads: 32
+  num_query_groups: 8
+  seq_length: 2048
+  max_padding_length: 2048 # (cutoff_len)max 262144, change according the dataset
+  # especial for qwen3-vl
+  enable_variable_seq_lengths: True
+  max_position_embeddings: 262144 # only useful for additional position embedding
+  swiglu: True
+  normalization: RMSNorm
+  norm_epsilon: 1e-6
+  init_method_std: 0.02
+  attention_dropout: 0.0
+  hidden_dropout: 0.0
+  clip_grad: 1.0
+  train_iters: 10
+  eval_iters: 0 # no valid
+  eval_interval: 1000
+  micro_batch_size: 1
+  global_batch_size: 4
+  allow_missing_vision_projection_checkpoint: False
+  apply_layernorm_1p: False
+  group_query_attention: True
+  no_masked_softmax_fusion: True
+  untie_embeddings_and_output_weights: True
+
+  # position embedding
+  position_embedding_type: mrope
+  rotary_percent: 1.0
+  rotary_base: 5000000
+  rotary_seq_len_interpolation_factor: 1
+  no_rope_fusion: False
+  mrope_section: [24, 20, 20]
+  eod_mask_loss: False
+
+  # vision model
+  patch_size: 16
+  freeze_LM: False
+  freeze_ViT: False
+  disable_vision_class_token: True
+  seed: 42
+
+  optimizer:
+    weight_decay: 0.1
+    adam_beta1: 0.9
+    adam_beta2: 0.999
+    lr_scheduler:
+      lr: 1.0e-5
+      min_lr: 1.0e-6
+      lr_warmup_fraction: .03
+      lr_decay_style: cosine
+
+data:
+  no_use_system_prompt: True
+  data_path: /home/gitlab-runner/data/blip_laion_cc_sbu_558k_first_5k/wds-2/
+  vision_root: /home/gitlab-runner/data/blip_laion_cc_sbu_558k_first_5k/
+  dataloader_type: external
+  split: 100,0,0
+  tokenizer:
+    tokenizer_type: Qwen2VLTokenizer
+    tokenizer_path: /home/gitlab-runner/tokenizers/Qwen3-VL-8B-Instruct
+    # vocab_size: 151936 #
+    extra_vocab_size: 293 # Qwen3-VL specific. total vocab size = 151643 + extra_vocab_size
+    make_vocab_size_divisible_by: 64

tests/functional_tests/train/qwen3_vl/gold_values/tp2.json

@@ -0,0 +1 @@
+{"lm loss:": {"values": [12.81162, 12.78231, 10.90357, 13.56564, 13.17456, 13.97026, 11.54165, 10.94413, 11.21150, 9.978001], "rtol": 0.1, "atol": 0.2}}