# Copyright 2025 the LlamaFactory team. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import os from dataclasses import dataclass from typing import Any import fire import torch from peft import PeftModel from torch.utils.data import Dataset from transformers import DataCollatorForSeq2Seq, Qwen2_5_VLProcessor from llamafactory.extras.constants import IGNORE_INDEX from llamafactory.hparams import get_train_args from llamafactory.model import load_model, load_tokenizer from llamafactory.train.callbacks import LogCallback from llamafactory.train.sft.trainer import CustomSeq2SeqTrainer class DummyDataset(Dataset): def __init__(self, size: int = 1000, seq_length: int = 1024, processor: Qwen2_5_VLProcessor = None): self.size = size self.seq_length = seq_length self.vocab_size = 32768 self.processor = processor image_token_num = 18 * 18 // (2 * 2) image_t = 2 self.text_seqlen = seq_length // 4 # 25% text video_seq_length = self.seq_length - self.text_seqlen - image_t * image_token_num video_t = video_seq_length // image_token_num self.image_size = [18 * 18 * image_t, 1176] self.image_grid_thw = torch.tensor([[1, 18, 18]] * image_t, dtype=torch.long) self.image_seqlen = image_t * image_token_num self.video_size = [18 * 18 * video_t, 1176] self.video_grid_thw = torch.tensor([[video_t, 18, 18]], dtype=torch.long) self.video_seqlen = video_t * image_token_num def __len__(self): return self.size def __getitem__(self, index: int): input_ids = torch.randint(low=0, high=self.vocab_size, size=(self.seq_length,)) input_ids[: self.image_seqlen] = self.processor.image_token_id input_ids[self.image_seqlen : self.image_seqlen + self.video_seqlen] = self.processor.video_token_id attention_mask = torch.ones((self.seq_length,), dtype=torch.long) labels = input_ids.clone() labels[: self.image_seqlen + self.video_seqlen] = IGNORE_INDEX pixel_values = torch.rand(self.image_size, dtype=torch.float32) pixel_values_videos = torch.rand(self.video_size, dtype=torch.float32) return { "input_ids": input_ids, "attention_mask": attention_mask, "labels": labels, "pixel_values": pixel_values, "pixel_values_videos": pixel_values_videos, "image_grid_thw": self.image_grid_thw, "video_grid_thw": self.video_grid_thw, } @dataclass class MultiModalDataCollatorForSeq2Seq(DataCollatorForSeq2Seq): def __post_init__(self): if isinstance(self.model, PeftModel): self.model = self.model.base_model.model if self.model is not None and hasattr(self.model, "get_rope_index"): # for qwen2vl mrope self.get_rope_func = self.model.get_rope_index # transformers < 4.52.0 or qwen2.5 omni elif self.model is not None and hasattr(self.model, "model") and hasattr(self.model.model, "get_rope_index"): self.get_rope_func = self.model.model.get_rope_index # transformers >= 4.52.0 else: self.get_rope_func = None def __call__(self, features: list[dict[str, Any]]) -> dict[str, "torch.Tensor"]: batch_pixel_values = [feature.pop("pixel_values") for feature in features] batch_pixel_values_videos = [feature.pop("pixel_values_videos") for feature in features] batch_image_grid_thw = [feature.pop("image_grid_thw") for feature in features] batch_video_grid_thw = [feature.pop("video_grid_thw") for feature in features] batch: dict[str, torch.Tensor] = super().__call__(features) batch["pixel_values"] = torch.cat(batch_pixel_values, dim=0) batch["pixel_values_videos"] = torch.cat(batch_pixel_values_videos, dim=0) batch["image_grid_thw"] = torch.cat(batch_image_grid_thw, dim=0) batch["video_grid_thw"] = torch.cat(batch_video_grid_thw, dim=0) if self.get_rope_func is not None: rope_index_kwargs = { "input_ids": batch["input_ids"], "image_grid_thw": batch["image_grid_thw"], "video_grid_thw": batch["video_grid_thw"], "attention_mask": (batch["attention_mask"] >= 1).float(), } batch["position_ids"], batch["rope_deltas"] = self.get_rope_func(**rope_index_kwargs) if "position_ids" not in batch or batch["position_ids"].dim() != 3: raise ValueError("Qwen2VL requires 3D position ids for mrope.") return batch def bench_qwen( model_name_or_path: str = "Qwen/Qwen2-VL-7B-Instruct", batch_size: int = 1, seq_length: int = 2048, liger_kernel: bool = False, deepspeed_stage: int = 3, ): os.environ["LLAMABOARD_ENABLED"] = "true" os.environ["LLAMABOARD_WORKDIR"] = "output/dummy_dir" args = { "model_name_or_path": model_name_or_path, "enable_liger_kernel": liger_kernel, "stage": "sft", "do_train": True, "finetuning_type": "full", "dataset": "alpaca_en_demo", "template": "qwen2_vl", "cutoff_len": seq_length, "output_dir": "output/dummy_dir", "logging_steps": 10, "save_strategy": "no", "save_only_model": True, "overwrite_output_dir": True, "per_device_train_batch_size": batch_size, "max_steps": 1000, "bf16": True, "include_num_input_tokens_seen": True, "report_to": "none", } if deepspeed_stage in [2, 3]: args["deepspeed"] = f"examples/deepspeed/ds_z{deepspeed_stage}_config.json" model_args, _, training_args, finetuning_args, _ = get_train_args(args) tokenizer_module = load_tokenizer(model_args) tokenizer = tokenizer_module["tokenizer"] trainset = DummyDataset(size=100000, seq_length=seq_length, processor=tokenizer_module["processor"]) model = load_model(tokenizer, model_args, finetuning_args, training_args.do_train) data_collator = MultiModalDataCollatorForSeq2Seq( tokenizer=tokenizer, model=model, pad_to_multiple_of=8, label_pad_token_id=IGNORE_INDEX ) trainer = CustomSeq2SeqTrainer( model=model, args=training_args, finetuning_args=finetuning_args, data_collator=data_collator, callbacks=[LogCallback()], train_dataset=trainset, **tokenizer_module, ) trainer.train(resume_from_checkpoint=training_args.resume_from_checkpoint) if __name__ == "__main__": fire.Fire(bench_qwen)