人工智能深度学习分布式训练【免费下载链接】accelerate A simple way to launch, train, and use PyTorch models on almost any device and distributed configuration, automatic mixed precision (including fp8), and easy-to-configure FSDP and DeepSpeed support项目地址https://gitcode.com/gh_mirrors/ac/accelerate点击查看免费下载本指南系统讲解如何在 Accelerate 分布式训练工作流中接入 PyTorch 2.0 的torch.compile先介绍TorchDynamoPlugin的完整配置方式与底层实现再深入区域编译Regional Compilation原理并结合仓库内基准测试与源码佐证帮助你掌握既享受编译加速、又显著压低 LLM 类模型首次编译冷启动时间的实战方案。一、为什么要在 Accelerate 中使用 torch.compilePyTorch 2.0 引入了torch.compile它通过 JIT 将 PyTorch 代码编译为优化后的内核从而显著加速模型执行。其核心特性可概括为三点性能提升Performance Improvement优化计算图显著加快模型执行速度易于使用Ease of Use对既有代码改动极小接入成本低兼容性强Compatibility与现有 PyTorch 代码和模型无缝协作。当与 Accelerate 配合使用时torch.compile会平滑地融入分布式训练流程使你在同一套代码中同时获得分布式执行与编译优化的收益。需要说明的是编译后的代码首次执行通常更慢其中包含编译耗时而后续运行会明显加快。为此torch.compile提供了多种编译模式modedefault默认、reduce-overhead利用 CUDA graphs 进一步压低调度开销以及max-autotune大规模自动调优为模型寻找最佳内核。二、TorchDynamoPlugin一行配置接入编译Accelerate 通过TorchDynamoPlugin提供对torch.compile的无缝集成。其定义位于 src/accelerate/utils/dataclasses.py继承自KwargsHandler会在模型预处理阶段把编译参数透传给torch.compile。2.1 最小可用示例from accelerate import Accelerator from accelerate.utils import TorchDynamoPlugin # 配置编译后端 dynamo_plugin TorchDynamoPlugin( backendinductor, # 可选inductor、aot_eager、aot_nvfuser 等 modedefault, # 可选default、reduce-overhead、max-autotune fullgraphTrue, dynamicFalse, ) # 使用插件初始化 accelerator accelerator Accelerator(dynamo_plugindynamo_plugin) # 这一步会对模型应用 torch.compile model accelerator.prepare(model)该插件与 Accelerate 的所有其他特性与插件兼容包括混合精度、分布式训练DDP、FSDP、DeepSpeed等。其兼容性在源码中得到印证Accelerator.prepare的常规路径src/accelerate/accelerator.py、FSDP2 路径src/accelerate/accelerator.py以及 DeepSpeed 路径src/accelerate/accelerator.py中均按各自方式触发编译。2.2 完整参数说明取自源码字段定义下表整理了 TorchDynamoPlugin 的全部可配置字段参数默认值含义backendNone回退到环境变量再回退noDynamo 编译后端取值见下方DynamoBackend枚举modeNone回退环境变量再回退default编译模式default、reduce-overhead或max-autotunefullgraphNone回退环境变量再回退False是否允许将模型拆分为多个子图True表示要求完整图dynamicNone仅当设置了对应环境变量时生效是否使用动态 shape 进行 tracingoptionsNone传给后端的一本参数字典disableFalse将torch.compile()变成 no-op便于测试use_regional_compilationNone回退环境变量再回退False是否启用区域编译见本文第三节其中backend的取值对应DynamoBackend枚举src/accelerate/utils/dataclasses.py包括NO不使用 torch dynamoEAGER用 PyTorch 直接运行提取出的 GraphModule适合调试 TorchDynamo 问题AOT_EAGER仅使用 AotAutograd 而无编译器适合调试通常不会有加速INDUCTOR默认推荐后端基于 AotAutograd、CUDA graphs 与 codegen 出的 Triton 内核是获取明显加速的主流选择AOT_TS_NVFUSER / NVPRIMS_NVFUSERnvFuser 与 AotAutograd/TorchScript 或 PrimTorch 组合CUDAGRAPHScudagraphs 配合 AotAutogradOFI / FX2TRT / ONNXRT / TENSORRT / TVM面向推理优化的后端TensorRT、ONNX Runtime、TVM 等AOT_TORCHXLA_TRACE_ONCE / TORCHXLA_TRACE_ONCEPyTorch/XLA 的 Dynamo 优化分别面向训练与推理HPU_BACKEND / NEURONHabana HPU 与 AWS NeuronTrainium/Inferentia后端。注意backendno时prepare不会触发任何编译DynamoBackend在内部通过.upper()归一化因此大小写不敏感。2.3 环境变量配置方式TorchDynamoPlugin.__post_init__src/accelerate/utils/dataclasses.py显示所有参数都支持通过ACCELERATE_DYNAMO_前缀的环境变量覆盖便于在不改代码的情况下切换配置环境变量默认值ACCELERATE_DYNAMO_BACKENDnoACCELERATE_DYNAMO_MODEdefaultACCELERATE_DYNAMO_USE_FULLGRAPHFalseACCELERATE_DYNAMO_USE_DYNAMIC未设置时不启用ACCELERATE_DYNAMO_USE_REGIONAL_COMPILATIONFalse例如export ACCELERATE_DYNAMO_BACKENDinductor export ACCELERATE_DYNAMO_MODEreduce-overhead export ACCELERATE_DYNAMO_USE_REGIONAL_COMPILATIONTrue随后直接Accelerator()即可生效。相关配置项同样可以在accelerate config生成的 YAML 或accelerate launch命令行参数如--dynamo_backend、--dynamo_mode中指定CLI 侧校验逻辑可参考 src/accelerate/commands/launch.py 与 src/accelerate/commands/config/cluster.py。三、区域编译Regional Compilation解决 LLM 冷启动痛点3.1 为什么要分区域编译对整个模型做torch.compile其优化问题空间通常非常巨大导致编译冷启动时间很长。区域编译的思路是不再整体编译整个模型而是定位模型中重复的同构 block逐个依次编译以命中编译器的缓存。以GPT2LMHeadModel为例重复块/类是GPT2Block可通过model.transformer.h[0]访问模型的其余部分如model.lm_head则单独编译。这种策略能显著加快 LLM 和 Transformer 类模型的编译开销/冷启动时间该方法源自 PyTorch 官方文档中介绍的 regional compilation 配方。3.2 源码级原理compile_regions 如何工作compile_regions的实现位于 src/accelerate/utils/other.py。它通过两个辅助函数识别重复块is_repeated_blocks(module)other.py判断一个模块是否是torch.nn.ModuleList且所有子模块为同一类has_repeated_blocks(module)other.py在模块层级树的任意位置递归查找是否存在上述重复块。核心递归逻辑_compile_regions分三种情况处理命中重复块is_repeated_blocks为真新建ModuleList对每个同构子模块逐一执行torch.compile(submodule, **compile_kwargs)。由于同构 block 依次编译会命中编译器缓存后续 block 的编译大幅提速包含重复块但本身不是has_repeated_blocks为真按类重建模块外壳、绑定方法并递归处理子模块仅对包含重复块的子层下钻叶子节点直接torch.compile(module, **compile_kwargs)。编译完成后若顶层没有_orig_mod引用则会把原始模型挂到new_module.__dict__[_orig_mod]上以便后续解包decompile/unwrap。从测试 tests/test_compile.py 可以看到该行为的直接验证compiled_model compile_regions(model, modereduce-overhead, backendbackend) # 编译后的模型保留对原始模型的引用 assert hasattr(compiled_model, _orig_mod) # transformer.h[i] 与 lm_head 分别被编译 assert isinstance(compiled_model.transformer.h[0], torch._dynamo.eval_frame.OptimizedModule) assert isinstance(compiled_model.lm_head, torch._dynamo.eval_frame.OptimizedModule)此外 tests/test_compile.py 中的test_regional_compilation_cold_start直接断言了区域编译的冷启动时间小于整体编译这一核心结论。3.3 如何在 Accelerate 中开启区域编译通过TorchDynamoPlugin的use_regional_compilationTrue开启# 配置编译后端 dynamo_plugin TorchDynamoPlugin( use_regional_compilationTrue, ... # 其他参数如 backend、mode ) # 使用插件初始化 accelerator accelerator Accelerator(dynamo_plugindynamo_plugin) # 这一步会对模型应用 compile_regions model accelerator.prepare(model)在Accelerator.prepare中当use_regional_compilation为真时常规分布式路径调用compile_regions(model, **dynamo_plugin.to_kwargs())accelerator.py。值得注意的细节是TorchDynamoPlugin.to_kwargs()dataclasses.py会把use_regional_compilation从透传给torch.compile的关键字参数中剔除保证该开关只由 Accelerate 消费、不会传给 PyTorch。3.4 直接使用 compile_regions 工具函数你也可以脱离Accelerator像使用torch.compile一样直接调用accelerate.utils.compile_regionsfrom accelerate.utils import compile_regions from transformers import AutoModelForCausalLM model AutoModelForCausalLM.from_pretrained(gpt2) compiled_model compile_regions(model, modereduce-overhead) # compiled_model.transformer.h[0] 现在是 OptimizedModule # 内部 (_orig_mod) 仍是原始 GPT2Block其余部分单独编译其 docstring 示例other.py展示了编译后compiled_model.transformer.h[0]成为OptimizedModule、内部保留原始GPT2Block结构的结果。3.5 DeepSpeed 与 FSDP2 场景下的专用变体由于不同分布式封装对torch.compile的敏感度不同仓库额外提供了两个变体compile_regions_deepspeedother.py针对DeepSpeedEngine.module。因为模型被DeepSpeedEngine包裹并挂载了大量 hooks、参数 offload 等直接torch.compile(...)会干扰这些机制因此改用原地的module.compile()方法逐块编译。在 accelerator.py 中当 DeepSpeed 版本满足条件且开启了编译与区域编译时会被调用compile_regions_fsdp2other.py与compile_regions类似但同样使用原地的module.compile()。其必要性在于torch.compile(module)返回的OptimizedModule的__call__会绕过nn.Module._call_impl导致之后由fully_shard添加的 forward/pre hooks 不再触发、逐层 all-gather/reshard 失效而原地的module.compile()保留_call_impl调用路径使 FSDP2 hooks 依然生效。该行为由 tests/fsdp/test_fsdp.py 的test_compile_regions_fsdp2_preserves_block_class测试覆盖验证isinstance检查通过block 类被保留。在 FSDP2 路径中开启区域编译时还会额外设置两个torch._dynamo.config开关accelerator.pycapture_scalar_outputs True与skip_fwd_side_effects_in_bwd_under_checkpoint True以匹配 per-block 编译配方适应 MoE token-choice 分派的动态 shape并保持激活检查点AC 编译边界的一致性。四、完整训练示例将上述要素组合成一个可运行的训练片段import torch from torch.utils.data import DataLoader from accelerate import Accelerator from accelerate.utils import TorchDynamoPlugin dynamo_plugin TorchDynamoPlugin( backendinductor, modereduce-overhead, # 利用 CUDA graphs 压低调度开销 use_regional_compilationTrue, # 针对重复 block 分区域编译 ) accelerator Accelerator(dynamo_plugindynamo_plugin) model, optimizer, train_dataloader accelerator.prepare( model, optimizer, train_dataloader ) for step, batch in enumerate(train_dataloader): with accelerator.accumulate(model): outputs model(**batch) loss outputs.loss accelerator.backward(loss) optimizer.step() optimizer.zero_grad()首次进入训练循环时模型会按区域依次编译从源码路径看编译发生在prepare阶段accelerator.py而实际触发 JIT 编译则是在第一次前向推理时。若后续需要解包编译后的模型例如保存权重可使用accelerator.unwrap_model(model, keep_torch_compileFalse)其底层通过 extract_model_from_parallel 配合_orig_mod引用还原原始模型——has_compiled_regions与is_compiled_moduleother.py正是用于判断模型是否带有编译包装。五、基准测试区域编译 vs 整体编译仓库在 benchmarks/torch.compile/README.md 中给出了系统性对比实验基准脚本为 benchmarks/torch.compile/regional_compilation.py。实验设计要点如下对比三种方案Baseline不编译纯 eager 执行、Full compilationtorch.compile()整体编译、Regional compilationaccelerate.utils.compile_regions()测试模型为 1B13B 参数的 LLaMA 系列因果语言模型如 Llama-3.2-1B、3B、8B 与 13B 级模型分别在 batch size 1 与 4、序列长度 128 下测量编译时间COMPILE_ITERS 2次取中位数语句中会先重置 dynamo 并清空 inductor 缓存与推理时间INFERENCE_ITERS 100次刻意在torch.no_grad()上下文之外运行前向以模拟训练环境需要梯度下的真实表现。运行方式python regional_compilation.py脚本会自动下载模型配置、构建模型并输出基于torch.utils.benchmark的横向对比表含 colorize 高亮。关键结论性能相当区域编译的加速效果与整体编译接近对更大规模模型尤其如此编译更快区域编译显著缩短模型编译耗时对部署场景更友好batch size 的影响随着 batch size 增大两种编译策略的性能差距缩小说明编译开销在大 batch 场景下的占比更低模型规模考量模型越大区域编译在编译时间上的节省越可观实践价值对真实应用而言区域编译是优化训练冷启动时间的务实选择尤其适合大模型场景。需要说明的是上述结论来自仓库内部基准脚本的设计与文档描述具体数值会随硬件、PyTorch 版本与模型结构而变化建议在自身环境上运行 regional_compilation.py 复现验证。此外该仓库还有一个相关的区域编译示例 benchmarks/torch.compile/regional_compilation.py 及其配套图片目录 benchmarks/torch.compile/imgs可供深入分析。六、结论整体编译Full Compilation与区域编译Regional Compilation都能显著加速模型执行而区域编译在编译时间与运行时性能之间提供了更务实的平衡尤其适合大模型 较大 batch size的训练冷启动优化。实践中追求最大推理/训练吞吐可选用backendinductor, modereduce-overhead或max-autotune追求更快的首次编译与部署体验建议叠加use_regional_compilationTrue使用 DeepSpeed 或 FSDP2 时Accelerate 已内置对应的区域编译变体无需手工处理 wrapper 带来的 hook 兼容问题所有参数均可通过ACCELERATE_DYNAMO_*环境变量无侵入切换便于在 CI 或多机环境中灵活调整。赞分享人工智能深度学习分布式训练【免费下载链接】accelerate A simple way to launch, train, and use PyTorch models on almost any device and distributed configuration, automatic mixed precision (including fp8), and easy-to-configure FSDP and DeepSpeed support项目地址https://gitcode.com/gh_mirrors/ac/accelerate点击查看免费下载相关推荐CopilotKit Tool Rendering 实战用 useRenderTool 把后端工具调用渲染成聊天内的 React 卡片CopilotKit Tool Rendering 实战用 useRenderTool 把后端工具调用渲染成聊天内的 React 卡片 本篇基于 Copilo人工智能深度学习分布式训练vLLM-Omni 区域编译Regional Compilation实践指南配置、动态形状与 MiniMax-H3 固定 packed shape 优化vLLM Omni 区域编译Regional Compilation实践指南配置、动态形状与 MiniMax H3 固定 packed shape 优化人工智能大模型模型推理服务多模态语音音频媒体生成本地部署PyTorch torch.compile 编译时间优化实战从避免重编译到嵌套编译区域PyTorch torch.compile 编译时间优化实战从避免重编译到嵌套编译区域 torch.compile 是一个即时编译器JIT首次调用编译后人工智能机器学习深度学习分布式训练模型编译上一篇KLayout开源免费的终极集成电路版图设计解决方案下一篇drizzle-orm-pg 0.13.1 解析node-postgres 命名预编译语句与 .prepare(name) 的使用指南创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考