support sp in vlm (#71)

* feat: add multimodal sequence parallelism support

* refactor: separate SFT/DPO training scripts and optimize data handling

- Remove data printing functions from SFT and DPO trainers for better performance
- Replace 360-example-vl.sh with separate SFT and DPO training scripts
- Add SFT visual-language demo dataset (data/sft-vl-demo/)
- Update dataset configuration to support new data structure

* refactor: restructure multimodal model forward functions and optimize code style

- Add multimodal_forwards module to centrally manage multimodal model forward logic
- Extract and optimize forward function implementations for Qwen2 VL and Qwen2.5 VL
- Improve sequence_parallel related code structure

* feat(vl): update VL training scripts and clean up demo data

* refactor: improve readability of sequence parallel attention check

---------

Co-authored-by: lilin3 <lilin3@360.cn>

Author: aleien95

360-example-vl-dpo.sh

@@ -0,0 +1,68 @@
+#!/bin/bash
+
+set -x
+
+# Environment setup
+export DS_SKIP_CUDA_CHECK=1
+export DISABLE_VERSION_CHECK=1
+export FORCE_TORCHRUN=1
+export CUDA_LAUNCH_BLOCKING=1
+
+# Parameters
+MODEL_PATH=""
+MODEL_SIZE="7B"
+DATA_NAME="dpo-vl-demo"
+NUM_NODES=1
+NUM_GPUS=8
+CUTOFF_LEN=20000
+LEARNING_RATE=1e-6
+PER_DEVICE_BATCH_SIZE=1
+GRADIENT_ACCUMULATION_STEPS=16
+
+# Output directory
+model_saved_name="demo-qwen25vl-${MODEL_SIZE}-len_${CUTOFF_LEN}-lr_${LEARNING_RATE}-data_${DATA_NAME}"
+OUTPUT_DIR="./output/${MODEL_SIZE}/${model_saved_name}"
+tensorboard_dir="${OUTPUT_DIR}/runs_${MODEL_SIZE}/${model_saved_name}"
+
+# Create directories
+mkdir -p ${OUTPUT_DIR}
+mkdir -p ${tensorboard_dir}
+
+# DPO Training
+deepspeed --hostfile=/etc/mpi.host src/train.py \
+    --stage dpo \
+    --do_train \
+    --model_name_or_path ${MODEL_PATH} \
+    --dataset ${DATA_NAME} \
+    --dataset_dir ./data \
+    --template qwen2_vl \
+    --finetuning_type full \
+    --freeze_vision_tower True \
+    --train_mm_proj_only False \
+    --image_resolution 1048576 \
+    --video_resolution 16384 \
+    --pref_beta 0.1 \
+    --pref_ftx 0.0 \
+    --output_dir ${OUTPUT_DIR} \
+    --overwrite_cache \
+    --overwrite_output_dir True \
+    --cutoff_len ${CUTOFF_LEN} \
+    --preprocessing_num_workers 128 \
+    --per_device_train_batch_size ${PER_DEVICE_BATCH_SIZE} \
+    --gradient_accumulation_steps ${GRADIENT_ACCUMULATION_STEPS} \
+    --learning_rate ${LEARNING_RATE} \
+    --lr_scheduler_type cosine_with_min_lr \
+    --lr_scheduler_kwargs "{\"min_lr_rate\": 0.1}" \
+    --num_train_epochs 1 \
+    --warmup_ratio 0.05 \
+    --logging_steps 1 \
+    --logging_dir "./output/runs_${MODEL_SIZE}/${model_saved_name}" \
+    --save_strategy epoch \
+    --plot_loss True \
+    --deepspeed examples/deepspeed/ds_z2_config.json \
+    --use_unsloth_gc True \
+    --bf16 \
+    --flash_attn fa2 \
+    --sequence_parallel_size 8 \
+    --ddp_timeout 180000000 \
+    --report_to tensorboard

360-example-vl-sft.sh

@@ -0,0 +1,66 @@
+#!/bin/bash
+
+set -x
+
+# Environment setup
+export DS_SKIP_CUDA_CHECK=1
+export DISABLE_VERSION_CHECK=1
+export FORCE_TORCHRUN=1
+export CUDA_LAUNCH_BLOCKING=1
+
+# Parameters
+MODEL_PATH=""
+MODEL_SIZE="7B"
+DATA_NAME="sft-vl-demo"
+NUM_NODES=1
+NUM_GPUS=8
+CUTOFF_LEN=20000
+LEARNING_RATE=6e-5
+PER_DEVICE_BATCH_SIZE=1
+GRADIENT_ACCUMULATION_STEPS=16
+
+# Output directory
+model_saved_name="demo-qwen25vl-${MODEL_SIZE}-len_${CUTOFF_LEN}-lr_${LEARNING_RATE}-data_${DATA_NAME}"
+OUTPUT_DIR="./output/${MODEL_SIZE}/${model_saved_name}"
+tensorboard_dir="${OUTPUT_DIR}/runs_${MODEL_SIZE}/${model_saved_name}"
+
+# Create directories
+mkdir -p ${OUTPUT_DIR}
+mkdir -p ${tensorboard_dir}
+
+# SFT Training
+deepspeed --hostfile=/etc/mpi.host src/train.py \
+    --stage sft \
+    --do_train \
+    --model_name_or_path ${MODEL_PATH} \
+    --dataset ${DATA_NAME} \
+    --dataset_dir ./data \
+    --template qwen2_vl \
+    --finetuning_type full \
+    --freeze_vision_tower True \
+    --train_mm_proj_only False \
+    --image_resolution 1048576 \
+    --video_resolution 16384 \
+    --output_dir ${OUTPUT_DIR} \
+    --overwrite_cache \
+    --overwrite_output_dir True \
+    --cutoff_len ${CUTOFF_LEN} \
+    --preprocessing_num_workers 128 \
+    --per_device_train_batch_size ${PER_DEVICE_BATCH_SIZE} \
+    --gradient_accumulation_steps ${GRADIENT_ACCUMULATION_STEPS} \
+    --learning_rate ${LEARNING_RATE} \
+    --lr_scheduler_type cosine_with_min_lr \
+    --lr_scheduler_kwargs "{\"min_lr_rate\": 0.1}" \
+    --num_train_epochs 1 \
+    --warmup_ratio 0.05 \
+    --logging_steps 1 \
+    --logging_dir "./output/runs_${MODEL_SIZE}/${model_saved_name}" \
+    --save_strategy epoch \
+    --plot_loss True \
+    --deepspeed examples/deepspeed/ds_z2_config.json \
+    --use_unsloth_gc True \
+    --bf16 \
+    --flash_attn fa2 \
+    --sequence_parallel_size 8 \
+    --ddp_timeout 180000000 \
+    --report_to tensorboard

data/dataset_info.json

@@ -624,5 +624,32 @@
       "prompt": "content"
     },
     "folder": "python"
