[megatron] support multimodal model CPT/SFT/DPO (Full/LoRA)

examples/megatron/multimodal/dense.sh

@@ -0,0 +1,31 @@
+# 4 * 56GiB; 2.3s/it
+PYTORCH_CUDA_ALLOC_CONF='expandable_segments:True' \
+NPROC_PER_NODE=4 \
+MAX_PIXELS=1003520 \
+CUDA_VISIBLE_DEVICES=0,1,2,3 \
+megatron sft \
+    --load Qwen2.5-VL-7B-Instruct-mcore \
+    --dataset 'AI-ModelScope/LaTeX_OCR:human_handwrite' \
+    --tensor_model_parallel_size 2 \
+    --sequence_parallel true \
+    --packing true \
+    --split_dataset_ratio 0.01 \
+    --micro_batch_size 1 \
+    --global_batch_size 4 \
+    --recompute_granularity full \
+    --recompute_method uniform \
+    --recompute_num_layers 1 \
+    --finetune true \
+    --cross_entropy_loss_fusion true \
+    --lr 1e-5 \
+    --lr_warmup_fraction 0.05 \
+    --min_lr 1e-6 \
+    --max_epochs 1 \
+    --save megatron_output/Qwen2.5-VL-7B-Instruct \
+    --save_interval 200 \
+    --vit_gradient_checkpointing true \
+    --max_length 2048 \
+    --num_workers 4 \
+    --no_save_optim true \
+    --no_save_rng true \
+    --dataset_num_proc 8

examples/train/multimodal/lora_llm_full_vit/sft.sh

@@ -1,5 +1,6 @@
 # 4 * 22GiB
 # vit/merger lr 1e-5; llm lora lr 1e-4
+# Note: not support resume_from_checkpoint (only support resume_only_model)
 NPROC_PER_NODE=4 \
 CUDA_VISIBLE_DEVICES=0,1,2,3 \
 MAX_PIXELS=1003520 \

swift/megatron/argument/megatron_args.py

@@ -94,6 +94,10 @@ class ExtraMegatronArguments(RLHFMegatronArgumentsMixin, MegatronTunerMixin):
     partial_rotary_factor: Optional[float] = None
     use_shared_expert_gate: Optional[bool] = None
 
+    # visual
+    vit_gradient_checkpointing: bool = True
+    gradient_checkpointing_kwargs: Optional[Union[dict, str]] = None
+
 
 @dataclass
 class MegatronArguments(ExtraMegatronArguments):
@@ -185,7 +189,8 @@ class MegatronArguments(ExtraMegatronArguments):
     group_query_attention: Optional[bool] = None
     num_query_groups: Optional[int] = None
     max_position_embeddings: Optional[int] = None
-    position_embedding_type: Literal['learned_absolute', 'rope', 'mrope', 'relative', 'none'] = 'rope'
+    position_embedding_type: Optional[Literal['learned_absolute', 'rope', 'mrope', 'relative', 'none']] = None
+    mrope_section: Optional[List[int]] = None
     rotary_base: Optional[int] = None
     rotary_percent: float = 1.
     rotary_interleaved: Optional[bool] = None
@@ -376,10 +381,14 @@ def __post_init__(self):
             self.rope_scaling = json_parse_to_dict(self.rope_scaling)
             if 'type' in self.rope_scaling and 'rope_type' not in self.rope_scaling:
                 self.rope_scaling['rope_type'] = self.rope_scaling['type']
+        if self.gradient_checkpointing_kwargs is not None:
+            self.gradient_checkpointing_kwargs = json_parse_to_dict(self.gradient_checkpointing_kwargs)
         if self.eval_interval is None:
             self.eval_interval = self.save_interval
         if self.seq_length is None:
             self.seq_length = self.max_position_embeddings
+        if self.position_embedding_type is None:
+            self.position_embedding_type = 'rope'
         if self.tensorboard_dir is None and self.save is not None:
             self.tensorboard_dir = f'{self.save}/runs'
         self._init_moe()

swift/megatron/argument/train_args.py

