QLoRA 微调Qwen1.5-0.5B-Chat

参考文章：

https://huggingface.co/blog/4bit-transformers-bitsandbytes

https://github.com/artidoro/qlora/tree/main

本文实战使用QLoRA技术微调阿里的Qwen1.5-0.5B-Chat模型，采用single-gpu 进行训练。

1. 核心Python包【python版本:3.10.0】

torch 2.2.2+cu118
accelerate 0.33.0
bitsandbytes 0.43.3
transformers 4.37.0

2. 使用数据集

https://github.com/DB-lost/self-llm/blob/master/dataset/huanhuan.json

3. 具体实现代码

# coding:utf-8
"""QLoRA Finetune Qwen1.5-0.5B-Chat"""from transformers import AutoTokenizer, AutoConfig, AutoModelForCausalLM, TrainingArguments, Trainer, BitsAndBytesConfig
from torch.utils.data import Dataset
import torch
from peft import LoraConfig, TaskType, get_peft_model, prepare_model_for_kbit_training
from typing import Dict
import transformers
import json
from transformers.trainer_pt_utils import LabelSmootherIGNORE_TOKEN_ID = LabelSmoother.ignore_indexmax_len = 512
data_json = json.load(open("./data/huanhuan.json", 'r', encoding='utf-8'))
train_json = []
lazy_preprocess = True
gradient_checkpointing = True
TEMPLATE = "{% for message in messages %}{% if loop.first and messages[0]['role'] != 'system' %}{{ '<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n' }}{% endif %}{{'<|im_start|>' + message['role'] + '\n' + message['content']}}{% if loop.last %}{{ '<|im_end|>'}}{% else %}{{ '<|im_end|>\n' }}{% endif %}{% endfor %}"def print_model_allarguments_name_dtype(model):for n, v in model.named_parameters():if v.requires_grad:print(f"trainable model arguments:{n}--{v.dtype}--{v.shape}")else:print(f"not trainable model arguments:{n}--{v.dtype}--{v.shape}")config = AutoConfig.from_pretrained("./models/Qwen1.5-0.5B-Chat",trust_remote_code=True)# kv cache 在推理的时候才用，训练时候不用
config.use_cache = Falsetokenizer = AutoTokenizer.from_pretrained("./models/Qwen1.5-0.5B-Chat",model_max_length=max_len,padding_side="right",use_fast=False
)model = AutoModelForCausalLM.from_pretrained("./models/Qwen1.5-0.5B-Chat",torch_dtype=torch.bfloat16,device_map="auto",quantization_config=BitsAndBytesConfig(load_in_4bit=True,bnb_4bit_use_double_quant=True,bnb_4bit_quant_type="nf4",bnb_4bit_compute_dtype=torch.bfloat16,),config=config,low_cpu_mem_usage=True
)print("Original Model: ")
print_model_allarguments_name_dtype(model)model = prepare_model_for_kbit_training(model, use_gradient_checkpointing=gradient_checkpointing)
print("kbit training: ")
print_model_allarguments_name_dtype(model)config = LoraConfig(task_type=TaskType.CAUSAL_LM,target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "up_proj", "gate_proj", "down_proj"],r=64, # Lora 秩lora_alpha=16, # Lora alaph，具体作用参见 Lora 原理lora_dropout=0.05, # Dropout 比例bias='none'
)
model = get_peft_model(model, config)
print("LoRA Model: ")
print_model_allarguments_name_dtype(model)
model.print_trainable_parameters()"""
这个函数调用启用了模型的梯度检查点。
梯度检查点是一种优化技术，可用于减少训练时的内存消耗。
通常，在反向传播期间，模型的中间激活值需要被保留以计算梯度。
启用梯度检查点后，系统只需在需要时计算和保留一部分中间激活值，从而减少内存需求。
这对于处理大型模型或限制内存的环境中的训练任务非常有用。
"""
if gradient_checkpointing:model.enable_input_require_grads()def preprocess(messages,tokenizer: transformers.PreTrainedTokenizer,max_len: int,
) -> Dict:"""Preprocesses the data for supervised fine-tuning."""texts = []for i, msg in enumerate(messages):texts.append(tokenizer.apply_chat_template(msg,chat_template=TEMPLATE,tokenize=True,add_generation_prompt=False,padding=True,max_length=max_len,truncation=True,))input_ids = torch.tensor(texts, dtype=torch.long)target_ids = input_ids.clone()target_ids[target_ids == tokenizer.pad_token_id] = IGNORE_TOKEN_IDattention_mask = input_ids.ne(tokenizer.pad_token_id)return dict(input_ids=input_ids, target_ids=target_ids, attention_mask=attention_mask)class LazySupervisedDataset(Dataset):"""Dataset for supervised fine-tuning."""def __init__(self, raw_data, tokenizer: transformers.PreTrainedTokenizer, max_len: int):super(LazySupervisedDataset, self).__init__()self.tokenizer = tokenizerself.max_len = max_lenself.tokenizer = tokenizerself.raw_data = raw_dataself.cached_data_dict = {}def __len__(self):return len(self.raw_data)def __getitem__(self, i) -> Dict[str, torch.Tensor]:if i in self.cached_data_dict:return self.cached_data_dict[i]ret = preprocess([self.raw_data[i]["conversations"]], self.tokenizer, self.max_len)ret = dict(input_ids=ret["input_ids"][0],labels=ret["target_ids"][0],attention_mask=ret["attention_mask"][0],)self.cached_data_dict[i] = retreturn retclass SupervisedDataset(Dataset):"""Dataset for supervised fine-tuning."""def __init__(self, raw_data, tokenizer: transformers.PreTrainedTokenizer, max_len: int):super(SupervisedDataset, self).__init__()sources = [example["conversations"] for example in raw_data]data_dict = preprocess(sources, tokenizer, max_len)self.input_ids = data_dict["input_ids"]self.labels = data_dict["labels"]self.attention_mask = data_dict["attention_mask"]def __len__(self):return len(self.input_ids)def __getitem__(self, i) -> Dict[str, torch.Tensor]:return dict(input_ids=self.input_ids[i],labels=self.labels[i],attention_mask=self.attention_mask[i],)for i, d in enumerate(data_json):t = {"id": f"identity_{i}","conversations": [{"role": "user","content": d['instruction'] + d['input']},{"role": "assistant","content": d['output']}]}train_json.append(t)dataset_cls = (LazySupervisedDataset if lazy_preprocess else SupervisedDataset
)train_dataset = dataset_cls(train_json, tokenizer=tokenizer, max_len=max_len)
eval_dataset = None
data_module = dict(train_dataset=train_dataset, eval_dataset=eval_dataset)args = TrainingArguments(output_dir="./output/Qwen1.5",per_device_train_batch_size=2,per_device_eval_batch_size=1,gradient_accumulation_steps=8,logging_steps=10,weight_decay=0.01,adam_beta2=0.95,num_train_epochs=5,save_steps=100,learning_rate=3e-4,save_on_each_node=True,gradient_checkpointing=True,lr_scheduler_type='cosine',warmup_ratio=0.01
)
trainer = Trainer(model=model,args=args,**data_module
)
trainer.train()