PyTorch——Accelerate使用总结

• 参考链接：
  • 官方文档：docs.pytorch.org/docs/stable/fsdp.html

---

整体说明

• HuggingFace Accelerate 是一个轻量级库（accelerate），专为简化 PyTorch 模型在各种硬件配置上的训练和推理而设计
• 它能自动处理分布式训练、混合精度训练等复杂设置，让开发者无需深入了解底层硬件细节，就能轻松将模型部署到单 GPU、多 GPU、TPU 甚至 CPU 集群等环境中（专注于模型逻辑和训练流程即可）
• 安装简便，仅需一行代码：pip install accelerate
• Accelerate 的核心功能包括下面几个
  • 自动识别可用硬件（GPU、TPU 等），并根据硬件情况优化训练配置
  • 无缝支持数据并行、模型并行等分布式训练模式 ，适配多 GPU 或集群环境
  • 可选择 FSDP 或 DeepSpeed 等底层框架，仅需简单修改启动命令即可
  • 支持 FP16、BF16 等混合精度训练 ，在减少显存占用的同时，保证模型训练精度
  • 只需对原有 PyTorch 代码进行少量修改 ，即可实现硬件加速和分布式训练
• 一般来说仅需要两行代码改动，其他都有命令行进行配置
  • accelerator.prepare() ：核心函数，用于包装模型、优化器、数据加载器等组件，自动适配分布式和混合精度设置
  • accelerator.backward() ：替代传统的 loss.backward()，在分布式环境中确保梯度正确同步

HuggingFace Accelerate 使用代码示例

• 以下是一个使用 Accelerate 进行模型训练的基础示例，训练代码（train.py）如下：

  import torch
  import torch.nn as nn
  import torch.optim as optim
  from torch.utils.data import Dataset, DataLoader
  from accelerate import Accelerator  # 导入 Accelerator
  
  class DiyModel(nn.Module):
      def __init__(self):
          super().__init__()
          self.fc = nn.Linear(10, 2) # 二分类任务
      def forward(self, x):
          return self.fc(x)
  
  class DiyDataset(Dataset):
      def __len__(self):
          return 1000
      def __getitem__(self, idx):
          x = torch.randn(10) 
          y = torch.randint(0, 2, (1,)).item()  # 随机标签（0 或 1）
          return x, y
  
  # 混合精度配置：可通过 `Accelerator(mixed_precision="fp16")` 启用 FP16 混合精度训练，减少显存占用
  accelerator = Accelerator() # 核心代码，初始化 Accelerator
  
  # 以下所有定义都不涉及使用 Accelerator
  model = DiyModel()
  criterion = nn.CrossEntropyLoss()
  optimizer = optim.Adam(model.parameters(), lr=1e-3)
  dataset = DiyDataset()
  dataloader = DataLoader(dataset, batch_size=32, shuffle=True)
  
  # 重点：使用 accelerator 包装多个组件（注：这一行会自动处理分布式和混合精度）
  model, optimizer, dataloader, criterion = accelerator.prepare(
      model, optimizer, dataloader, criterion
  )
  
  # 特别注意的一点不同是：训练循环时，使用 accelerator.backward()
  model.train()
  for epoch in range(3):
      total_loss = 0.0
      for x, y in dataloader:
          optimizer.zero_grad()
          outputs = model(x)
          loss = criterion(outputs, y)
          accelerator.backward(loss)  # 替代 loss.backward()
          optimizer.step()
          total_loss += loss.item()
      avg_loss = total_loss / len(dataloader)
      print(f"Epoch {epoch+1}, Loss: {avg_loss:.4f}")

Accelerate 启动训练

• 可简单通过 Accelerate 命令行工具配置训练环境并启动：

  accelerate launch --num_processes=2 train.py  # 使用 2 个进程（如 2 个 GPU）

  • --num_processes：指定进程数（通常等于 GPU 数量）
  • 若使用单 GPU 或 CPU，可直接运行 python train.py，Accelerator 会自动适配环境

• 更多启动命令参见下文

---

accelerate launch 命令详细说明

• accelerate launch 是 HuggingFace Accelerate 库的核心命令，用于启动分布式训练脚本，它能自动处理多卡、多机等复杂分布式环境的配置

• accelerate launch 的核心作用是：

  • 1）初始化分布式环境（进程组、通信后端等）
  • 2）根据参数自动选择分布式策略（数据并行/FSDP/DeepSpeed 等）
  • 3）将环境配置传递给训练脚本中的 Accelerator 实例，使其能正确处理模型、数据的分布式适配

• accelerate launch 命令基本语法

  accelerate launch [启动参数] your_script.py [脚本参数]

  • [启动参数]：控制分布式训练的配置（如使用的 GPU 数量、分布式策略等）
  • [脚本参数]：传递给你的训练脚本（your_script.py）的自定义参数（如 --epochs 10、--batch_size 32 等）