@@ -28,6 +28,9 @@ def init_model_args(self, tokenizer, config):
                 setattr(self, k, v)
         MegatronArguments.__post_init__(self)
         self.extra_args = self.parse_to_megatron()
+        self.extra_args['model_info'] = self.model_info
+        self.extra_args['model_meta'] = self.model_meta
+        self.extra_args['megatron_model_meta'] = self.megatron_model_meta
 
     def _init_save(self):
         init_process_group(backend=self.ddp_backend, timeout=self.ddp_timeout)

swift/megatron/model/__init__.py

@@ -1,4 +1,4 @@
 # Copyright (c) Alibaba, Inc. and its affiliates.
-from . import gpt
+from . import gpt, qwen2_5_vl
 from .constant import MegatronModelType
 from .register import MegatronModelMeta, get_megatron_model_meta, register_megatron_model

swift/megatron/model/constant.py

@@ -1,3 +1,4 @@
 # Copyright (c) Alibaba, Inc. and its affiliates.
 class MegatronModelType:
     gpt = 'gpt'
+    qwen2_5_vl = 'qwen2_5_vl'

swift/megatron/model/gpt/__init__.py

@@ -1,4 +1,10 @@
 # Copyright (c) Alibaba, Inc. and its affiliates.
+from dataclasses import dataclass
+from typing import Any, Callable, Dict
+
+from torch import nn
+from transformers import PretrainedConfig
+
 from swift.llm import ModelType
 from ..constant import MegatronModelType
 from ..register import MegatronModelMeta, register_megatron_model
@@ -53,3 +59,11 @@
         ModelType.glm4_5,
         ModelType.deepseek_v3_1,
     ], model_provider, convert_gpt_hf_config, convert_mcore2hf, convert_hf2mcore))
+
+
+@dataclass
+class GptMegatronModelMeta(MegatronModelMeta):
+    model_provider: Callable[[], nn.Module] = model_provider
+    convert_hf_config: Callable[[PretrainedConfig], Dict[str, Any]] = convert_gpt_hf_config
+    convert_mcore2hf: Callable[[nn.Module, nn.Module], None] = convert_mcore2hf
+    convert_hf2mcore: Callable[[nn.Module, nn.Module], None] = convert_hf2mcore

swift/megatron/model/gpt/config.py

@@ -39,6 +39,9 @@ def convert_gpt_hf_config(config) -> Dict[str, Any]:
             res['rotary_interleaved'] = True
         elif architectures == 'Glm4MoeForCausalLM':
             res['moe_router_score_function'] = 'sigmoid'
+        elif architectures == 'Qwen2_5_VLForConditionalGeneration':
+            res['position_embedding_type'] = 'mrope'
+            res['mrope_section'] = res['rope_scaling']['mrope_section']
     if first_k_dense_replace is not None:
         res['moe_layer_freq'] = f'[0]*{first_k_dense_replace}+[1]*{res["num_layers"] - first_k_dense_replace}'
     if res.get('moe_router_score_function', 'softmax') == 'sigmoid':

swift/megatron/model/gpt/hf2mcore.py

@@ -91,7 +91,7 @@ def set_mlp_state(args, mg_mlp, hf_mlp):
 
 def set_layer_state(args, mg_model, hf_model, layer_idx):
     mg_layer = mg_model.decoder.layers[layer_idx]
-    hf_layer = hf_model.model.layers[layer_idx]
+    hf_layer = hf_model.layers[layer_idx]
     if args.multi_latent_attention:
         set_mla_attn_state(args, mg_layer.self_attention, hf_layer.self_attn)
         mg_layer.input_layernorm.weight.data.copy_(hf_layer.input_layernorm.weight)
@@ -115,4 +115,4 @@ def convert_hf2mcore(hf_model, mg_model):
         mg_model.output_layer.weight.data.copy_(hf_model.lm_head.weight)
     mg_model.decoder.final_layernorm.weight.data.copy_(hf_model.model.norm.weight)
     for layer_idx in range(args.num_layers):
