大模型训练之计算量和内存优化篇------混合精度/量化/FSDP&cpu offload/Gradient Checkpointing

zhurui_xiaozhuzaizai

已于 2023-09-20 16:38:02 修改

阅读量2.8k

点赞数 4

分类专栏：自然语言处理文章标签：性能优化

于 2023-05-24 19:59:40 首次发布

自然语言处理专栏收录该内容

78 篇文章 26 订阅

订阅专栏

1 混合精度

主要思想是使用较低精度的浮点数（如FP16）来表示神经网络中的权重和激活值，从而减少内存使用和计算开销，进而加速训练过程。

精度数值范围

FP16：半精度浮点数，使用16位二进制数表示，其中1位表示符号位，5位表示指数位，10位表示尾数位，能够表示的数值范围为±2¹⁵。
FP32：单精度浮点数，使用32位二进制数表示，其中1位表示符号位，8位表示指数位，23位表示尾数位，能够表示的数值范围为±3.4×10³⁸。
FP64：双精度浮点数，使用64位二进制数表示，其中1位表示符号位，11位表示指数位，52位表示尾数位，能够表示的数值范围为±1.8×10³⁰⁸。
INT8：8位整数，能够表示的数值范围为-128到127。
INT4：4位整数，能够表示的数值范围为-8到7。

流程

混合精度训练的流程如下：

将FP32的权重转换为FP16格式，然后进行前向计算，得到FP32的损失（loss）。
使用FP16计算梯度。
将梯度转换为FP32格式，并将其更新到权重上。

由于FP16精度较低，可能会导致精度损失，因此在混合精度训练中需要进行一些技巧来保持模型的准确性。例如，可以使用梯度缩放（GradScaler）来控制梯度的大小，以避免梯度下降过快而影响模型的准确性。

FP32保存权重的原因

梯度的更新值太小，FP16直接变为了0
FP16表示权重，梯度的计算结果可能变成0
用FP16保存权重会造成80%的精度损失

使用 Huggingface Transformers：在 TrainingArguments 里声明 fp16=True
https://huggingface.co/docs/transformers/perf_train_gpu_one#fp16-training
CUDA Automatic Mixed Precision examples
https://pytorch.org/docs/stable/notes/amp_examples.html

使用apex
apex.parallel.DistributedDataParallel 与apex.amp 已被以下两个API取代
torch.nn.parallel.DistributedDataParallel和torch.amp

model = Net().cuda()
optimizer = optim.SGD(model.parameters(), ...)
# 创建scaler
scaler = GradScaler()

for epoch in epochs:
    for inputs, target in data:
        optimizer.zero_grad()
        with autocast(device_type='cuda', dtype=torch.float16):  # 自动进行转化
            output = model(inputs)
            loss = loss_fn(output, target)
        scaler.scale(loss).backward()  # 应用于损失函数
        scaler.step(optimizer)  # 应用于优化函数
        scaler.update()

2 量化

量化是一种通过减少数字表示的位数来减小模型存储量和计算量的方法。在深度学习中，通常使用32位浮点数来表示权重和激活值。但是，这种精度可能会导致计算和存储的开销非常高。因此，量化使用更短的整数表示权重和激活值，从而减少内存和计算开销。

在量化过程中，可以使用两种方法：动态量化和静态量化。
动态量化在运行时收集数据，并根据数据动态地量化模型。
静态量化在训练过程中对模型进行量化，并在推理时应用量化。
量化会导致模型准确度下降，因为更低的精度可能会导致舍入误差。因此，在量化期间，需要进行一些技巧来保持模型的准确程度，例如：对权重进行缩放或使用动态范围量化。

Huggingface 在这篇文章中用动图解释了 quantization 的实现：
https://huggingface.co/blog/hf-bitsandbytes-integration
借助 Huggingface PEFT，使用 int8 训练 opt-6.5B 的完整流程：
https://github.com/huggingface/peft/blob/main/examples/int8_training/Finetune_opt_bnb_peft.ipynb

量化推理

使用load_in_8bit方法可以实现模型的量化。该方法可以将模型权重和激活值量化为8位整数，从而减少内存和计算开销。具体实现方法如下：

需要注意的是，使用load_in_8bit方法量化模型可能会导致模型准确度下降。因此，在量化模型之前，需要对模型进行测试，确保准确度可以接受。另外，不是所有的模型都可以被量化，只有支持动态量化的模型才可以使用该方法进行量化。

import torch
from transformers import AutoModel

# 加载模型
model = AutoModel.from_pretrained('bert-base-uncased',load_in_8bit=True)

量化训练
在深度学习中，量化是一种通过减少数字表示的位数来减小模型存储量和计算量的方法。在使用混合精度训练时，可以将模型权重和梯度从FP32转换为FP16，以节省内存和加速训练。同样的思路，量化训练可以将激活值转换为更短的整数，从而减少内存和计算开销。

PyTorch中提供了一些量化训练的工具和API，例如QAT（量化感知训练），使用动态范围量化等。其中，使用Adam8bit进行量化训练是一种方法。

下面是使用Adam8bit、trainer进行量化训练的示例代码：

import torch
from torch.optim import Adam
from torch.optim.lr_scheduler import LambdaLR
from torch.utils.data import DataLoader
from torchvision.transforms import Compose, ToTensor, Normalize
from torch.cuda.amp import GradScaler, autocast
from torch.quantization import QuantWrapper, Adam8Bit

# 初始化模型
model = Net().cuda()

# 初始化Adam优化器
optimizer = Adam(model.parameters(), lr=0.001)