硬件与进程配置参数

• --num_processes N：指定总进程数（通常等于参与训练的 GPU 总数）
  • 例如：--num_processes 4 表示使用 4 个 GPU
• --num_machines N：指定机器数量（多机分布式训练时使用），默认值为 1（单机器）
• --machine_rank N：当使用多机时，指定当前机器的序号（从 0 开始）
  • 例如：主节点用 --machine_rank 0，从节点用 --machine_rank 1
• --main_process_ip IP地址：多机训练时，主节点的 IP 地址（供从节点连接）
• --main_process_port 端口号：主节点的通信端口（默认 29500，需确保端口未被占用）

分布式策略选择参数

• Accelerate 支持多种分布式策略，通过参数指定：
• 默认策略（自动选择） ：不指定任何策略时，Accelerate 会根据硬件自动选择当前硬件下的最佳策略：
  • 单卡：直接使用单进程训练
  • 多卡：默认使用 PyTorch 的 nn.DataParallel 或 DistributedDataParallel（数据并行）
• --use_fsdp：
  • 启用 FSDP（完全分片数据并行） ，适合超大规模模型（需 PyTorch ≥ 1.11）；
  • 常用搭配参数如下：
    • --fsdp_fully_shard：完全分片参数、梯度和优化器状态（最大程度节省内存）
    • --fsdp_transformer_layer_cls_to_wrap "类名"：指定 Transformer 层的类名（如 GPT2 的 GPT2Layer、BERT 的 BertLayer），用于自动分片模型层
    • --fsdp_sharding_strategy 策略：分片策略，可选 FULL_SHARD（完全分片）、SHARD_GRAD_OP（梯度和优化器分片）等
• --use_deepspeed：
  • 启用 DeepSpeed 分布式框架，支持 ZeRO 优化、混合精度等（需提前安装 deepspeed）
  • 通常需配合 DeepSpeed 配置文件使用，通过 --deepspeed 配置文件路径 指定

混合精度训练参数

• --mixed_precision [mode]：指定混合精度策略
• 可选模式 [mode] 为：
  • no：不使用混合精度（默认）
  • fp16：使用 FP16 混合精度
  • bf16：使用 BF16 混合精度（需 GPU 支持，如 A100、RTX 3090 等）
  • fp8：使用 FP8 混合精度（需 PyTorch ≥ 2.0 且 GPU 支持）

其他实用参数

• --config_file 配置文件路径：通过 YAML 配置文件指定所有参数（推荐复杂场景使用），无需在命令行逐个输入
• --debug：启用调试模式，打印详细的分布式初始化日志，便于排查问题
• --gradient_accumulation_steps N：指定梯度累积步数（等价于在代码中设置，但通过命令行更灵活）

---

附录：在启动命令中使用配置文件用法（推荐）

• 对于复杂配置（如 FSDP/DeepSpeed 细节），建议使用 YAML 配置文件，步骤如下：

• 第一步：生成默认配置文件 ：

  accelerate launch --config_file accelerate_config.yaml --generate_config

  • --generate_config 表示运行后会交互式提问，自动生成配置文件
  • 注：也可以自己编辑 accelerate_config.yaml 文件

• 第二步：示例配置文件（FSDP 场景） ：：

  compute_environment: LOCAL_MACHINE  # 本地机器环境
  distributed_type: FSDP             # 使用 FSDP 策略
  fsdp_config:
      fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP  # 自动包装 Transformer 层
      fsdp_transformer_layer_cls_to_wrap: "GPT2Layer"  # 模型层类名
      fsdp_sharding_strategy: FULL_SHARD  # 完全分片
  mixed_precision: fp16              # 启用 FP16 混合精度
  num_processes: 4                   # 4 个进程（4 卡）

• 第三步：使用配置文件启动 ：

  accelerate launch --config_file accelerate_config.yaml your_script.py --epochs 10

---

附录：启动命令的常见场景示例

• 示例一：单机器多卡基础数据并行示例

  accelerate launch --num_processes 4 train.py --batch_size 32

• 示例二：启用 FSDP 训练大模型

  accelerate launch \
      --num_processes 4 \
      --use_fsdp \
      --fsdp_fully_shard \
      --fsdp_transformer_layer_cls_to_wrap "BertLayer" \
      --mixed_precision bf16 \
      train_bert.py

• 示例三：启用 DeepSpeed 与 ZeRO-3 优化（其中 ds_config.json 为 DeepSpeed 配置文件，定义 ZeRO 阶段、梯度裁剪等）

  accelerate launch \
      --num_processes 8 \
      --use_deepspeed \
      --deepspeed ds_config.json \
      train.py