+  },
+  "sft-vl-demo": {
+      "file_name": "sft-vl-demo/train.jsonl",
+      "formatting": "sharegpt",
+      "columns": {
+          "messages": "conversations",
+          "images": "images",
+          "videos": "videos"
+      },
+      "tags": {
+          "role_tag": "from",
+          "content_tag": "value",
+          "user_tag": "human",
+          "assistant_tag": "assistant"                                                                                                                                             
+      }
+  },
+  "dpo-vl-demo": {                                                                                                                                        
+      "file_name": "dpo-vl-demo/train.jsonl",
+      "ranking": true,
+      "formatting": "sharegpt",
+      "columns": {
+          "messages": "conversations",
+          "chosen": "chosen",
+          "rejected": "rejected",
+          "images": "images",
+          "videos": "videos"
+      }
   }
 }

src/llamafactory/data/collator.py

@@ -154,6 +154,8 @@ def __call__(self, features: Sequence[Dict[str, Any]]) -> Dict[str, "torch.Tenso
                 if self.require_position_ids:
                     # if requires, would be padded to cutoff_len in preprocessing
                     target_feature["position_ids"] = feature[f"{key}_position_ids"]
+                if "image_position_maps" in feature:
+                    target_feature["image_position_maps"] = feature[f"image_position_maps"]
                 concatenated_features.append(target_feature)
 
         return super().__call__(concatenated_features)

src/llamafactory/data/preprocess.py

@@ -119,9 +119,9 @@ def get_sequence_parallel_preprocess(
     tokenizer: "PreTrainedTokenizer",
 ) -> Tuple[Callable, Callable]:
     if stage == "pad":
-        preprocess_func = partial(pad_sequence, data_args=data_args, tokenizer=tokenizer)
+        preprocess_func = partial(pad_sequence, data_args=data_args, tokenizer=tokenizer, model_args=model_args)
     elif stage == "split":
-        preprocess_func = partial(sp_split, model_args=model_args)
+        preprocess_func = partial(sp_split, model_args=model_args, tokenizer=tokenizer)
     else:
         raise NotImplementedError(f"Unexpected stage in sequence_parallel_preprocess: {stage}")
 

src/llamafactory/data/processors/sequence_parallel.py

@@ -2,12 +2,28 @@
 from ..data_utils import preprocess_sp_dataset
 
 
-def pad_sequence(examples, data_args, tokenizer):
+def get_max_lengths(examples):
+    valid_lists = []
+    for key, value in examples.items():
+        if key.endswith('input_ids') and value is not None:
+            valid_lists.append(value)
+    
+    if not valid_lists:
+        return []
+    
+    max_lengths = [max(len(lst) if lst is not None else 0 for lst in group) 
+                   for group in zip(*valid_lists)]
+    
+    return max_lengths
+
+
+def pad_sequence(examples, data_args, tokenizer, model_args):
     max_length = data_args.cutoff_len
     input_pad_token_id = tokenizer.pad_token_id
     assert data_args.ignore_pad_token_for_loss
     label_pad_token_id = IGNORE_INDEX if data_args.ignore_pad_token_for_loss else tokenizer.pad_token_id
 
+    max_input_ids_length_list = get_max_lengths(examples)
     for k, v in examples.items():
         if k.endswith("input_ids"):
             pad_token_id = input_pad_token_id
@@ -25,25 +41,59 @@ def pad_sequence(examples, data_args, tokenizer):
             continue  # TODO: haven't tested multi-modal yet
         else:
             raise NotImplementedError(f"Unexpected dataset key: {k}")
+
         for i in range(len(v)):
-            v[i].extend([pad_token_id] * (max_length - len(v[i])))
+             tmp_sp_len = max_input_ids_length_list[i] // model_args.sequence_parallel_size
+             closest_cutoff_len = int(tmp_sp_len + (8 - tmp_sp_len % 8)) * model_args.sequence_parallel_size
+             max_length = min(closest_cutoff_len, data_args.cutoff_len)
+
+             v[i].extend([pad_token_id] * (max_length - len(v[i])))
         examples[k] = v
 
     return examples
 
 
+def create_image_position_info(seq_ids, image_token_id):
+    """为整个序列创建图像位置信息"""
+    info = []
+    global_image_pos = 0  # 全局连续的图像位置计数器
+    
+    for token_id in seq_ids:
+        if token_id == image_token_id:
+            info.append(global_image_pos)
+            global_image_pos += 1
+        else:
+            info.append(-1)
+    return info
+
+
 # sp for Sequence Parallel
-def sp_split(examples, model_args):
+def sp_split(examples, model_args, tokenizer):
+    all_image_position_maps = list()
+    new_examples = dict()
+
     for k, v in examples.items():
         chunks = list()
         for row in v:
-            if k.endswith("attention_mask"):
-                chunks.extend([row] * model_args.sequence_parallel_size)
-            elif row is None:
+            if row is None:
                 chunks.extend([None] * model_args.sequence_parallel_size)
+            elif k in ['images']:
+                chunks.extend([row] * model_args.sequence_parallel_size)
             else:
                 chunks.extend(
                     preprocess_sp_dataset(row, model_args.sequence_parallel_size, model_args.sequence_parallel_mode)
                 )
-        examples[k] = chunks
-    return examples
+                if k.endswith("input_ids") and len(all_image_position_maps) < (len(v) * model_args.sequence_parallel_size):
+                    image_position_info = create_image_position_info(row, tokenizer.image_token_id)
+                    all_image_position_maps.extend(
+                        preprocess_sp_dataset(image_position_info, model_args.sequence_parallel_size, model_args.sequence_parallel_mode)
+                    )
+        new_examples[k] = chunks
+    
+    if len(all_image_position_maps)>0:
+        new_examples['image_position_maps'] = all_image_position_maps
+        for index in range(len(new_examples['images'])):
+            if all(image_position==-1 for image_position in new_examples['image_position_maps'][index]):
+                new_examples['images'][index] = None
+
+    return new_examples 

src/llamafactory/hparams/parser.py

@@ -41,7 +41,7 @@
 logger = logging.get_logger(__name__)
 
 
-check_dependencies()
+# check_dependencies()
 
 
 _TRAIN_ARGS = [ModelArguments, DataArguments, Seq2SeqTrainingArguments, FinetuningArguments, GeneratingArguments]

src/llamafactory/model/loader.py

@@ -87,6 +87,9 @@ def load_tokenizer(model_args: "ModelArguments") -> "TokenizerModule":
     except Exception as e:
         raise OSError("Failed to load tokenizer.") from e
 
+    if hasattr(config, 'image_token_id'):
+        tokenizer.image_token_id = config.image_token_id
+
     if model_args.new_special_tokens is not None:
         num_added_tokens = tokenizer.add_special_tokens(
             dict(additional_special_tokens=model_args.new_special_tokens),
@@ -144,7 +147,7 @@ def load_model(
         config.attention_dropout = 0.0
 
     apply_liger_kernel(config, model_args, is_trainable, require_logits=(finetuning_args.stage not in ["pt", "sft"]))
-    sequence_parallel_group = apply_sequence_parallel(model_args, full_determinism)  # monkey patching, similar to liger_kernel
+    sequence_parallel_group = apply_sequence_parallel(model_args, config, full_determinism)  # monkey patching, similar to liger_kernel
 
     model = None
     lazy_load = False
@@ -157,7 +160,9 @@ def load_model(
     if model is None and not lazy_load:
         init_kwargs["config"] = config
         init_kwargs["pretrained_model_name_or_path"] = model_args.model_name_or_path
-        if sequence_parallel_group is not None and is_transformers_version_greater_than("4.51.0"):
+        if (sequence_parallel_group is not None 
+            and is_transformers_version_greater_than("4.51.0")
+            and config.model_type not in ['qwen2_vl', 'qwen2_5_vl']):
             init_kwargs["attn_implementation"] = "sequence_parallel_attention"
 
         if model_args.mixture_of_depths == "load":

src/llamafactory/model/model_utils/checkpointing.py

@@ -68,7 +68,9 @@ def backward(ctx: "torch.autograd.Function", grad_output: "torch.Tensor") -> "to
             hidden_states = hidden_states.to("cuda", non_blocking=True).detach()
             hidden_states.requires_grad_(True)
             with torch.enable_grad():
-                (output,) = ctx.forward_function(hidden_states, *ctx.args)
+                #(output,) = ctx.forward_function(hidden_states, *ctx.args)
+                outputs = ctx.forward_function(hidden_states, *ctx.args)
+                output = outputs[0] if isinstance(outputs, tuple) else outputs
 
             torch.autograd.backward(output, grad_output)
             return (None, hidden_states.grad) + (None,) * len(ctx.args)

src/llamafactory/model/model_utils/multimodal_forwards/__init__.py

@@ -0,0 +1,8 @@
+#!/usr/bin/python
+#!-*-coding:utf8-*-
+
+
+from .qwen2_vl_forward import patched_qwen2_vl_forward
+from .qwen2_5_vl_forward import patched_qwen2_5_vl_forward
+
+__all__ = ['patched_qwen2_vl_forward', 'patched_qwen2_5_vl_forward']