import logging import torch from datasets import load_dataset from peft import get_peft_model, LoraConfig, TaskType from transformers import ( AutoTokenizer, AutoModelForCausalLM, Trainer, TrainingArguments, DataCollatorForLanguageModeling, ) logging.basicConfig(level=logging.INFO) logger = logging.getLogger(__name__) class LoRATrainer: def __init__(self, config): self.config = config self.model_path = config['model_path'] self.train_file = config['train_file'] self.output_dir = config['output_dir'] self.max_steps = config.get('max_steps', 100) self._load_tokenizer() self._load_dataset() self._load_model() def _load_tokenizer(self): logger.info("[1] 加载 Tokenizer...") self.tokenizer = AutoTokenizer.from_pretrained( self.model_path, trust_remote_code=True ) if self.tokenizer.pad_token is None: self.tokenizer.pad_token = self.tokenizer.eos_token def _load_dataset(self): logger.info("[2] 加载并 Tokenize 数据...") raw_dataset = load_dataset("text", data_files={"train": self.train_file}) def tokenize_function(example): return self.tokenizer( example["text"], truncation=True, max_length=2048, # 移除padding padding=False ) self.tokenized_dataset = raw_dataset.map( tokenize_function, batched=True, remove_columns=["text"] )["train"] def _load_model(self): logger.info("[3] 加载模型 + LoRA 适配...") # 配置4-bit量化 self.model = AutoModelForCausalLM.from_pretrained( self.model_path, device_map="auto", torch_dtype=torch.float16, trust_remote_code=True ) lora_config = LoraConfig( r=8, lora_alpha=32, target_modules=["q_proj", "v_proj"], lora_dropout=0.05, bias="none", task_type=TaskType.CAUSAL_LM ) self.model = get_peft_model(self.model, lora_config) # 启用梯度检查点 self.model.gradient_checkpointing_enable() self.data_collator = DataCollatorForLanguageModeling( tokenizer=self.tokenizer, mlm=False, pad_to_multiple_of=8 # 提升GPU效率 ) def train(self): logger.info("[4] 开始 LoRA 微调训练...") training_args = TrainingArguments( output_dir=self.output_dir, per_device_train_batch_size=1, gradient_accumulation_steps=4, logging_steps=10, max_steps=self.max_steps, save_strategy="steps", save_steps=50, # 定期保存 report_to="none", fp16=True # 混合精度训练 ) trainer = Trainer( model=self.model, args=training_args, train_dataset=self.tokenized_dataset, tokenizer=self.tokenizer, data_collator=self.data_collator ) trainer.train() if __name__ == "__main__": config = { "model_path": r"C:\Users\vipuser\Documents\ai_writer_project_final_with_fixed_output_ui\models\Yi-6B", "train_file": "./data/train_cleaned.txt", "output_dir": "./output/lora", "max_steps": 100 } trainer = LoRATrainer(config) trainer.train()

<think>我们正在使用LoRA对Yi-6B模型进行微调。需要实现以下步骤： 1. 数据加载：使用一个适合因果语言建模的数据集（例如，wikitext, 或者自定义的文本数据） 2. tokenizer配置：使用与Yi-6B模型对应的tokenizer 3. 4-bit量化：使用bitsandbytes库进行4-bit量化，以降低显存占用 4. 梯度检查点（gradient checkpointing）：用于减少显存占用 5. 混合精度训练：使用fp16或bf16进行训练 6. LoRA配置：使用PEFT库添加LoRA适配器 我们将使用Hugging Face的transformers库，以及peft和bitsandbytes库。 步骤概述： 1. 安装必要的库 2. 加载模型（使用4-bit量化）和tokenizer 3. 准备数据集（使用datasets库加载并进行tokenize） 4. 配置LoRA参数 5. 配置训练参数（使用transformers.TrainingArguments） 6. 定义训练器（Trainer）并训练 注意：Yi-6B是一个较大的模型，需要足够的显存（即使使用4-bit，也需要大约10GB左右显存，具体取决于序列长度和batch size）。如果显存不足，可以调整batch size和序列长度，或者使用梯度累积。 代码实现如下：</think>### 使用LoRA微调Yi-6B模型的完整实现 以下是使用LoRA对Yi-6B模型进行因果语言建模微调的完整实现，包含数据加载、tokenizer配置、4-bit量化、梯度检查点和混合精度训练： ```python # 安装必要库 !pip install transformers accelerate bitsandbytes peft datasets import torch from transformers import ( AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig, TrainingArguments, Trainer ) from peft import LoraConfig, get_peft_model from datasets import load_dataset # 1. 配置4-bit量化 bnb_config = BitsAndBytesConfig( load_in_4bit=True, # 启用4-bit量化 bnb_4bit_quant_type="nf4", # 使用NormalFloat4量化类型 bnb_4bit_compute_dtype=torch.bfloat16, # 计算时使用bfloat16 bnb_4bit_use_double_quant=True # 启用双重量化减少内存 ) # 2. 加载模型和tokenizer model = AutoModelForCausalLM.from_pretrained( "01-ai/Yi-6B", quantization_config=bnb_config, device_map="auto", # 自动分配设备 trust_remote_code=True ) tokenizer = AutoTokenizer.from_pretrained("01-ai/Yi-6B") tokenizer.pad_token = tokenizer.eos_token # 设置pad token # 3. 配置LoRA适配器 lora_config = LoraConfig( r=8, # LoRA秩 lora_alpha=32, # 缩放因子 target_modules=["q_proj", "v_proj"], # 目标模块 lora_dropout=0.05, bias="none", task_type="CAUSAL_LM" ) model = get_peft_model(model, lora_config) model.print_trainable_parameters() # 打印可训练参数占比 # 4. 加载并预处理数据集 dataset = load_dataset("wikitext", "wikitext-2-raw-v1") def tokenize_function(examples): return tokenizer( examples["text"], truncation=True, max_length=512, # 最大序列长度 padding="max_length" ) tokenized_dataset = dataset.map( tokenize_function, batched=True, remove_columns=["text"] ) # 5. 配置训练参数 training_args = TrainingArguments( output_dir="./yi6b-lora-output", per_device_train_batch_size=4, # 根据显存调整 gradient_accumulation_steps=4, # 梯度累积 learning_rate=2e-4, num_train_epochs=3, fp16=True, # 混合精度训练 gradient_checkpointing=True, # 梯度检查点 logging_steps=50, optim="paged_adamw_8bit", # 分页优化器 report_to="none", save_strategy="epoch", evaluation_strategy="epoch", logging_dir="./logs" ) # 6. 创建Trainer并训练 trainer = Trainer( model=model, args=training_args, train_dataset=tokenized_dataset["train"], eval_dataset=tokenized_dataset["validation"] ) trainer.train() ``` #### 关键配置说明： 1. **4-bit量化**： - 使用`NF4`量化类型和双重量化技术，显著降低显存占用 - 计算公式：$ \text{显存节省} = \frac{32}{\text{4-bit}} = 8\times $ - 计算时使用`bfloat16`保持精度[^1] 2. **LoRA配置**： - 仅微调`q_proj`和`v_proj`层 - 可训练参数占比通常小于原模型的$1\%$ - 数学原理：$ W = W_0 + BA $，其中$ B \in \mathbb{R}^{d \times r}, A \in \mathbb{R}^{r \times k} $ 3. **梯度检查点**： - 用时间换空间，显存减少约$ \frac{2}{3} $ - 前向传播时不保存中间激活值，反向传播时重新计算 4. **混合精度训练**： - 使用`fp16`或`bf16`加速计算 - 梯度缩放防止下溢：$ \text{gradient} \times \text{scale} $ #### 注意事项： 1. 根据显存调整`per_device_train_batch_size`和`gradient_accumulation_steps` 2. 序列长度影响显存：$ \text{显存} \propto (\text{seq\_len})^2 $ 3. 使用`model.push_to_hub()`保存适配器权重 4. 对于自定义数据集，需预处理为`{'text': ...}`格式 此实现结合了QLoRA的核心优化技术[^1]，在保持模型性能的同时显著降低资源需求，适合在单卡消费级GPU上运行。