-        set_layer_state(args, mg_model, hf_model, layer_idx)
+        set_layer_state(args, mg_model, hf_model.model, layer_idx)

swift/megatron/model/gpt/mcore2hf.py

@@ -88,7 +88,7 @@ def set_mlp_state(args, mg_mlp, hf_mlp):
 
 def set_layer_state(args, mg_model, hf_model, layer_idx):
     mg_layer = mg_model.decoder.layers[layer_idx]
-    hf_layer = hf_model.model.layers[layer_idx]
+    hf_layer = hf_model.layers[layer_idx]
 
     if args.multi_latent_attention:
         set_mla_attn_state(args, mg_layer.self_attention, hf_layer.self_attn)
@@ -113,4 +113,4 @@ def convert_mcore2hf(hf_model, mg_model):
         hf_model.lm_head.weight.data.copy_(mg_model.output_layer.weight)
     hf_model.model.norm.weight.data.copy_(mg_model.decoder.final_layernorm.weight)
     for layer_idx in range(args.num_layers):
-        set_layer_state(args, mg_model, hf_model, layer_idx)
+        set_layer_state(args, mg_model, hf_model.model, layer_idx)

swift/megatron/model/gpt_model.py

@@ -4,13 +4,14 @@
 from typing import Any, Dict, Literal, Optional
 
 import torch
-from megatron.core import InferenceParams
+from megatron.core import InferenceParams, mpu
 from megatron.core.config_logger import has_config_logger_enabled, log_config_to_disk
 from megatron.core.models.common.embeddings.rotary_pos_embedding import RotaryEmbedding
 from megatron.core.models.gpt import GPTModel as McoreGPTModel
 from megatron.core.packed_seq_params import PackedSeqParams
 from megatron.core.transformer.spec_utils import ModuleSpec
 from megatron.core.transformer.transformer_config import TransformerConfig
+from megatron.training import get_args
 
 from swift.utils import get_logger
 from .rope import dynamic_rope_update, get_rope_inv_freq