# 初始化Adam8Bit
optimizer = Adam8Bit(optimizer)

# 初始化学习率调度器
lr_scheduler = LambdaLR(optimizer, lr_lambda=lambda epoch: 0.1 ** (epoch // 20))

# 初始化GradScaler
scaler = GradScaler()

# 开始训练
for epoch in range(10):
    for i, (inputs, labels) in enumerate(train_loader):
        optimizer.zero_grad()
        with autocast():
            outputs = model(inputs.cuda())
            loss = torch.nn.functional.cross_entropy(outputs, labels.cuda())
        scaler.scale(loss).backward()
        scaler.step(optimizer)
        scaler.update()
        lr_scheduler.step()
        if i % 100 == 0:
            print('Epoch:{}, Iteration:{}, Loss:{:.4f}'.format(epoch, i, loss.item()))

# 量化模型
quantized_model = QuantWrapper(model).to(torch.device('cuda'))

其中torch.quantization.default_qconfig是PyTorch中提供的一个默认的量化配置，包含了一些默认的量化参数。在一些简单的量化任务中，可以使用这个默认配置，而不需要自己手动指定每个参数的值。其中：

activation_8bit: 激活值的8位量化配置。默认值为torch.quantization.QConfig(activation=torch.quantization.default_observer.with_args(dtype=torch.qint8), weight=torch.quantization.default_per_channel_qconfig)。

weight_8bit: 权重的8位量化配置。默认值为torch.quantization.default_per_channel_qconfig.with_args(dtype=torch.qint8)。

这些默认值可以在使用**quantize()**函数时进行自定义，以满足特定的量化需求。

import torch
import transformers

# 加载已经训练好的BERT模型
model = transformers.BertForSequenceClassification.from_pretrained('bert-base-uncased')

# 定义训练参数
training_args = transformers.TrainingArguments(
    output_dir='./results',
    evaluation_strategy='steps',
    eval_steps=1000,
    save_total_limit=3,
    learning_rate=1e-5,
    num_train_epochs=3,
    per_device_train_batch_size=8,
    per_device_eval_batch_size=8,
    warmup_steps=100,
    weight_decay=0.01,
    logging_dir='./logs',
    logging_steps=1000,
    load_best_model_at_end=True,
    metric_for_best_model='eval_loss',
    greater_is_better=False,
    quantization_config={
        'activations': torch.quantization.default_qconfig,
        'weights': torch.quantization.default_qconfig
    }
)

# 定义Trainer并进行量化训练
trainer = transformers.Trainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset,
    eval_dataset=val_dataset
)
trainer.quantize()
trainer.train()

3 CPU offload

在深度学习训练过程中，GPU需要大量的数据进行计算，但是如果数据没有及时传输到GPU就会导致GPU处于等待状态，浪费GPU的计算能力。因此，CPU卸载技术就应运而生，通过让CPU负责将数据传输到GPU，可以让GPU专心计算，提高训练速度。

卸载分为多种类型，包括数据卸载、模型卸载、梯度卸载等。数据卸载是指将数据卸载到存储设备（例如硬盘）中，用的时候传输到GPU中，其他的类似。

CPU
深度学习的训练需要用到GPU，因为它的计算速度比较快，但是如果GPU没有及时得到数据，就会处于等待状态，浪费了它的计算能力。CPU卸载技术就是让CPU来负责把数据传输到GPU，这样GPU就可以专心计算，提高训练速度。CPU卸载还可以帮助减少GPU内存的使用，降低训练过程中的内存压力。

nvme
使用 NVMe 固态硬盘扩展 GPU 是一种有效的 CPU 卸载技术。NVMe 固态硬盘具有高速的数据传输速度和低延迟，可以大大提高数据传输效率。此外，NVMe 固态硬盘还支持多通道和多队列，可以实现并行传输和处理，进一步提高数据传输速度。在使用 NVMe 固态硬盘时，可以通过在 CPU 和 GPU 之间使用高带宽 PCIe 总线，进一步优化数据传输速度，提高训练效率和速度。

FSDP
Fully Sharded Data Paralle（FSDP）和 DeepSpeed 类似，均通过 ZeRO 等分布优化算法，减少内存的占用量。其将模型参数，梯度和优化器状态分布至多个 GPU 上，而非像 DDP 一样，在每个 GPU 上保留完整副本。
Huggingface 这篇博文解释了 ZeRO 的大致实现方法：
https://huggingface.co/blog/zero-deepspeed-fairscale
借助 torch 实现 FSDP，只需要将 model 用 FSDPwarp 一下；同样，cpu_offload 也只需要一行代码：
https://pytorch.org/blog/introducing-pytorch-fully-sharded-data-parallel-api/
在这个可以查看 FSDP 支持的模型：
https://pytorch.org/docs/stable/fsdp.html
在 Huggingface Transformers 中使用 Torch FSDP：
https://huggingface.co/docs/transformers/v4.27.2/en/main_classes/trainer#transformers.Trainin
根据某些 issue，shard_grad_op（只分布保存 optimizer states 和 gradients）模式可能比 fully_shard 更稳定：
https://github.com/tatsu-lab/stanford_alpaca/issues/32

4 Gradient Checkpointing

在 torch 中使用 - 把 model 用一个 customize 的 function 包装一下即可，详见：
Explore Gradient-Checkpointing in PyTorch
https://qywu.github.io/2019/05/22/explore-gradient-checkpointing.html
在 Huggingface Transformers 中使用：
https://huggingface.co/docs/transformers/v4.27.2/en/perf_train_gpu_one#gradient-checkpointing