模型训练 llama模型训练与微调实战

该项目实现4位量化+LoRA轻量级微调Llama-2模型,一般适用于小资源环境下的大模型适配任务。(需要peft,datasets和transformerku)

一、代码框架

核心流程:

  1. 加载4位量化的Llama-2基础模型
  2. 配置LoRA参数并注入模型
  3. 加载SQuAD数据集并预处理
  4. 启动训练与评估

二、代码实操

1. 第一步是导入依赖库喵

# 数据集加载工具
from datasets import load_dataset
import torch
# LoRA配置与模型封装工具
from peft import LoraConfig, get_peft_model
# Transformers核心工具:模型、分词器、训练参数、训练器
from transformers import (
    AutoModelForCausalLM,    # 因果语言模型(Llama-2是自回归模型)
    AutoTokenizer,           # 自动分词器
    BitsAndBytesConfig,      # 4/8位量化配置
    TrainingArguments,       # 训练参数配置
    Trainer                  # 训练器
)

2. 参数定义

output_dir写你自己的模型保存路径(这里用相对路径./outputs)
别忘了喵

# 文本最大长度(需根据模型支持的上下文窗口调整,Llama-2-7B默认是4096)
max_length = 128

# ---------------------- 模型加载参数 ----------------------
load_in_4bit = True  # 启用4位量化(减少显存占用)

# ---------------------- LoRA微调参数 ----------------------
lora_alpha = 16  # LoRA的缩放因子(控制LoRA参数的权重占比)
lora_dropout = 0.1  # LoRA层的dropout率(防止过拟合)
lora_r = 16  # LoRA的低秩矩阵维度(r越小,参数量越少)
lora_bias = "all"  # LoRA是否训练偏置项("all"/"none")
model_type = "llama"  # 模型类型(可选"falcon"/"llama")
# LoRA要注入的目标模块(需匹配Llama-2的模型结构)
lora_target_modules = [
    "q_proj",          # 查询投影层
    "k_proj",          # 键投影层
    "v_proj",          # 值投影层
    "dense",           # 全连接层
    "dense_h_to_4h",   # 前馈层升维层
]

# ---------------------- 训练参数 ----------------------
output_dir = "./outputs"  # 模型保存目录
learning_rate = 0.002     # 学习率
weight_decay = 0.0005     # 权重衰减(防止过拟合)
per_device_train_batch_size = 1  # 单设备训练批次大小
per_device_eval_batch_size = 1   # 单设备评估批次大小
gradient_accumulation_steps = 1  # 梯度累积步数(模拟大批次)
warmup_steps = 10               # 学习率预热步数
save_steps = 10                 # 模型保存步数间隔
logging_steps = 10              # 日志打印步数间隔

3. 加载4位量化模型(这里用llama2)

def load_4bit_model():
    # 配置4位量化参数
    bnb_config = BitsAndBytesConfig(
        load_in_4bit=load_in_4bit,  # 启用4位加载
        bnb_4bit_use_double_quant=True,  # 双量化(进一步压缩)
        bnb_4bit_quant_type="nf4",  # 量化类型(NF4适配大模型)
        bnb_4bit_compute_dtype=torch.bfloat16  # 计算精度
    )

    # 加载Llama-2-7B基础模型(需提前获取Meta的授权)
    if model_type == "llama":
        model = AutoModelForCausalLM.from_pretrained(
            "meta-llama/Llama-2-7b-hf",# 加载llama2预训练模型
            quantization_config=bnb_config,  # 应用4位量化
            device_map="auto",  # 自动分配设备(CPU/GPU)
            trust_remote_code=True  # 允许加载远程代码(Llama-2需要)
        )
	# 当模型类型指定为 "falcon" 时执行该代码块
	elif model_type == "falcon":
	    # 加载 Falcon 系列因果语言模型(CausalLM)
	    # AutoModelForCausalLM 是 Hugging Face Transformers 的自动模型加载类会根据模型名称自动识别并加载对应的因果语言模型架构
	    model = AutoModelForCausalLM.from_pretrained(
	        "tiiuae/falcon-7b", 
	        quantization_config=bnb_config,  # 量化配置参数:使用 bitsandbytes(bnb)库进行模型量化
	        device_map="auto",  # 设备分配策略:自动将模型层分配到可用设备(CPU/GPU)
	        trust_remote_code=True) # 允许执行模型仓库中的自定义代码(.py/.json 等)
	                               # 因为 Falcon 模型的架构实现不在 Transformers 官方内置中,需要从模型仓库加载自定义代码才能正确初始化模型

	return model
