huggingface笔记：使用accelerate加速

1 介绍

随着模型规模的增大，并行处理已成为在有限硬件上训练大型模型和提高训练速度的重要策略。
Hugging Face 创建了Accelerate库，帮助用户在任何类型的分布式环境中轻松训练Transformers模型，无论是单机多GPU还是跨多机的多GPU

2 创建Accelerator对象

from accelerate import Accelerator

accelerator = Accelerator()

Accelerator类初始化分布式训练所需的一切，并且会根据代码启动的方式自动检测你的训练环境（单机单GPU、单机多GPU、多机多GPU或TPU等）
默认情况下，Accelerator类会自动将对象放置到合适的设备上
- 可以通过在初始化Accelerator时传递device_placement=False来关闭自动设备放置
- 如果想明确地将对象放置到某个设备上，确保使用accelerator.device

3 训练对象传递给prepare方法

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

4 反向传播

唯一不同的是把loss.backward()替换成Accelerate的backward

5 总结

绿色的是加上的，红色的是去掉的

6 训练

运行以下命令创建并保存配置文件

accelerate config

然后启动训练：

accelerate launch train.py

7 保存和加载

使用 save_state() 将所有内容保存到一个文件夹位置
使用 load_state() 加载之前通过 save_state() 保存的所有内容

通过使用 register_for_checkpointing()，可以注册自定义对象以便通过前两个函数自动存储或加载

7.1 unwrap_model

解包模型，在保存模型之前，用于从 prepare() 可能添加的额外层中解包模型
- 因为 prepare() 方法将模型包装成适合分布式训练的接口。
- 如果不解包模型，保存模型状态字典时也会保存大模型中可能存在的额外层，之后就无法将权重加载回基础模型

# 假设有两个GPU进程
from torch.nn.parallel import DistributedDataParallel
from accelerate import Accelerator

accelerator = Accelerator()
model = accelerator.prepare(MyModel())
print(model.__class__.__name__)
#DistributedDataParallel 

model = accelerator.unwrap_model(model)
print(model.__class__.__name__)
#MyModel

torch.save(model.state_dict(), path)
#解包之后，可以用torch的save函数

还可以使用使用 accelerate提供的save_model() 方法来解包并保存模型状态字典。

accelerator.wait_for_everyone()
accelerator.save_model(model, save_directory)

7.1.1 使用save_pretrained 保存模型

这样可以后续用from_pretrained的方式再得到模型

from transformers import AutoModel

unwrapped_model = accelerator.unwrap_model(model)
unwrapped_model.save_pretrained(
    "path/to/my_model_directory",
    is_main_process=accelerator.is_main_process,
    save_function=accelerator.save,
)

model = AutoModel.from_pretrained("path/to/my_model_directory")

8 分布式评估

8.1 gather_for_metrics

在分布式系统中，每个设备只接收一部分评估数据，这意味着你应该使用 gather_for_metrics() 方法将你的预测结果组合在一起。
此方法会收集 input_data 并可能在最后一个批次中删除重复项。
可应用于收集用于指标计算的输入和目标。
此方法要求每个进程上的所有张量大小相同，因此如果张量在每个进程上的大小不同（例如在批处理中动态填充到最大长度时），应该使用 pad_across_processes() 方法将你的张量填充到跨进程的最大大小

for inputs, targets in validation_dataloader:
    predictions = model(inputs)


    # Gather all predictions and targets
    all_predictions, all_targets = accelerator.gather_for_metrics((predictions, targets))
    

    # Example of use with a *Datasets.Metric*
    metric.add_batch(all_predictions, all_targets)

8.2 pad_across_processes

( tensordim = 0，pad_index = 0，pad_first = False ) → torch.Tensor,

递归地填充来自所有设备的嵌套列表/元组/字典中的张量，以便它们可以安全地被聚合。

tensor	要聚合的数据
dim	填充的维度
pad_index	用于填充的值
pad_first	是否在开始或结束时进行填充

# 假设有两个进程，第一个进程有一个大小为1的张量，第二个进程的张量大小为2
import torch
from accelerate import Accelerator

accelerator = Accelerator()
process_tensor = torch.arange(accelerator.process_index + 1).to(accelerator.device)
'''
accelerator.process_index：这是当前进程的索引。

在两个进程的设置中，第一个进程的索引为 0，第二个进程的索引为 1。
'''
padded_tensor = accelerator.pad_across_processes(process_tensor)
'''
它确保所有进程中的张量都被填充到相同的大小，这样它们就可以安全地聚合或进行其他并行操作
'''
padded_tensor.shape
#torch.Size([2])

9 accelerate原理

在内部，Accelerate 首先分析启动脚本的环境，以确定使用的分布式设置类型、有多少个不同的进程以及当前脚本位于哪个进程中
然后，调用 prepare() 时，库会：
- 将模型封装在适用于分布式设置的容器中，
- 将优化器封装在 AcceleratedOptimizer 中，
- 将调度器封装在 AcceleratedScheduler 中，
- 在 DataLoaderShard 或 DataLoaderDispatcher 中创建数据加载器的新版本。
  - 创建的原因是， PyTorch 不允许用户在数据加载器创建后更改其批处理采样器
  - DataLoaderShard 继承自 DataLoader，添加了以下功能：
    - 它在每次新迭代时同步所有进程的适当随机数生成器，以确保任何随机化（如洗牌）都以完全相同的方式跨进程完成。
    - 它在产生批次之前将批次放置在适当的设备上（除非选择了 device_placement=True）。
  - DataLoaderDispatcher 与 DataLoaderShard 的不同之处在于，在迭代 DataLoader 时，数据都是从进程 0 开始的，然后分割并发送到每个进程，而不是在数据集级别进行

10 混合精度

混合精度通过使用较低精度的数据类型（如fp16，即半精度）来计算梯度，从而加速训练。
为了在Accelerate中获得最佳性能，应该在模型内部计算损失（如在Transformers模型中），因为模型外部的计算是以全精度进行的。
- 设置Accelerator中使用的混合精度类型，然后使用autocast()上下文管理器自动将值转换为指定的数据类型。
- Accelerate启用了自动混合精度，因此只有在除了backward()已经处理的损失上的混合精度操作之外，还有其他混合精度操作时，才需要使用autocast()。

accelerator = Accelerator(mixed_precision="fp16")
with accelerator.autocast():
    loss = complex_loss_function(outputs, target)