update qwen3vl

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(

tools/checkpoint/qwen2_5_vl/utils.py

@@ -174,7 +174,9 @@ def safe_copy(src_tensor: torch.Tensor, dst_tensor: torch.Tensor, skip_dtype_ass
             raise ValueError(
                 f"Get source dtype {src_tensor.dtype}, but target dtype {dst_tensor.dtype}"
             )
-    assert src_tensor.shape == dst_tensor.shape
+    assert src_tensor.shape == dst_tensor.shape, (
+        f"Get source shape {src_tensor.shape}, but target shape {dst_tensor.shape}"
+    )
     dst_tensor.data.copy_(src_tensor.data)
     return src_tensor.numel()
 

tools/checkpoint/qwen3_vl/hf2mcore_qwen3_vl.py

@@ -208,7 +208,9 @@ def load_megatron_model(args):
 
 
 @torch.inference_mode()
-def convert_checkpoint_from_megatron_to_transformers(mgmodel, hfmodel, args):
+def convert_checkpoint_from_megatron_to_transformers(
+    mgmodel, hfmodel, args, hf_pretrained_model=None
+):
     if args.fp16:
         mgmodel = mgmodel.half()
         hfmodel = hfmodel.half()
@@ -230,7 +232,14 @@ def convert_checkpoint_from_megatron_to_transformers(mgmodel, hfmodel, args):
     vision_head_dim = vision_hidden_size // mgvision.config.num_attention_heads
     copied_numel = 0
     safe_copy(mgvision.rotary_pos_emb.inv_freq, hfvision.rotary_pos_emb.inv_freq)
-    copied_numel += safe_copy(mgvision.patch_embed.proj.weight, hfvision.patch_embed.proj.weight)
+    # patch_embed: mcore may use Linear (enable_linear) while hf uses Conv3d
+    mg_patch_weight = mgvision.patch_embed.proj.weight
+    hf_patch_weight = hfvision.patch_embed.proj.weight
+    if mg_patch_weight.dim() == 2 and hf_patch_weight.dim() == 5:
+        # mcore Linear [embed_dim, flat_dim] -> hf Conv3d [embed_dim, C, T, H, W]
+        copied_numel += safe_copy(mg_patch_weight.view(hf_patch_weight.shape), hf_patch_weight)
+    else:
+        copied_numel += safe_copy(mg_patch_weight, hf_patch_weight)
     copied_numel += safe_copy(mgvision.patch_embed.proj.bias, hfvision.patch_embed.proj.bias)
     copied_numel += safe_copy(mgvision.pos_embed.weight, hfvision.pos_embed.weight)
 
@@ -367,7 +376,7 @@ def convert_checkpoint_from_megatron_to_transformers(mgmodel, hfmodel, args):
 
     copied_numel += safe_copy(mgllm.decoder.final_layernorm.weight, hfllm.norm.weight)
     if args.untie_embeddings_and_output_weights:
-        safe_copy(mgllm.output_layer.weight, hfmodel.lm_head.weight)
+        safe_copy(mgllm.output_layer.weight, hf_pretrained_model.lm_head.weight)
 
     n_params = sum(
         [
@@ -382,7 +391,9 @@ def convert_checkpoint_from_megatron_to_transformers(mgmodel, hfmodel, args):
 
 
 @torch.inference_mode()
-def convert_checkpoint_from_transformers_to_megatron(hfmodel, mgmodel, args):
+def convert_checkpoint_from_transformers_to_megatron(
+    hfmodel, mgmodel, args, hf_pretrained_model=None
+):
     if args.fp16:
         mgmodel = mgmodel.half()
         hfmodel = hfmodel.half()
@@ -410,7 +421,14 @@ def convert_checkpoint_from_transformers_to_megatron(hfmodel, mgmodel, args):
     copied_numel = 0
     # rotary_pos_emb.inv_freq is buffer not parameter
     safe_copy(hfvision.rotary_pos_emb.inv_freq, mgvision.rotary_pos_emb.inv_freq)
-    copied_numel += safe_copy(hfvision.patch_embed.proj.weight, mgvision.patch_embed.proj.weight)
+    # patch_embed: hf uses Conv3d while mcore may use Linear (enable_linear)
+    hf_patch_weight = hfvision.patch_embed.proj.weight
+    mg_patch_weight = mgvision.patch_embed.proj.weight
+    if hf_patch_weight.dim() == 5 and mg_patch_weight.dim() == 2:
+        # hf Conv3d [embed_dim, C, T, H, W] -> mcore Linear [embed_dim, flat_dim]
+        copied_numel += safe_copy(hf_patch_weight.reshape(mg_patch_weight.shape), mg_patch_weight)
+    else:
+        copied_numel += safe_copy(hf_patch_weight, mg_patch_weight)
     copied_numel += safe_copy(hfvision.patch_embed.proj.bias, mgvision.patch_embed.proj.bias)
     copied_numel += safe_copy(hfvision.pos_embed.weight, mgvision.pos_embed.weight)
 
@@ -550,7 +568,7 @@ def convert_checkpoint_from_transformers_to_megatron(hfmodel, mgmodel, args):
     copied_numel += safe_copy(hfllm.norm.weight, mgllm.decoder.final_layernorm.weight)
     if args.untie_embeddings_and_output_weights:
         # lm_head not in hfllm
-        safe_copy(hfmodel.lm_head.weight, mgllm.output_layer.weight)
+        safe_copy(hf_pretrained_model.lm_head.weight, mgllm.output_layer.weight)
 
     n_params = sum([t.numel() for t in hfllm.state_dict().values()])
     assert n_params == copied_numel, (
@@ -873,15 +891,19 @@ def main():
             args.hf_ckpt_path, torch_dtype=config.torch_dtype
         )
         mg_model = load_megatron_model(args)
-        convert_checkpoint_from_megatron_to_transformers(mg_model, hf_model, args)
+        convert_checkpoint_from_megatron_to_transformers(
+            mg_model, hf_model.model, args, hf_pretrained_model=hf_model
+        )
         save_hfmodel(args, hf_model)
     else:
         config = AutoConfig.from_pretrained(args.load)
         hf_model = Qwen3VLForConditionalGeneration.from_pretrained(
             args.load, torch_dtype=config.torch_dtype
         )
         mg_model = model_provider()
-        convert_checkpoint_from_transformers_to_megatron(hf_model, mg_model, args)
+        convert_checkpoint_from_transformers_to_megatron(
+            hf_model.model, mg_model, args, hf_pretrained_model=hf_model
+        )
         save_mgmodel(mg_model, args)
 
 

tools/datasets/qwenvl/data/dataset_helpers.py

@@ -381,14 +381,19 @@ def encode_chatml(self, sample: ChatMLSample):
 
         # NOTE: we need to mask all system/user input tokens and assistant generation prefix tokens
         input_ids = self.tokenizer.apply_chat_template(
-            conversation, tokenize=True, return_tensors="np"
+            conversation,
+            tokenize=True,
+            return_tensors="np",
+            return_dict=False,
         )[0]
         target = input_ids.copy()
 
         conversation_start_idx = 0
         if self.use_system_prompt:
             system_prompt_prefix = len(
-                self.tokenizer.apply_chat_template([conversation[0]], tokenize=True)
+                self.tokenizer.apply_chat_template(
+                    [conversation[0]], tokenize=True, return_dict=False
+                )
             )
             conversation_start_idx = 1
         else:
@@ -403,7 +408,10 @@ def encode_chatml(self, sample: ChatMLSample):
         for turn_idx, turn in enumerate(conversation[conversation_start_idx:]):
             if self._single_turn_template:
                 turn_tokens = self.tokenizer.apply_chat_template(
-                    [turn], tokenize=True, return_tensors="np"
+                    [turn],
+                    tokenize=True,
+                    return_tensors="np",
+                    return_dict=False,
                 )[0]
                 turn_content = turn_tokens[system_prompt_prefix:]
                 n_tokens = len(turn_content)
@@ -414,7 +422,10 @@ def encode_chatml(self, sample: ChatMLSample):
             else:
                 prefix = conversation[: conversation_start_idx + turn_idx + 1]
                 prefix_tokens = self.tokenizer.apply_chat_template(
-                    prefix, tokenize=True, return_tensors="np"
+                    prefix,
+                    tokenize=True,
+                    return_tensors="np",
+                    return_dict=False,
                 )[0]
                 n_tokens = len(prefix_tokens) - prev_prefix_len
 
@@ -545,8 +556,16 @@ def encode_vqa(self, sample: VQASample):
             {"role": "assistant", "content": answer},
         ]
 
-        user_inputs = self.tokenizer.apply_chat_template(conversation[:-1], tokenize=False)
-        text = self.tokenizer.apply_chat_template(conversation, tokenize=False)
+        user_inputs = self.tokenizer.apply_chat_template(
+            conversation[:-1],
+            tokenize=False,
+            return_dict=False,
+        )
+        text = self.tokenizer.apply_chat_template(
+            conversation,
+            tokenize=False,
+            return_dict=False,
+        )
 
         # text, target = self.tokenizer.tokenize_conversation(conversation, False, False)
         # replace <image> token by <image> * (thw)