from unsloth import FastVisionModel # FastLanguageModel for LLMs import torch from datasets import load_dataset from unsloth import is_bf16_supported from unsloth.trainer import UnslothVisionDataCollator from trl import SFTTrainer, SFTConfig from transformers import TextStreamer from PIL import Image import requests from io import BytesIO import datetime timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S") model, tokenizer = FastVisionModel.from_pretrained( model_name = "/home/rzhong/project/FSTSPrune/model_pretrain_sft_20250303_125849-Pruned-hf", load_in_4bit = False, # Use 4bit to reduce memory use. False for 16bit LoRA. # 模型已经选了4bit量化后的，这里还需不需要再以4bit加载？建议实验一下 | 应该是一个意思 use_gradient_checkpointing = "unsloth", # True or "unsloth" for long context max_seq_length = 2048, # unsloth支持4x的上下文微调，如果原模型支持8192的上下文，这里只需要设置为2048 dtype = torch.bfloat16, # A100支持bfloat16，可以减少显存占用。默认为None，也可以选择torch.float16 ) model = FastVisionModel.get_peft_model( model, finetune_vision_layers = True, # False if not finetuning vision layers finetune_language_layers = True, # False if not finetuning language layers finetune_attention_modules = True, # False if not finetuning attention layers finetune_mlp_modules = True, # False if not finetuning MLP layers r = 16, # The larger, the higher the accuracy, but might overfit lora_alpha = 16, # Recommended alpha == r at least lora_dropout = 0, bias = "none", random_state = 3407, use_rslora = False, # We support rank stabilized LoRA loftq_config = None, # And LoftQ # target_modules = "all-linear", # Optional now! Can specify a list if needed target_modules = ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"], ) # 数据集1处理 def process_dataset1(): dataset = load_dataset("/home/share/rzhong/dataset/Qwen10k", split="train").select(range(200)) def convert(sample): return { "messages": [ { "role": "user", "content": [ {"type": "text", "text": sample["question"]}, {"type": "image", "image": sample["image"]} ] }, { "role": "assistant", "content": [{"type": "text", "text": sample["chosen"]}] } ] } return [convert(sample) for sample in dataset] # 数据集2处理（带图片下载和压缩） def process_dataset2(): dataset = load_dataset( "json", data_files="/home/share/rzhong/dataset/ALLaVA-4V-Chinese/allava_laion/ALLaVA-Instruct-LAION-4V_Chinese.json", split="train" ).select(range(200)) def convert(sample): try: # 下载并处理图片 url = sample['url'].split('?')[0] response = requests.get(url, timeout=5) response.raise_for_status() image = Image.open(BytesIO(response.content)) # 图片压缩处理 image = image.resize((320, 240), Image.LANCZOS).convert('RGB') return { "messages": [ { "role": "user", "content": [ {"type": "text", "text": sample['conversations'][0]['value']}, {"type": "image", "image": image} ] }, { "role": "assistant", "content": [{"type": "text", "text": sample['conversations'][1]['value']}] } ] } except Exception as e: print(f"跳过样本: {url}，错误: {str(e)}") return None # 过滤无效样本 return [conv for sample in dataset if (conv := convert(sample)) is not None] # 数据集3处理 def process_dataset3(): dataset = load_dataset("/home/share/rzhong/dataset/google-landmark/dataset_4/dataset_file", split="train") instruction = "描述这张图片。" def convert(sample): return { "messages": [ { "role": "user", "content": [ {"type": "text", "text": instruction}, {"type": "image", "image": sample["image"]} ] }, { "role": "assistant", "content": [{"type": "text", "text": sample["text"]}] } ] } return [convert(sample) for sample in dataset] # 数据集4处理 def process_dataset4(): dataset = load_dataset("/home/share/rzhong/dataset/R1-Vision-PixMo-Cap-QA-zh", split="train").select(range(200)) def convert(sample): try: image = Image.open(BytesIO(requests.get(sample['image_url'], timeout=5).content)) image = image.resize((320, 240), Image.LANCZOS).convert('RGB') return { "messages": [ { "role": "user", "content": [ {"type": "text", "text": sample['question']}, {"type": "image", "image": image} ] }, { "role": "assistant", "content": [{"type": "text", "text": sample['r1_solution']}] } ] } except Exception as e: print(f"跳过样本: {sample['image_url']}，错误: {str(e)}") return None return [conv for sample in dataset if (conv := convert(sample)) is not None] # 合并所有数据集 def merge_datasets(): dataset1 = process_dataset1() dataset2 = process_dataset2() dataset3 = process_dataset3() dataset4 = process_dataset4() merged_dataset = dataset1 + dataset2 + dataset3 + dataset4 print(f"总样本数: {len(merged_dataset)}") print("前3个样本预览:") for i in range(min(3, len(merged_dataset))): print(merged_dataset[i]) return merged_dataset # 执行合并 merged_data = merge_datasets() FastVisionModel.for_training(model) # Enable for training! trainer = SFTTrainer( model = model, tokenizer = tokenizer, data_collator = UnslothVisionDataCollator(model, tokenizer), # Must use! train_dataset = merged_data, args = SFTConfig( per_device_train_batch_size = 2, gradient_accumulation_steps = 4, # 原来是4。可以增加，相当于提高batch size，但不会影响内存消耗。增加会使loss曲线更平滑 warmup_steps = 500, # max_steps = None, num_train_epochs = 10, # Set this instead of max_steps for full training runs learning_rate = 5e-5, # 2e-4 1e-4 5e-5 2e-5 fp16 = not is_bf16_supported(), bf16 = is_bf16_supported(), logging_steps = 1, optim = "adamw_8bit", weight_decay = 0.01, lr_scheduler_type = "linear", seed = 3407, output_dir = "model_pretrain_sft_20250303_125849-Pruned-hf-lora-output6", report_to = "none", # For Weights and Biases # You MUST put the below items for vision finetuning: remove_unused_columns = False, dataset_text_field = "", dataset_kwargs = {"skip_prepare_dataset": True}, dataset_num_proc = 16, max_seq_length = 2048, ), ) trainer_stats = trainer.train() model.save_pretrained("model_pretrain_sft_20250303_125849-Pruned-hf-lora6") # Local saving tokenizer.save_pretrained("model_pretrain_sft_20250303_125849-Pruned-hf-lora6")