PyTorch TorchTune 教程：使用LoRA微调Llama2模型-优快云博客

本文链接：https://blog.youkuaiyun.com/gitblog_00250/article/details/148506183

PyTorch TorchTune 教程：使用LoRA微调Llama2模型

torchtune A Native-PyTorch Library for LLM Fine-tuning 项目地址: https://gitcode.com/gh_mirrors/to/torchtune

概述

本文将详细介绍如何在PyTorch TorchTune框架中使用LoRA（Low-Rank Adaptation）技术高效微调Llama2模型。LoRA是一种参数高效的微调方法，可以显著减少训练过程中的内存消耗，同时保持模型性能。

LoRA技术原理

什么是LoRA？

LoRA是一种基于适配器的参数高效微调技术，其核心思想是在神经网络的不同层添加可训练的低秩分解矩阵，同时冻结模型原有参数。在Transformer模型中，LoRA通常应用于自注意力机制中的线性投影层。

LoRA工作原理

LoRA通过低秩近似替代权重更新矩阵。具体来说，对于一个nn.Linear(in_dim, out_dim)层，LoRA训练两个矩阵A和B：

A矩阵将输入投影到更小的维度（通常为4或8）
B矩阵将低维表示投影回原始输出维度

数学上，LoRA将参数数量从in_dim*out_dim减少到r*(in_dim+out_dim)，其中r是远小于输入输出维度的秩。

内存优势

LoRA的主要内存优势体现在：

梯度存储量大幅减少
优化器状态内存占用降低（特别是使用AdamW等含动量的优化器时）

在TorchTune中实现LoRA

基本实现

以下是LoRA线性层的简化实现：

import torch
from torch import nn

class LoRALinear(nn.Module):
    def __init__(self, in_dim, out_dim, rank, alpha, dropout):
        super().__init__()
        self.linear = nn.Linear(in_dim, out_dim, bias=False)
        self.lora_a = nn.Linear(in_dim, rank, bias=False)
        self.lora_b = nn.Linear(rank, out_dim, bias=False)
        self.rank = rank
        self.alpha = alpha
        self.dropout = nn.Dropout(p=dropout)
        self.linear.weight.requires_grad = False

    def forward(self, x):
        frozen_out = self.linear(x)
        lora_out = self.lora_b(self.lora_a(self.dropout(x)))
        return frozen_out + (self.alpha / self.rank) * lora_out

应用于Llama2模型

在TorchTune中，可以轻松地为Llama2模型添加LoRA层：

from torchtune.models.llama2 import lora_llama2_7b

# 默认配置下创建LoRA Llama2模型
lora_model = lora_llama2_7b(lora_attn_modules=["q_proj", "v_proj"])

参数设置

加载基础模型权重后，需要设置可训练参数：

from torchtune.modules.peft.peft_utils import get_adapter_params, set_trainable_params

lora_params = get_adapter_params(lora_model)
set_trainable_params(lora_model, lora_params)

LoRA微调实践

基本配置

TorchTune提供了LoRA微调配方，典型配置如下：

# 模型参数
model:
  _component_: lora_llama2_7b
  lora_attn_modules: ['q_proj', 'v_proj']
  lora_rank: 8
  lora_alpha: 16

运行微调

使用2个GPU运行分布式微调：

tune run --nnodes 1 --nproc_per_node 2 lora_finetune_distributed --config llama2/7B_lora

单设备微调

对于资源有限的环境，可以在单设备上运行：

tune run lora_finetune_single_device --config llama2/7B_lora_single_device

实验与调优

配置实验

我们可以调整以下参数进行实验：

应用LoRA的层（自注意力、MLP、输出投影）
LoRA秩大小（通常与alpha值同步调整）

实验结果示例

| LoRA层配置 | 秩 | Alpha | 峰值内存 | 准确率 | |-----------------|----|-------|---------|-------| | 仅Q和V | 8 | 16 | 15.57GB | 0.475 | | 所有层 | 8 | 16 | 15.87GB | 0.508 | | 仅Q和V | 64 | 128 | 15.86GB | 0.504 | | 所有层 | 64 | 128 | 17.04GB | 0.514 |