# 函数返回加载完成的模型实例

4. 加载分词器

(这里也有类似的条件分支判断,这里的部分注释就不写了喵)

def load_tokenizer():
    if model_type == "llama":
        tokenizer = AutoTokenizer.from_pretrained(
            "meta-llama/Llama-2-7b-hf",
            trust_remote_code=True # 允许执行模型仓库中的自定义代码
        )
    elif model_type == "falcon":
        tokenizer = AutoTokenizer.from_pretrained(
            "tiiuae/falcon-7b",
            trust_remote_code=True
        )
    # 设置填充token(Llama-2默认没有pad_token,用eos_token代替)
    tokenizer.pad_token = tokenizer.eos_token
    return tokenizer
    # 返回分词器实例

5. 数据预处理(SQuAD数据集)

# 加载SQuAD问答数据集
dataset = load_dataset("squad")

def map_function(example):
    """
    数据预处理函数:将SQuAD的question和answer格式化为对话形式
    example: 数据集的单条样本(包含question、answers等字段)
    """
    # 构造对话格式(匹配Llama-2的指令微调模板)即转换数据集格式
    question = f"##### Human: {example['question'].strip()}"
    output = f"##### Assistant: {example['answers']['text'][0].strip()}"

    # 对question编码(不截断、不填充)
    question_encoded = tokenizer(question)
    # 对output编码(需保证question+output总长度不超过max_length)
    output_encoded = tokenizer(
        output,
        max_length=max_length - len(question_encoded["input_ids"]),
        truncation=True,  # 超过长度则截断
        padding="max_length"  # 不足长度则填充
    )

    # ---------------------- 构造训练用的input_ids ----------------------
    # 将question和output的input_ids拼接(模型输入)
    input_ids = question_encoded["input_ids"] + output_encoded["input_ids"]

    # ---------------------- 构造训练用的labels ----------------------
    # labels中,question部分标记为-100(PyTorch交叉熵损失会忽略-100)
    # output部分保留原token_id(模型需要学习预测的部分)
    labels = [
        -100 if i < len(question_encoded["input_ids"]) 
        else output_encoded["input_ids"][i - len(question_encoded["input_ids"])]
        for i in range(len(input_ids))
    ]

    # ---------------------- 构造attention_mask ----------------------
    # 注意力掩码:1表示有效token,0表示填充token
    attention_mask = question_encoded["attention_mask"] + output_encoded["attention_mask"]

    return {
        "input_ids": input_ids,
        "labels": labels,
        "attention_mask": attention_mask
    }

# 对训练集和验证集应用预处理
data_train = dataset["train"].map(map_function)
data_test = dataset["validation"].map(map_function)
 

6. 注入LoRA权重

def apply_lora(model):
    # 配置LoRA参数
    peft_config = LoraConfig(
        lora_alpha=lora_alpha,  # LoRA的缩放因子(与r相关),通常设置为 2*r
                                # 作用:控制LoRA适配器对模型输出的影响程度,alpha越大影响越显著
        lora_dropout=lora_dropout,  # LoRA层的dropout概率(如0.05/0.1)
                                    # 作用:防止过拟合,随机丢弃部分LoRA参数的梯度更新
        r=lora_r,  # LoRA的低秩维度(核心参数),通常取值范围8-64(如8/16/32)
                   # 原理:将高维权重矩阵分解为两个低维矩阵(A: d*r,B: r*d),训练量从O(d²)降至O(dr)
                   # r越小,训练参数越少、速度越快,但微调能力越弱;需根据任务调整
        bias=lora_bias,  # LoRA的偏置项训练策略(可选值:"none"/"all"/"lora_only")
        task_type="CAUSAL_LM",  # 任务类型:因果语言模型(与基础模型类型一致)
        inference_mode=False,   # 训练模式开关:False=训练模式(启用梯度更新),True=推理模式(冻结参数)
        target_modules=lora_target_modules  # 注入LoRA适配器的目标模块名称列表
                                          
    )
    # 将LoRA注入基础模型
    peft_model = get_peft_model(model, peft_config)
    # 打印可训练参数占比(LoRA仅训练少量参数)
    peft_model.print_trainable_parameters()
    return peft_model

7. 启动模型训练