@@ -91,6 +92,11 @@ def __init__(
             logger.warning('`apply_rope_fusion` does not support `attention_scaling`. '
                            f'Setting `config.apply_rope_fusion`: {config.apply_rope_fusion}')
 
+        args = get_args()
+        self.visual = None
+        if args.megatron_model_meta.visual is not None:
+            self.visual = args.megatron_model_meta.visual(config)
+
     @contextmanager
     def _patch_apply_rotary_pos_emb(self):
         if self.attention_scaling == 1.:
@@ -138,10 +144,19 @@ def forward(
         # Otherwise, apply embedding layer on input_ids and position_ids to get decoder_input.
 
         # Decoder embedding.
+        args = get_args()
         if decoder_input is not None:
             pass
         elif self.pre_process:
             decoder_input = self.embedding(input_ids=input_ids, position_ids=position_ids)
+            if self.visual is not None:
+                if args.tensor_model_parallel_size > 1 and args.sequence_parallel:
+                    input_ids = input_ids.chunk(
+                        args.tensor_model_parallel_size, dim=-1)[mpu.get_tensor_model_parallel_rank()]
+                kwargs.update({'input_ids': input_ids})
+                decoder_input = decoder_input.transpose(0, 1)
+                decoder_input = self.visual.get_inputs_embeds(decoder_input, **kwargs)
+                decoder_input = decoder_input.transpose(0, 1)
         else:
             # intermediate stage of pipeline
             # decoder will get hidden_states from encoder.input_tensor
@@ -172,6 +187,13 @@ def forward(
                     rotary_seq_len,
                     packed_seq=packed_seq_params is not None and packed_seq_params.qkv_format == 'thd',
                 )
+        elif self.position_embedding_type in 'mrope':
+            if self.training or not self.config.flash_decode:
+                rotary_pos_emb = self.rotary_pos_emb(position_ids, self.mrope_section)
+            else:
+                # Flash decoding uses precomputed cos and sin for RoPE
+                raise NotImplementedError('Flash decoding uses precomputed cos and sin for RoPE, not implmented in '
+                                          'MultimodalRotaryEmbedding yet.')
         if ((self.config.enable_cuda_graph or self.config.flash_decode) and rotary_pos_cos is not None
                 and inference_params):
             sequence_len_offset = torch.tensor(

swift/megatron/model/qwen2_5_vl/__init__.py

@@ -0,0 +1,15 @@
+from swift.llm import ModelType
+from ..constant import MegatronModelType
+from ..gpt import GptMegatronModelMeta
+from ..register import MegatronModelMeta, register_megatron_model
+from .convert import convert_hf2mcore_qwen2_5_vl, convert_mcore2hf_qwen2_5_vl
+from .vit import Qwen2_5VL_Vit
+
+register_megatron_model(
+    GptMegatronModelMeta(
+        MegatronModelType.qwen2_5_vl, [
+            ModelType.qwen2_5_vl,
+        ],
+        convert_hf2mcore=convert_hf2mcore_qwen2_5_vl,
+        convert_mcore2hf=convert_mcore2hf_qwen2_5_vl,
+        visual=Qwen2_5VL_Vit))

swift/megatron/model/qwen2_5_vl/convert.py

@@ -0,0 +1,31 @@
+from megatron.training import get_args
+
+from swift.utils import deep_getattr
+from ..gpt.hf2mcore import set_layer_state as set_layer_state_hf2mcore
+from ..gpt.mcore2hf import set_layer_state as set_layer_state_mcore2hf
+
+
+def convert_hf2mcore_qwen2_5_vl(hf_model, mg_model):
+    language_model = hf_model.model.language_model
+    args = get_args()
+    # language_model
+    mg_model.embedding.word_embeddings.weight.data.copy_(language_model.embed_tokens.weight)
+    if args.untie_embeddings_and_output_weights:
+        mg_model.output_layer.weight.data.copy_(hf_model.lm_head.weight)
+    mg_model.decoder.final_layernorm.weight.data.copy_(language_model.norm.weight)
+    for layer_idx in range(args.num_layers):
+        set_layer_state_hf2mcore(args, mg_model, language_model, layer_idx)
+    mg_model.visual.model.load_state_dict(hf_model.model.visual.state_dict())
+
+
+def convert_mcore2hf_qwen2_5_vl(hf_model, mg_model):
+    language_model = hf_model.model.language_model
+    args = get_args()
+    # language_model
+    language_model.embed_tokens.weight.data.copy_(mg_model.embedding.word_embeddings.weight)
+    if args.untie_embeddings_and_output_weights:
+        hf_model.lm_head.weight.data.copy_(mg_model.output_layer.weight)
+    language_model.norm.weight.data.copy_(mg_model.decoder.final_layernorm.weight)
+    for layer_idx in range(args.num_layers):
+        set_layer_state_mcore2hf(args, mg_model, language_model, layer_idx)
+    hf_model.model.visual.load_state_dict(mg_model.visual.model.state_dict())

swift/megatron/model/qwen2_5_vl/vit.py

@@ -0,0 +1,90 @@
+from contextlib import contextmanager
+
+import torch
+from megatron.core.models.huggingface import HuggingFaceModule
+from megatron.training import get_args
+
+from swift.llm import get_model_tokenizer, to_device
+
+
+@contextmanager
+def patch_device_map_meta(model_cls):
+    __origin_init__ = model_cls.__init__
+
+    def __init__(self, *args, **kwargs):
+        with torch.device('meta'):
+            __origin_init__(self, *args, **kwargs)
+
+    model_cls.__init__ = __init__
+    try:
+        yield
+    finally:
+        model_cls.__init__ = __origin_init__
+
+
+class Qwen2_5VL_Vit(HuggingFaceModule):
+
+    def __init__(self, config):
+        from transformers.models.qwen2_5_vl import Qwen2_5_VLTextModel
+        super().__init__(config)
+        args = get_args()
+        model_dir = args.model_info.model_dir
+        kwargs = {'attn_impl': 'flash_attn'} if args.attention_backend.name == 'flash' else {}
+        with patch_device_map_meta(Qwen2_5_VLTextModel):
+            model, _ = get_model_tokenizer(model_dir, args.torch_dtype, return_dummy_model=True, **kwargs)
+        self.model = model.visual.to('cuda')
+        self.model_config = model.config
+
+    def forward(self, *args, **kwargs):
+        return self.model(*args, **kwargs)
+
+    def get_inputs_embeds(self, inputs_embeds, **kwargs):
+        input_ids = kwargs['input_ids']
+        pixel_values = kwargs.get('pixel_values')
+        pixel_values_videos = kwargs.get('pixel_values_videos')
+        image_grid_thw = kwargs.get('image_grid_thw')
+        video_grid_thw = kwargs.get('video_grid_thw')
+        dtype = self.model.dtype
+        if pixel_values is None and pixel_values_videos is None:  # plain-text
+            from PIL import Image
+            images = [Image.new('RGB', (32, 32), (0, 0, 0))]
+            media_inputs = self.processor.image_processor(images=images, return_tensors='pt')
+            device = input_ids.device
+            media_inputs = to_device(media_inputs, device)
+            pixel_values = media_inputs['pixel_values'].type(dtype)
+            image_embeds = self.model(pixel_values, grid_thw=media_inputs['image_grid_thw'])
+            inputs_embeds += image_embeds.mean() * 0.
+        else:
+            if pixel_values is None:
+                pixel_values_mixed = pixel_values_videos
+                grid_thw = video_grid_thw
+            elif pixel_values_videos is None:
+                pixel_values_mixed = pixel_values
+                grid_thw = image_grid_thw
+            else:
+                pixel_values_mixed = torch.concat([pixel_values, pixel_values_videos], dim=0)
+                grid_thw = torch.concat([image_grid_thw, video_grid_thw], dim=0)
+            pixel_values_mixed = pixel_values_mixed.type(dtype)
+            mixed_embeds = self.model(pixel_values_mixed, grid_thw=grid_thw)
+            if pixel_values is None:
+                image_embeds = None
+                video_embeds = mixed_embeds
+            elif pixel_values_videos is None:
+                image_embeds = mixed_embeds
+                video_embeds = None
+            else:
+                merge_length = self.processor.image_processor.merge_size**2
+                image_tokens = (image_grid_thw.prod(dim=-1) // merge_length).sum()
+                image_embeds = mixed_embeds[:image_tokens]
+                video_embeds = mixed_embeds[image_tokens:]
+
+            if image_embeds is not None:
+                image_mask = (input_ids == self.model_config.image_token_id).unsqueeze(-1).expand_as(inputs_embeds)
+                image_embeds = image_embeds.to(inputs_embeds.device, inputs_embeds.dtype)
+                inputs_embeds = inputs_embeds.masked_scatter(image_mask, image_embeds)
+
+            if video_embeds is not None:
+                video_mask = (input_ids == self.model_config.video_token_id).unsqueeze(-1).expand_as(inputs_embeds)
+                video_embeds = video_embeds.to(inputs_embeds.device, inputs_embeds.dtype)
+                inputs_embeds = inputs_embeds.masked_scatter(video_mask, video_embeds)
+        return inputs_embeds

swift/megatron/model/register.py

@@ -1,7 +1,7 @@
 # Copyright (c) Alibaba, Inc. and its affiliates.
 from argparse import ArgumentParser
 from dataclasses import dataclass
-from typing import Any, Callable, Dict, List, Optional
+from typing import Any, Callable, Dict, List, Optional, Type
 
 import torch.nn as nn
 from transformers import PretrainedConfig
@@ -20,6 +20,7 @@ class MegatronModelMeta:
     convert_hf_config: Callable[[PretrainedConfig], Dict[str, Any]]
     convert_mcore2hf: Callable[[nn.Module, nn.Module], None]
     convert_hf2mcore: Callable[[nn.Module, nn.Module], None]
+    visual: Optional[Type[nn.Module]] = None
 
     extra_args_provider: Optional[Callable[[ArgumentParser], ArgumentParser]] = None