用LoRA进行高效微调

大型语言模型的低秩自适应 (LoRA) 用于解决微调大型语言模型 (LLM) 的挑战。GPT 和 Llama 等模型拥有数十亿个参数，通常对于特定任务或领域进行微调的成本过高。LoRA 保留了预训练的模型权重，并在每个模型块中加入了可训练层。这显著减少了需要微调的参数数量，并大大降低了 GPU 内存需求。LoRA 的主要优势在于，它大幅减少了可训练参数的数量——有时最多可减少 10,000 倍——从而大大降低了 GPU 资源需求。

1、LoRA 为何有效

预训练的 LLM 在适应新任务时具有较低的“固有维度”，这意味着数据可以通过低维空间有效地表示或近似，同时保留其大部分基本信息或结构。我们可以将适应任务的新权重矩阵分解为低维（较小）矩阵，而不会丢失大量重要信息。我们通过低秩近似实现这一点。

矩阵的秩是一个可以让你了解矩阵复杂度的值。矩阵的低秩近似旨在尽可能接近原始矩阵，但秩较低。低秩矩阵降低了计算复杂度，从而提高了矩阵乘法的效率。低秩分解是指通过推导矩阵 A 的低秩近似来有效近似矩阵 A 的过程。奇异值分解 (SVD) 是一种常用的低秩分解方法。

假设 W 表示给定神经网络层中的权重矩阵，假设 ΔW 是经过完全微调后 W 的权重更新。然后，我们可以将权重更新矩阵 ΔW 分解为两个较小的矩阵：ΔW = WA*WB，其中 WA 是 A × r 维矩阵，WB 是 r × B 维矩阵。在这里，我们保持原始权重 W 不变，只训练新矩阵 WA 和 WB。这总结了 LoRA 方法，如下图所示。

LoRA 的优势如下：

• 减少资源消耗。对深度学习模型进行微调通常需要大量计算资源，这可能既昂贵又耗时。LoRA 可在保持高性能的同时减少对资源的需求。
• 更快的迭代。LoRA 可实现快速迭代，从而更轻松地尝试不同的微调任务并快速调整模型。
• 改进迁移学习。LoRA 提高了迁移学习的有效性，因为带有 LoRA 适配器的模型可以用更少的数据进行微调。这在标记数据稀缺的情况下尤其有价值。
• 广泛适用。LoRA 用途广泛，可应用于自然语言处理、计算机视觉和语音识别等不同领域。
• 降低碳排放。通过减少计算要求，LoRA 有助于实现更环保、更可持续的深度学习方法。

2、使用 LoRA 技术训练神经网络

在此博客中，我们利用 CIFAR-10 数据集，使用几个epoch从头开始训练基本图像分类器。之后，我们进一步使用 LoRA 训练模型，说明将 LoRA 纳入训练过程的优势。

2.1 设置

此演示使用以下设置创建。有关全面的支持详细信息，请参阅 ROCm 文档。

硬件和操作系统：

• AMD Instinct GPU
• Ubuntu 22.04.3 LTS

软件：

• ROCm 5.7.0+
• Pytorch 2.0+

2.2 训练初始模型

导入软件包：

import torch
import torchvision
import torchvision.transforms as transforms

加载数据集并设置目标设备：