def train_model(model, tokenizer):
    # 配置训练核心参数
    training_args = TrainingArguments(
        output_dir=output_dir,  # 模型及训练日志保存目录
        evaluation_strategy="epoch",  # 按每个epoch进行一次验证
        optim="paged_adamw_8bit",  # 8位优化器,节省显存
        learning_rate=learning_rate,  # 学习率
        weight_decay=weight_decay,  # 权重衰减,防止过拟合
        per_device_train_batch_size=per_device_train_batch_size,  # 单设备训练批次大小
        per_device_eval_batch_size=per_device_eval_batch_size,  # 单设备验证批次大小
        gradient_accumulation_steps=gradient_accumulation_steps,  # 梯度累积步数
        do_train=True,  # 开启训练模式
        warmup_steps=warmup_steps,  # 学习率热身步数
        save_steps=save_steps,  # 模型保存间隔步数
        logging_steps=logging_steps  # 训练日志打印间隔步数
    )

    # 初始化训练器(整合模型、参数、数据、分词器)
    trainer = Trainer(
        model=model,  # 待训练的模型(含LoRA的PEFT模型)
        args=training_args,  # 训练参数配置
        train_dataset=data_train,  # 训练数据集
        eval_dataset=data_test,    # 验证数据集
        tokenizer=tokenizer        # 分词器(用于批量数据填充等预处理)
    )

    # 启动模型训练
    trainer.train()
    # 保存训练完成的最终模型
    trainer.save_model(output_dir + "/final_model")# 模型文件最终会保存到根目录的/outputs/final_model

8. 主函数(执行)

最后就是主函数了喵,这里是整个项目的主要运行流程

if __name__ == "__main__":
    # 1. 加载4位量化模型
    model = load_4bit_model()
    # 2. 加载分词器
    tokenizer = load_tokenizer()
    # 3. 注入LoRA
    peft_model = apply_lora(model)
    # 4. 启动训练
    train_model(peft_model, tokenizer)

三、训练输出说明

  1. 命令行输出(对应原文图6-9)

训练过程中会打印:

  • 数据加载进度
  • 可训练参数占比(如 trainable params: 29,987,360 || all params: 6,869,233,408 || trainable%: 0.436 )
  • 每步的训练损失、速度等
  1. wandb监控(对应原文图6-10/6-11)

通过wandb可查看:

  • 训练/验证损失曲线
  • 学习率变化
  • 显存占用等资源指标
  1. 模型保存目录(对应原文6.8.8)

训练完成后, outputs 目录下会生成多个检查点(如 checkpoint-10 ),每个检查点包含:

adapter_config.json  # LoRA配置
adapter_model.bin    # LoRA权重(仅几MB)
tokenizer.model      # 分词器文件

四、模型合并(LoRA+基础模型)

别忘了喵,微调得到的是LoRA增量模型,推理前需与基础模型合并:

from peft import PeftModel

# 加载基础模型(同训练时的4位量化模型)
base_model = load_4bit_model()
# 加载LoRA增量模型
peft_model = PeftModel.from_pretrained(base_model, "./outputs/final_model")
# 合并模型(仅需执行一次)
merged_model = peft_model.merge_and_unload()
# 保存合并后的模型
merged_model.save_pretrained("./merged_model")

我会把完整代码文件放到资源里面的 喵喵

五、模型加载

LoRA微调模型加载代码示例

# 导入依赖库
from peft import PeftModel  # LoRA模型工具
from transformers import (
    AutoModelForCausalLM,  # 因果语言模型(如Llama-2)
    AutoTokenizer,         # 自动分词器
    BitsAndBytesConfig     # 4/8位量化配置
)

# ---------------------- 加载基础模型(4位量化) ----------------------
# 配置4位量化参数(减少显存占用)
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,               # 启用4位加载
    bnb_4bit_use_double_quant=True,  # 双量化(进一步压缩)
    bnb_4bit_quant_type="nf4",       # 量化类型(适配大模型)
    bnb_4bit_compute_dtype=torch.bfloat16  # 计算精度
)

# 加载Llama-2基础模型(需提前获取Meta授权)
base_model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-2-7b-hf",
    quantization_config=bnb_config,  # 应用4位量化
    device_map="auto",               # 自动分配设备(CPU/GPU)
    trust_remote_code=True           # 允许加载远程代码
)

# ---------------------- 加载LoRA增量模型 ----------------------
# 从本地加载LoRA微调后的权重(仅几MB)
peft_model = PeftModel.from_pretrained(
    base_model, 
    "./outputs/final_lora_model"  # LoRA模型保存路径
)

# ---------------------- 合并模型(推理前需执行) ----------------------
# 将LoRA权重与基础模型合并(仅需执行一次)
merged_model = peft_model.merge_and_unload()

