HuggingFace Accelerate 是一个面向 PyTorch 的分布式训练统一 API。其定位是把分布式训练简化为只需少量代码改动:资料称只需 4 行即可为任意 PyTorch 脚本添加分布式支持,并用同一套 API 统一 DeepSpeed/FSDP/Megatron/DDP,自动处理设备放置。
安装与改造方式 安装命令为 pip install accelerate。改造脚本的基本方式是在原有 PyTorch 训练代码中导入 Accelerator、创建实例 accelerator = Accelerator()、用 accelerator.prepare(model, optimizer, dataloader) 包装模型与优化器及数据加载器,并将 loss.backward() 替换为 accelerator.backward(loss)。改造后,数据移动到设备等操作交由框架自动处理,无需再手动调用 .to('cuda')。启动使用统一命令 accelerate launch train.py。
主要能力
- 硬件与规模适配:同一脚本可运行于 CPU、单 GPU、多 GPU、多节点、TPU 与 Apple MPS。多卡可用 accelerate launch --multi_gpu --num_processes 8 train.py,多机可配合 --num_machines、--machine_rank、--main_process_ip 等参数。
- 交互式配置:通过 accelerate config 回答单机/多机、机器数量、混合精度、是否启用 DeepSpeed 等问题,无需手工搭建启动器。
- 混合精度训练:支持 mixed_precision='fp16'(含梯度缩放)、'bf16'(无缩放、更稳定)与 'fp8'(资料标注为 H100+),并可用 accelerator.autocast() 控制前向计算。
- DeepSpeed ZeRO 集成:可通过 deepspeed_plugin 参数配置 zero_stage、offload_optimizer、gradient_accumulation_steps 等,也可通过 accelerate config 或 JSON 配置文件启用,并用 accelerate launch --config_file 指定。
- FSDP 支持:使用 FullyShardedDataParallelPlugin 配置 sharding_strategy(如 FULL_SHARD)、auto_wrap_policy、cpu_offload 等,或通过交互式配置选择。
- 梯度累积:Accelerator(gradient_accumulation_steps=4) 配合 with accelerator.accumulate(model) 上下文管理器实现,有效批大小为 batch_size × GPU 数 × 累积步数。
- 分布式检查点:用 accelerator.is_main_process 判断主进程保存 accelerator.save_state(),各进程用 accelerator.load_state() 加载;用 accelerate.utils.set_seed(42) 统一随机种子以保持 FSDP 下结果一致。
适用场景 资料建议在以下情况使用:希望以最简单方式做分布式训练;需要单一脚本适配任意硬件;使用 HuggingFace 生态(Transformers、TRL、PEFT 等均基于它);需要在 DDP/DeepSpeed/FSDP/Megatron 间保持灵活切换;需要快速原型验证。
对比与限制 若需要回调与高层抽象可选 PyTorch Lightning;需要多节点编排与超参调优可选 Ray Train;需要直接 API 控制与高级特性可选 DeepSpeed;需要最大控制权与最少抽象可选原生 DDP。启动依赖方面:DDP 使用内置 torch.distributed.run,DeepSpeed 需 pip install deepspeed,FSDP 需 PyTorch 1.12+,Megatron 需自定义环境。资料注明版本为 1.11.0+,Megatron 集成、自定义插件与性能调优细节分别见其参考文档。