# 10 classes from CIFAR10 dataset
classes = ('airplane', 'automobile', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck')

# batch size
batch_size = 8

# image preprocessing
preprocessor = transforms.Compose(
    [transforms.ToTensor(),
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])

# training dataset
train_set = torchvision.datasets.CIFAR10(root='./dataset', train=True,
                                        download=True, transform=preprocessor)
train_loader = torch.utils.data.DataLoader(train_set, batch_size=batch_size,
                                        shuffle=True, num_workers=8)
# test dataset
test_set = torchvision.datasets.CIFAR10(root='./dataset', train=False,
                                    download=True, transform=preprocessor)
test_loader = torch.utils.data.DataLoader(test_set, batch_size=batch_size,
                                        shuffle=False, num_workers=8)

# Define the device
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

展示数据集的样本：

import matplotlib.pyplot as plt
import numpy as np

# helper function to display image
def image_display(images):
    # get the original image
    images = images * 0.5 + 0.5
    plt.imshow(np.transpose(images.numpy(), (1, 2, 0)))
    plt.axis('off')
    plt.show()

# get a batch of images
images, labels = next(iter(train_loader))
# display images
image_display(torchvision.utils.make_grid(images))
# show ground truth labels
print('Ground truth labels: ', ' '.join(f'{classes[labels[j]]}' for j in range(images.shape[0])))

输出显示：

Ground truth labels:  cat ship ship airplane frog frog automobile frog

创建一个用于图像分类的基本三层神经网络，注重简单性以说明 LoRA 效果：

import torch.nn as nn
import torch.nn.functional as F

class net(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(3*32*32, 4096)
        self.fc2 = nn.Linear(4096, 2048)
        self.fc3 = nn.Linear(2048, 10)

    def forward(self, x):
        x = torch.flatten(x, 1)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x

# move the model to device
classifier = net().to(device)

接下来训练模型。

我们使用交叉熵损失和 Adam 作为损失函数和优化器。

import torch.optim as optim

def train(train_loader, classifier, start_epoch = 0, epochs=1, device="cuda:0"):
    classifier = classifier.to(device)
    classifier.train()
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.Adam(classifier.parameters(), lr=0.001)
    
    for epoch in range(epochs):  # training loop

        loss_log = 0.0
        for i, data in enumerate(train_loader, 0):
            inputs, labels = data[0].to(device), data[1].to(device)
            # Resets the parameter gradients
            optimizer.zero_grad()
    
            outputs = classifier(inputs)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()
    
            # print loss after every 1000 mini-batches
            loss_log += loss.item()
            if i % 2000 == 1999:    
                print(f'[{start_epoch + epoch}, {i+1:5d}] loss: {loss_log / 2000:.3f}')
                loss_log = 0.0

现在开始训练：

import time

start_epoch = 0
epochs = 1
# warm up the gpu with one epoch
train(train_loader, classifier, start_epoch=start_epoch, epochs=epochs, device=device)

# run another epoch to record the time
start_epoch += epochs
epochs = 1
start = time.time()
train(train_loader, classifier, start_epoch=start_epoch, epochs=epochs, device=device)
torch.cuda.synchronize()
end = time.time()
train_time = (end - start)

print(f"One epoch takes {train_time:.3f} seconds")

输出如下：

    [0,  2000] loss: 1.987
    [0,  4000] loss: 1.906
    [0,  6000] loss: 1.843
    [1,  2000] loss: 1.807
    [1,  4000] loss: 1.802
    [1,  6000] loss: 1.782
    One epoch takes 31.896 seconds

一个 epoch 大约需要 31 秒。

保存模型：

model_path = './classifier_cira10.pth'
torch.save(classifier.state_dict(), model_path)

稍后我们将使用 LoRA 训练相同的模型，并检查训练一个 epoch 需要多长时间。

加载保存的模型并进行快速测试：

# Prepare the test data.
images, labels = next(iter(test_loader))
# display the test images
image_display(torchvision.utils.make_grid(images))
# show ground truth labels
print('Ground truth labels: ', ' '.join(f'{classes[labels[j]]}' for j in range(images.shape[0])))

# Load the saved model and have a test
model = net()
model.load_state_dict(torch.load(model_path))
model = model.to(device)
images = images.to(device)
outputs = model(images)
_, predicted = torch.max(outputs, 1)

print('Predicted: ', ' '.join(f'{classes[predicted[j]]}'
                            for j in range(images.shape[0])))

输出：

Ground truth labels:  cat ship ship airplane frog frog automobile frog
Predicted:  deer truck airplane ship deer frog automobile bird

我们观察到，仅对模型进行两个阶段的训练并不能产生令人满意的结果。让我们检查一下该模型在整个测试数据集上的表现。

def test(model, test_loader, device):
    model=model.to(device)
    model.eval()
    correct = 0
    total = 0
    with torch.no_grad():
        for data in test_loader:
            images, labels = data[0].to(device), data[1].to(device)
            # images = images.to(device)
            # labels = labels.to(device)
            # inference
            outputs = model(images)
            # get the best prediction
            _, predicted = torch.max(outputs.data, 1)
            
            total += labels.size(0)
            correct += (predicted == labels).sum().item()
    
    print(f'Accuracy of the given model on the {total} test images is {100 * correct // total} %')

test(model, test_loader, device)

输出：

    Accuracy of the given model on the 10000 test images is 32 %

这一结果表明，通过进一步训练，模型有很大的改进潜力。在以下部分中，我们将把 LoRA 应用于模型，并继续使用这种方法进行训练。

2.3 将 LoRA 应用于模型

定义用于将 LoRA 应用于模型的辅助函数：

class ParametrizationWithLoRA(nn.Module):
    def __init__(self, features_in, features_out, rank=1, alpha=1, device='cpu'):
        super().__init__()

        # Create A B and scale used in ∆W = BA x α/r
        self.lora_weights_A = nn.Parameter(torch.zeros((rank,features_out)).to(device))
        nn.init.normal_(self.lora_weights_A, mean=0, std=1)
        self.lora_weights_B = nn.Parameter(torch.zeros((features_in, rank)).to(device))
        self.scale = alpha / rank
        
        self.enabled = True

    def forward(self, original_weights):
        if self.enabled:
            return original_weights + torch.matmul(self.lora_weights_B, self.lora_weights_A).view(original_weights.shape) * self.scale
        else:
            return original_weights

def apply_parameterization_lora(layer, device, rank=1, alpha=1):
    """
    Apply loRA to a given layer
    """
    features_in, features_out = layer.weight.shape
    return ParametrizationWithLoRA(
        features_in, features_out, rank=rank, alpha=alpha, device=device
    )
    
def enable_lora(model, enabled=True):
    """
    enabled = True: incorporate the the lora parameters to the model
    enabled = False: the lora parameters have no impact on the model
    """
    for layer in [model.fc1, model.fc2, model.fc3]:
        layer.parametrizations["weight"][0].enabled = enabled

将 LoRA 应用于我们的模型：

import torch.nn.utils.parametrize as parametrize
parametrize.register_parametrization(model.fc1, "weight", apply_parameterization_lora(model.fc1, device))
parametrize.register_parametrization(model.fc2, "weight", apply_parameterization_lora(model.fc2, device))
parametrize.register_parametrization(model.fc3, "weight", apply_parameterization_lora(model.fc3, device))

现在，我们的模型参数由两部分组成：原始参数和 LoRA 引入的参数。由于我们尚未训练此更新后的模型，因此 LoRA 权重的初始化方式不会影响模型的准确性（请参阅“ParametrizationWithLoRA”）。因此，禁用或启用 LoRA 应该会导致模型的准确性相同。让我们来测试一下这个假设。

enable_lora(model, enabled=False)
test(model, test_loader, device)

输出：

    Accuracy of the network on the 10000 test images: 32 %

enable_lora(model, enabled=True)
test(model, test_loader, device)

输出：

    Accuracy of the network on the 10000 test images: 32 %

这正是我们所期望的。

现在让我们看看 LoRA 添加了多少参数。

total_lora_params = 0
total_original_params = 0
for index, layer in enumerate([model.fc1, model.fc2, model.fc3]):
    total_lora_params += layer.parametrizations["weight"][0].lora_weights_A.nelement() + layer.parametrizations["weight"][0].lora_weights_B.nelement()
    total_original_params += layer.weight.nelement() + layer.bias.nelement()

print(f'Number of parameters in the model with LoRA: {total_lora_params + total_original_params:,}')
print(f'Parameters added by LoRA: {total_lora_params:,}')
params_increment = (total_lora_params / total_original_params) * 100
print(f'Parameters increment: {params_increment:.3f}%')

输出：

    Number of parameters in the model with LoRA: 21,013,524
    Parameters added by LoRA: 15,370
    Parameters increment: 0.073%

LoRA 只为我们的模型添加了 0.073% 的参数。

接下来继续使用 LoRA 训练模型。

在继续训练模型之前，我们希望冻结模型的所有原始参数，如论文所述。通过这样做，我们只更新 LoRA 引入的权重，这是原始模型参数数量的 0.073%。

for name, param in model.named_parameters():
    if 'lora' not in name:
        param.requires_grad = False

继续应用 LoRA 训练模型。

# make sure the loRA is enabled 
enable_lora(model, enabled=True)

start_epoch += epochs
epochs = 1
# warm up the GPU with the new model (loRA enabled) one epoch for testing the training time
train(train_loader, model, start_epoch=start_epoch, epochs=epochs, device=device)

start = time.time()
# run another epoch to record the time
start_epoch += epochs
epochs = 1
import time
start = time.time()
train(train_loader, model, start_epoch=start_epoch, epochs=epochs, device=device)
torch.cuda.synchronize()
end = time.time()
train_time = (end - start)
print(f"One epoch takes {train_time} seconds")

输出：

    [2,  2000] loss: 1.643
    [2,  4000] loss: 1.606
    [2,  6000] loss: 1.601
    [3,  2000] loss: 1.568
    [3,  4000] loss: 1.560
    [3,  6000] loss: 1.585
    One epoch takes 16.622623205184937 seconds

你可能会注意到，现在只需大约 16 秒即可完成一个 epoch 的训练，这大约是训练原始模型所需时间（31 秒）的 53%。

损失的减少表明模型已经从更新 LoRA 引入的参数中学习到了知识。现在，如果我们在启用 LoRA 的情况下测试模型，准确率应该高于我们之前使用原始模型实现的准确率（32%）。如果我们禁用 LoRA，该模型应该能够达到与原始模型相同的准确率。让我们继续进行这些测试。

enable_lora(model, enabled=True)
test(model, test_loader, device)
enable_lora(model, enabled=False)
test(model, test_loader, device)

输出：

    Accuracy of the given model on the 10000 test images is 42 %
    Accuracy of the given model on the 10000 test images is 32 %

使用之前的图像再次测试更新后的模型。

# display the test images
image_display(torchvision.utils.make_grid(images.cpu()))
# show ground truth labels
print('Ground truth labels: ', ' '.join(f'{classes[labels[j]]}' for j in range(images.shape[0])))

# Load the saved model and have a test
enable_lora(model, enabled=True)
images = images.to(device)
outputs = model(images)
_, predicted = torch.max(outputs, 1)

print('Predicted: ', ' '.join(f'{classes[predicted[j]]}'
                            for j in range(images.shape[0])))

输出：

    Ground truth labels:  cat ship ship airplane frog frog automobile frog
    Predicted:  cat ship ship ship frog frog automobile frog

我们可以观察到，与步骤 6 中获得的结果相比，新模型的表现更好，表明参数确实学到了有意义的信息。

3、结束语

在这篇博文中，我们探索了 LoRA 算法，深入研究了它的原理以及在带有 ROCm 的 AMD GPU 上的实现。我们从头开始开发了一个基本网络和 LoRA 模块，以展示 LoRA 如何有效地减少可训练参数和训练时间。

原文链接：Using LoRA for efficient fine-tuning: Fundamental principles

BimAnt翻译整理，转载请标明出处