# ---------------------- 推理示例 ----------------------
tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-hf")
tokenizer.pad_token = tokenizer.eos_token  # 补充pad_token

# 输入文本(以羊驼介绍为例)
instruction = "请介绍羊驼(Vicugna pacos)"
inputs = tokenizer(
    instruction, 
    return_tensors="pt", 
    truncation=True, 
    max_length=128
).to("cuda")

# 生成输出
outputs = merged_model.generate(
    **inputs,
    max_new_tokens=100,  # 生成最大长度
    temperature=0.7      # 随机性
)
print(tokenizer.decode(outputs[0], skip_special_tokens=True))

六、加速技术和工具

1. DeepSpeed(微软)

核心定位:大模型分布式训练优化库,通过多种技术降低显存占用、提升训练速度。
关键特性:

  • ZeRO(核心功能):按GPU分片模型参数、梯度、优化器状态,单卡仅存部分数据,解决大模型显存不足问题;
  • 模型并行化:将模型拆分到多个GPU运行,支持超大规模模型(超过单卡显存);
  • 激活值检查点:反向传播时仅保存部分激活值,大幅减少显存占用(需额外计算);
  • 混合精度训练:同时用float32(精度)和float16(速度)运算,平衡效率与精度;
  • Pipeline并行化:模型不同层分布在不同GPU并行计算,提升硬件利用率;
  • 1-bit Adam:压缩通信带宽,加速分布式训练。

2. FairScale(Meta)

核心定位:大模型分布式训练效率优化库,提供轻量化并行训练工具。
关键特性:

  • 模型并行化:支持 ShardedDDP 等工具,实现模型在多GPU上的分片并行;
  • 优化器状态分片(OSS):拆分优化器状态到多卡,减少单卡显存占用;
  • 梯度累积:多步累积梯度后一次性更新,支持小显存训练大模型;
  • ZeRO实现:复用DeepSpeed的ZeRO思想,减少训练冗余数据;
  • 通信优化:优化分布式训练中的通信操作,提升训练效率。

3. FSDP(PyTorch)

核心定位:PyTorch原生完全分片数据并行工具,专为大模型训练设计。
关键特性:

  • 参数全分片:模型参数、梯度、优化器状态全部分片到多GPU,单卡显存占用极低;
  • 自动通信管理:框架自动处理GPU间参数同步,无需手动编写通信逻辑;
  • 弹性扩展:支持动态调整GPU数量,适配不同硬件资源;
  • 兼容性强:无缝对接PyTorch现有训练代码,仅需少量修改。

4. GPTQ

核心定位:大模型量化工具,将模型权重压缩为4/8位,大幅降低推理显存占用。
关键特性:

  • 低精度量化:在几乎不损失精度的前提下,将模型权重量化为4bit;
  • 推理加速:量化后模型体积更小,推理速度更快,适配普通GPU/CPU;
  • 易用性高:HuggingFace集成相关API,可直接在预训练模型上执行量化。

5、工具对比与适用场景

工具 核心优势 适用场景
DeepSpeed 功能最全、显存优化最强 DeepSpeed 功能最全、显存优化最强 超大规模模型(10B+参数)训练
FairScale 轻量、易集成 中等规模模型分布式训练
FSDP PyTorch原生、兼容性好 PyTorch生态下的大模型训练
GPTQ 量化精度高、推理速度快 大模型低资源环境推理

七、超长上下文技术

大模型默认上下文窗口有限(如Llama-2是4096),超长上下文技术可扩展到10万+ tokens,核心方法包括:

1. LongLoRA

通过LoRA微调扩展上下文窗口,无需全量训练模型:

  • 原理:在注意力层注入LoRA权重,调整位置编码适应更长序列;
  • 优势:训练成本低(仅需微调少量参数),可将Llama-2扩展到10万tokens;
  • 代码示例(核心配置):
from peft import LoraConfig
lora_config = LoraConfig(
    r=16,
    lora_alpha=32,
    target_modules=["q_proj", "k_proj"],  # 仅微调注意力层
    task_type="CAUSAL_LM"
)

2. StreamingLLM

通过注意力池化机制实现超长上下文推理:

  • 原理:仅保留最近的注意力标记(KV缓存),自动丢弃早期无关信息;
  • 优势:无需微调模型,直接扩展上下文窗口到百万tokens;
  • 特点:推理时动态调整KV缓存大小,平衡显存与性能。
Logo

有“AI”的1024 = 2048,欢迎大家加入2048 AI社区

更多推荐