ARTICLE · INTELLIGENCE

战地情报 · 详情页

来自尧图项目组的一线实战观察与深度解析

使用 [特殊字符] Accelerate 在 Transformers 中实现分布式训练:从手写 PyTorch 训练循环到 Trainer 集成

使用 [特殊字符] Accelerate 在 Transformers 中实现分布式训练:从手写 PyTorch 训练循环到 Trainer 集成 使用 Accelerate 在 Transformers 中实现分布式训练从手写 PyTorch 训练循环到 Trainer 集成【免费下载链接】transformers Transformers: the model-definition framework for state-of-the-art machine learning models in text, vision, audio, and multimodal models, for both inference and training.项目地址: https://gitcode.com/GitHub_Trending/tra/transformers随着模型参数量持续增长并行parallelism已成为在有限硬件上训练大模型、并把训练速度提升数个数量级的关键策略。本指南以 Accelerate 为核心完整讲解两条实战路径一是如何用 Accelerate 改造原生 PyTorch 训练循环仅需新增数行代码即可支持单机多卡与多机多卡二是如何通过 Accelerate 配置文件与TrainingArguments让Trainer自动接管 FSDP、DeepSpeed、DDP 等分布式后端的全部底层机制。读完后你将掌握分布式训练的配置、改造与启动全流程并理解其与 Transformers 源码的集成原理。为什么需要 Accelerate统一分布式训练的复杂性在 Hugging Face 团队看来并行化是解决两大现实问题的核心手段在受限硬件上训练大模型以及在多卡环境下成倍提升训练吞吐。但分布式训练的后端五花八门——FSDP、DeepSpeed、DDP、混合精度、多机通信……每种都有自己的配置方式和 API。 Accelerate 提供的正是这样一个统一接口它自动探测你的环境GPU 数量、分布式后端、混合精度设置等无论你是在单卡上用 DDP还是在 8 卡上用 FSDP都能自动完成配置。Accelerate 与 Transformers 的分工非常清晰Accelerate 负责把模型包进合适的分布式包装器、移动到正确的设备、创建兼容的优化器训练过程中Accelerate 使用自己的Accelerator.backward方法处理混合精度下的梯度缩放Trainer调用相应的 Accelerate API把所有分布式机制委托给 Accelerate 完成。因此无论是手写训练循环还是使用高层 TrainerAccelerate 都是分布式训练能力背后的统一引擎。路径一改造原生 PyTorch 训练循环这是关联文档 es/accelerate.md 的主体内容在不引入 Trainer的情况下用三步改造让一个普通 PyTorch 训练脚本跑在任意分布式环境中。第一步安装并初始化 Accelerator安装 Acceleratepip install accelerate然后导入并创建Accelerator对象from accelerate import Accelerator accelerator Accelerator()Accelerator会自动检测你当前可用的分布式配置类型并初始化训练所需的全部组件。你无需手动指定模型应该放到哪个设备——设备放置cuda/cpu/多卡完全由 Accelerate 接管。第二步用 prepare 准备训练对象把训练相关的所有对象传给prepare方法包括训练/评估的 DataLoader、模型和优化器train_dataloader, eval_dataloader, model, optimizer accelerator.prepare( train_dataloader, eval_dataloader, model, optimizer )prepare是分布式改造的核心动作它会依据探测到的环境把 DataLoader 包装成分片采样保证每个进程拿到不重叠的数据分片、把模型包进 DDP/FSDP 等分布式包装器并移动到正确设备、为优化器创建兼容版本。这也是后文 Trainer 集成中create_accelerator_and_postprocess内部实际执行的逻辑。第三步用 accelerator.backward 替换 loss.backward()在训练循环中把常规的loss.backward()替换为 Accelerate 的backward方法。它负责在混合精度场景下正确地进行梯度缩放gradient scaling这是手写循环最容易出错的地方for epoch in range(num_epochs): for batch in train_dataloader: outputs model(**batch) loss outputs.loss accelerator.backward(loss) optimizer.step() lr_scheduler.step() optimizer.zero_grad() progress_bar.update(1)完整改造对比只需 4 行代码下面的 diff 展示了一个基于AutoModelForSequenceClassification的情感分类训练脚本的完整改造过程。注意三处变化新增accelerator Accelerator()删除手动的device torch.device(cuda) ...和model.to(device)以及循环内batch {k: v.to(device) for k, v in batch.items()}用accelerator.prepare(...)一次性包装 DataLoader、模型和优化器用accelerator.backward(loss)替换loss.backward()。 from accelerate import Accelerator from transformers import AdamW, AutoModelForSequenceClassification, get_scheduler accelerator Accelerator() model AutoModelForSequenceClassification.from_pretrained(checkpoint, num_labels2) optimizer AdamW(model.parameters(), lr3e-5) - device torch.device(cuda) if torch.cuda.is_available() else torch.device(cpu) - model.to(device) train_dataloader, eval_dataloader, model, optimizer accelerator.prepare( train_dataloader, eval_dataloader, model, optimizer ) num_epochs 3 num_training_steps num_epochs * len(train_dataloader) lr_scheduler get_scheduler( linear, optimizeroptimizer, num_warmup_steps0, num_training_stepsnum_training_steps ) progress_bar tqdm(range(num_training_steps)) model.train() for epoch in range(num_epochs): for batch in train_dataloader: - batch {k: v.to(device) for k, v in batch.items()} outputs model(**batch) loss outputs.loss - loss.backward() accelerator.backward(loss) optimizer.step() lr_scheduler.step() optimizer.zero_grad() progress_bar.update(1)核心收益在于你不再需要关心设备字符串、to(device)调用和数据搬运Accelerate 统一处理了设备放置、梯度同步与混合精度。这正是文档强调的只需要往训练循环里加四行代码。启动训练脚本方式与 Notebook 方式改造完成后可以从脚本或 Notebook如 Colaboratory启动训练。方式一脚本启动推荐适合多卡/多机首先运行交互式配置向导回答关于硬件和训练设置的问题它会生成并保存一份配置文件accelerate config然后启动训练accelerate launch train.py在 examples/pytorch/README.md 中还可以看到官方推荐的完整三步流程accelerate config回答向导问题 →accelerate test验证环境就绪 →accelerate launch path_to_script.py --args_to_script正式启动。仓库内大量的*_no_trainer.py示例脚本如 run_qa_no_trainer.py都遵循这一模式例如 run_swag_no_trainer.sh 中直接用accelerate launch run_swag_no_trainer.py拉起训练。方式二Notebook 启动适合 Colab TPU如果计划使用 Colaboratory 的 TPU可以把负责训练的代码封装进一个函数再传给notebook_launcherfrom accelerate import notebook_launcher notebook_launcher(training_function)路径二通过 Trainer 集成 Accelerate当使用高层Trainer时你不再需要手写prepare/backward——Trainer 会在内部完成这一切。配置方式有两种Accelerate 配置文件或TrainingArguments参数。方式 AAccelerate 配置文件先运行accelerate config并回答关于硬件与训练设置的问题它会在缓存目录生成一份default_config.yaml。以下是一份 FSDP 配置示例来自 en/accelerate.mdcompute_environment: LOCAL_MACHINE distributed_type: FSDP fsdp_config: fsdp_version: 2 fsdp_reshard_after_forward: true fsdp_cpu_offload: false fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP fsdp_cpu_ram_efficient_loading: true fsdp_activation_checkpointing: false fsdp_state_dict_type: SHARDED_STATE_DICT fsdp_transformer_layer_cls_to_wrap: LlamaDecoderLayer mixed_precision: bf16 num_machines: 1 num_processes: 4关键字段速览compute_environment计算环境类型本地机器为LOCAL_MACHINEdistributed_type分布式后端可选FSDP、DEEPSPEED、MULTI_GPUDDP等fsdp_configFSDP 专属配置块包括分片策略fsdp_auto_wrap_policy、CPU 卸载fsdp_cpu_offload、激活检查点fsdp_activation_checkpointing以及需要包装的 Transformer 层类fsdp_transformer_layer_cls_to_wrap: LlamaDecoderLayer等mixed_precision混合精度类型如bf16、fp16num_machines/num_processes机器数量与总进程数通常等于 GPU 总数。之后用accelerate launch运行基于 Trainer 的脚本Accelerate 会读取该配置文件完成训练环境搭建accelerate launch train.py值得注意使用 Accelerate 配置文件时TrainingArguments中的fsdp_config和deepspeed参数就不再需要设置了因为配置文件已覆盖相同的配置项避免双重配置引发冲突。通过 accelerator_config 传入无专用参数的设置TrainingArguments的accelerator_config参数用于接收那些没有独立顶层参数的 Accelerate 设置。一个典型例子是non_blockingTrue把它与dataloader_pin_memoryTrue搭配可以让数据搬运与计算重叠提升 GPU 吞吐from transformers import TrainingArguments TrainingArguments( ..., dataloader_pin_memoryTrue, accelerator_config{ non_blocking: True, }, )这一设计在源码中有直接印证accelerator_config字段接受一个 JSON 配置文件路径或字典。在 create_accelerator_and_postprocess 中Trainer 会从accelerator_config中提取出split_batches、dispatch_batches、even_batches、use_seedable_sampler等 DataLoader 相关参数构造DataLoaderConfiguration并单独取出non_blocking——同时还会校验如果开启了non_blocking却没有开启dataloader_pin_memory会打印性能建议警告源码见 trainer.py。方式 B通过 TrainingArguments 配置后端也可以直接向TrainingArguments传入后端专属配置由create_accelerator_and_postprocess读取并据此完成训练配置。三种常用后端对比如下。FSDP全分片数据并行把 JSON 配置文件路径或字典传给fsdp_configfrom transformers import TrainingArguments TrainingArguments( ..., fsdpTrue, fsdp_configpath/to/fsdp.json, )完整指南与配置参考见 FSDP 指南。DeepSpeedZeRO 优化与卸载把 JSON 配置文件路径或字典传给deepspeedfrom transformers import TrainingArguments TrainingArguments( ..., deepspeedpath/to/ds_config.json, )完整指南与配置参考见 DeepSpeed 指南。DDP数据并行模型能放进单卡时DDP 直接通过TrainingArguments字段配置from transformers import TrainingArguments TrainingArguments( ..., ddp_backendnccl, ddp_find_unused_parametersFalse, ddp_bucket_cap_mb25, ddp_timeout1800, )常用 DDP 参数说明ddp_backend通信后端NVIDIA GPU 用nccl默认且最快CPU 训练或调试用gloo其他硬件分别有xccl、hccl、cncl等ddp_find_unused_parameters在前向结束遍历 autograd 图为不会收到梯度的参数标记 ready避免阻塞 all-reduce不要与梯度检查点gradient_checkpointing同用ddp_bucket_cap_mb反向时将梯度打包进单次 all-reduce 的桶大小更大的桶意味着更少的 all-reduce 调用和更低的启动开销ddp_timeout所有进程与通信操作all-reduce、broadcast的时间上限进程挂起时抛出错误而非无限阻塞。更完整的 DDP 机制讲解含梯度累积与 all-reduce 时序见 DDP 指南。源码视角Trainer 如何创建 Accelerator理解了两种配置方式后再深入源码看清 Trainer 内部的集成流程这能帮助你诊断分布式训练中的问题。create_accelerator_and_postprocesstrainer.py是集成的枢纽方法其执行流程大致如下梯度累积交由 Trainer 自己处理显式地把gradient_accumulation_kwargs[num_steps]设为 1避免与TrainingArguments.gradient_accumulation_steps双重计算源码见 trainer.py解析accelerator_config拆出 DataLoader 配置与non_blocking等参数调用_build_accelerator_argstrainer.py组装 Accelerator 关键字参数包括mixed_precision、deepspeed_plugin、由DistributedDataParallelKwargs包装的 DDP 参数find_unused_parameters、bucket_cap_mb、broadcast_buffers、static_graph、dataloader_config、gradient_accumulation_plugin、可选的fsdp_plugin与dynamo_plugintorch.compile 后端以及要求 Accelerate ≥ 1.12.0 的张量并行parallelism_config创建 Accelerator 实例self.accelerator Accelerator(**args)trainer.py后处理通过accelerator.state检测是否启用了 DeepSpeed/FSDPis_deepspeed_enabled、is_fsdp_enabled并做冲突校验——例如 FSDP 的activation_checkpointing与TrainingArguments.gradient_checkpointing不能同时为 Truetrainer.py。值得留意的一点是 DDP 的find_unused_parameters默认值推导逻辑当模型是PreTrainedModel时Trainer 会自动取not (model.is_gradient_checkpointing or args.gradient_checkpointing)即开启梯度检查点时自动关闭该选项这与官方文档不要与 gradient_checkpointing 同用的告诫完全一致trainer.py。下一步深入方向模型能放进单卡时阅读 DDP 指南 掌握数据并行训练的梯度同步细节模型超出单卡显存时阅读 FSDP 指南 学习跨卡分片参数、梯度与优化器状态需要 ZeRO 优化与 CPU/异构卸载时阅读 DeepSpeed 指南。总而言之无论你选择手写训练循环还是高层 TrainerAccelerate 都是 Transformers 分布式训练的统一底座前者通过Acceleratorpreparebackward三步改造即可跑通多卡后者只需一份accelerate config配置文件或几个TrainingArguments参数Trainer 便会自动完成分布式后端的创建、包装与校验。理解这条集成链路你就能在任何分布式环境中自如地训练 Transformers 模型。【免费下载链接】transformers Transformers: the model-definition framework for state-of-the-art machine learning models in text, vision, audio, and multimodal models, for both inference and training.项目地址: https://gitcode.com/GitHub_Trending/tra/transformers创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考
RELATED READING

延伸阅读

更多一线实战笔记与深度复盘,助您持续精进