近期大模型层出不穷，大家对于大模型的微调也在跃跃欲试，像Lijia的BELLE，斯坦福的Alpaca[1], 清华的ChatGLM[2]，中文的Chinese-Vicuna[3]，让我这样的普通玩家也能训练自己的微调模型。

在微调和推理的时候仍然需要加速，有哪些方法可以加速微调呢？

Part1 LoRA

低秩矩阵分解 LoRA[4]原理：冻结预训练模型权重，并将可训练的秩分解矩阵注入到Transformer层的每个权重中，大大减少了下游任务的可训练参数数量。LoRA 开源代码[5]见文末。

原理图：公式：

结合原理图和公式，我们可以很容易明白LoRA了：

左侧是预训练模型的权重，输入输出维度都是d，在训练期间被冻结，不接受梯度更新。

右侧，对A使用随机的高斯初始化，B在训练开始时为零，r是秩，会对△Wx做缩放 α/r。

HuggingFace的包peft[6]对LoRA做了封装支持，几步即可使用：

from peft import get_peft_model, LoraConfig, TaskType

peft_config = LoraConfig(
    task_type=TaskType.CAUSAL_LM, 
    inference_mode=False, 
    r=8, 
    lora_alpha=32, 
    lora_dropout=0.1,
    target_modules=['query_key_value']
)

model = "加载的模型"
model = get_peft_model(model, peft_config)
# 打印参数情况
model.print_trainable_parameters()
接下来和正常训练模型一样

论文中提到了LoRA的诸多优点：

Part2 Accelerate 和 deepspeed

Accelerate[7]库提供了简单的 API，使我们可以在任何类型的单节点或分布式节点(单CPU、单GPU、多GPU 和 TPU)上运行，也可以在有或没有混合精度(fp16)的情况下运行。

这里是我用Accelerator和DeepSpeedPlugin做个示例：

需要提前知道梯度累积步骤 gradient_accumulation_steps 和 梯度累积计算

from accelerate import Accelerator, DeepSpeedPlugin
import tqdm

model = ...

deepspeed_plugin = DeepSpeedPlugin(
    zero_stage=2, 
    gradient_accumulation_steps=2)

accelerator = Accelerator(
    mixed_precision='fp16', 
    gradient_accumulation_steps=2, 
    deepspeed_plugin=deepspeed_plugin)

device = accelerator.device
... ...
optimizer = ...
lr_scheduler = ...

model, optimizer, train_dataloader = accelerator.prepare(model, optimizer, train_dataloader)

for epoch in range(epochs):
    total_loss = 0
    for step, batch in enumerate(t:=tqdm.tqdm(train_dataloader)):
        with accelerator.accumulate(model):
            outputs = model(**batch)
            loss_detach = outputs.loss.detach().cpu().float()
            t.set_description(f"loss: {loss_detach}")
            total_loss += loss_detach
            loss = outputs.loss
            # 不再是 loss.backward()
            accelerator.backward(loss)
            optimizer.step()
            lr_scheduler.step()
            optimizer.zero_grad()
    # 每个epoch 保存
    accelerator.wait_for_everyone()
    if accelerator.is_main_process:
        accelerator.save(model.state_dict(accelerator.unwrap_model(model), '/saved/model.pt')

# 其他参考保存方法
accelerator.wait_for_everyone()
unwrapped_model = accelerator.unwrap_model(model)
unwrapped_model.save_pretrained(save_dir, 
                                save_function=accelerator.save, 
                                state_dict=accelerator.get_state_dict(model))

Part3 Autocast 自动混合精度

autocast是在GPU上训练时一种用于降低显存消耗的技术。原理是用更短的总位数来保存浮点数，能够有效将显存消耗降低，从而设置更大的batch来加速训练。但会造成精度的损失，导致收敛效果也会变差。

PyTorch的AMP有2种精度是torch.FloatTensor和torch.HalfTensor。

使用方法：

from torch.cuda.amp import autocast as autocast, GradScaler

dataloader = ...
model = model.cuda()
optimizer = ...
scheduler = ...
# scaler的大小在每次迭代中动态估计，为了尽可能减少梯度，scaler应该更大；
# 但太大，半精度浮点型又容易 变成inf或NaN.
# 动态估计原理就是在不出现if或NaN梯度的情况下，尽可能的增大scaler值。 
scaler = GradScaler()

for epoch in range(epochs):
    for batch_idx, (data, targets) in enumerate(train_dataloader):
        optimizer.zero_grad()
        data = data.cuda(0)
        with autocast(dtype=torch.bfloat16): # 自动混精度
            logits = model(data)
            loss = loss(logits, targets)
        # 反向传播梯度放大
        scaler.scale(loss).backward()
        # 首先 把梯度值unscale回来, 优化器中的值也需要放缩
        # 如果梯度值不是inf或NaN, 则调用optimizer.step()来更新权重，否则，忽略step调用，从而保证权重不更新。
        scaler.step(optimizer)
        # 看是否要增大scaler, 更新scaler
        scaler.update()

Part4 单机多GPU、多机多卡

如果条件允许的话，可以使用单机多卡和多机多卡分布式训练。

那么：

• 模型怎么同步参数与梯度？
• 数据怎么划分到多个GPU中？

pytorch框架给我们封装了对应的接口函数：

import torch.distributed as dist
from torch.utils.data.distributed import DistributedSampler
from torch.nn.parallel import DistributedDataParallel

PyTorch提供的torchrun命令以及一些API封装了多进程的实现。 我们只要在普通【单进程程序前后】加入： 开头 setup