人工智能深度學(xué)習(xí)分布式訓(xùn)練【免費(fèi)下載鏈接】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點(diǎn)擊查看免費(fèi)下載本指南系統(tǒng)講解如何在 Accelerate 分布式訓(xùn)練工作流中接入 PyTorch 2.0 的torch.compile先介紹TorchDynamoPlugin的完整配置方式與底層實(shí)現(xiàn)再深入?yún)^(qū)域編譯Regional Compilation原理并結(jié)合倉庫內(nèi)基準(zhǔn)測試與源碼佐證幫助你掌握既享受編譯加速、又顯著壓低 LLM 類模型首次編譯冷啟動時間的實(shí)戰(zhàn)方案。一、為什么要在 Accelerate 中使用 torch.compilePyTorch 2.0 引入了torch.compile它通過 JIT 將 PyTorch 代碼編譯為優(yōu)化后的內(nèi)核從而顯著加速模型執(zhí)行。其核心特性可概括為三點(diǎn)性能提升Performance Improvement優(yōu)化計算圖顯著加快模型執(zhí)行速度易于使用Ease of Use對既有代碼改動極小接入成本低兼容性強(qiáng)Compatibility與現(xiàn)有 PyTorch 代碼和模型無縫協(xié)作。當(dāng)與 Accelerate 配合使用時torch.compile會平滑地融入分布式訓(xùn)練流程使你在同一套代碼中同時獲得分布式執(zhí)行與編譯優(yōu)化的收益。需要說明的是編譯后的代碼首次執(zhí)行通常更慢其中包含編譯耗時而后續(xù)運(yùn)行會明顯加快。為此torch.compile提供了多種編譯模式modedefault默認(rèn)、reduce-overhead利用 CUDA graphs 進(jìn)一步壓低調(diào)度開銷以及max-autotune大規(guī)模自動調(diào)優(yōu)為模型尋找最佳內(nèi)核。二、TorchDynamoPlugin一行配置接入編譯Accelerate 通過TorchDynamoPlugin提供對torch.compile的無縫集成。其定義位于 src/accelerate/utils/dataclasses.py繼承自KwargsHandler會在模型預(yù)處理階段把編譯參數(shù)透傳給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) # 這一步會對模型應(yīng)用 torch.compile model accelerator.prepare(model)該插件與 Accelerate 的所有其他特性與插件兼容包括混合精度、分布式訓(xùn)練DDP、FSDP、DeepSpeed等。其兼容性在源碼中得到印證Accelerator.prepare的常規(guī)路徑src/accelerate/accelerator.py、FSDP2 路徑src/accelerate/accelerator.py以及 DeepSpeed 路徑src/accelerate/accelerator.py中均按各自方式觸發(fā)編譯。2.2 完整參數(shù)說明取自源碼字段定義下表整理了 TorchDynamoPlugin 的全部可配置字段參數(shù)默認(rèn)值含義backendNone回退到環(huán)境變量再回退noDynamo 編譯后端取值見下方DynamoBackend枚舉modeNone回退環(huán)境變量再回退default編譯模式default、reduce-overhead或max-autotunefullgraphNone回退環(huán)境變量再回退False是否允許將模型拆分為多個子圖True表示要求完整圖dynamicNone僅當(dāng)設(shè)置了對應(yīng)環(huán)境變量時生效是否使用動態(tài) shape 進(jìn)行 tracingoptionsNone傳給后端的一本參數(shù)字典disableFalse將torch.compile()變成 no-op便于測試use_regional_compilationNone回退環(huán)境變量再回退False是否啟用區(qū)域編譯見本文第三節(jié)其中backend的取值對應(yīng)DynamoBackend枚舉src/accelerate/utils/dataclasses.py包括NO不使用 torch dynamoEAGER用 PyTorch 直接運(yùn)行提取出的 GraphModule適合調(diào)試 TorchDynamo 問題AOT_EAGER僅使用 AotAutograd 而無編譯器適合調(diào)試通常不會有加速INDUCTOR默認(rèn)推薦后端基于 AotAutograd、CUDA graphs 與 codegen 出的 Triton 內(nèi)核是獲取明顯加速的主流選擇AOT_TS_NVFUSER / NVPRIMS_NVFUSERnvFuser 與 AotAutograd/TorchScript 或 PrimTorch 組合CUDAGRAPHScudagraphs 配合 AotAutogradOFI / FX2TRT / ONNXRT / TENSORRT / TVM面向推理優(yōu)化的后端TensorRT、ONNX Runtime、TVM 等AOT_TORCHXLA_TRACE_ONCE / TORCHXLA_TRACE_ONCEPyTorch/XLA 的 Dynamo 優(yōu)化分別面向訓(xùn)練與推理HPU_BACKEND / NEURONHabana HPU 與 AWS NeuronTrainium/Inferentia后端。注意backendno時prepare不會觸發(fā)任何編譯DynamoBackend在內(nèi)部通過.upper()歸一化因此大小寫不敏感。2.3 環(huán)境變量配置方式TorchDynamoPlugin.__post_init__src/accelerate/utils/dataclasses.py顯示所有參數(shù)都支持通過ACCELERATE_DYNAMO_前綴的環(huán)境變量覆蓋便于在不改代碼的情況下切換配置環(huán)境變量默認(rèn)值A(chǔ)CCELERATE_DYNAMO_BACKENDnoACCELERATE_DYNAMO_MODEdefaultACCELERATE_DYNAMO_USE_FULLGRAPHFalseACCELERATE_DYNAMO_USE_DYNAMIC未設(shè)置時不啟用ACCELERATE_DYNAMO_USE_REGIONAL_COMPILATIONFalse例如export ACCELERATE_DYNAMO_BACKENDinductor export ACCELERATE_DYNAMO_MODEreduce-overhead export ACCELERATE_DYNAMO_USE_REGIONAL_COMPILATIONTrue隨后直接Accelerator()即可生效。相關(guān)配置項同樣可以在accelerate config生成的 YAML 或accelerate launch命令行參數(shù)如--dynamo_backend、--dynamo_mode中指定CLI 側(cè)校驗邏輯可參考 src/accelerate/commands/launch.py 與 src/accelerate/commands/config/cluster.py。三、區(qū)域編譯Regional Compilation解決 LLM 冷啟動痛點(diǎn)3.1 為什么要分區(qū)域編譯對整個模型做torch.compile其優(yōu)化問題空間通常非常巨大導(dǎo)致編譯冷啟動時間很長。區(qū)域編譯的思路是不再整體編譯整個模型而是定位模型中重復(fù)的同構(gòu) block逐個依次編譯以命中編譯器的緩存。以GPT2LMHeadModel為例重復(fù)塊/類是GPT2Block可通過model.transformer.h[0]訪問模型的其余部分如model.lm_head則單獨(dú)編譯。這種策略能顯著加快 LLM 和 Transformer 類模型的編譯開銷/冷啟動時間該方法源自 PyTorch 官方文檔中介紹的 regional compilation 配方。3.2 源碼級原理compile_regions 如何工作compile_regions的實(shí)現(xiàn)位于 src/accelerate/utils/other.py。它通過兩個輔助函數(shù)識別重復(fù)塊is_repeated_blocks(module)other.py判斷一個模塊是否是torch.nn.ModuleList且所有子模塊為同一類has_repeated_blocks(module)other.py在模塊層級樹的任意位置遞歸查找是否存在上述重復(fù)塊。核心遞歸邏輯_compile_regions分三種情況處理命中重復(fù)塊is_repeated_blocks為真新建ModuleList對每個同構(gòu)子模塊逐一執(zhí)行torch.compile(submodule, **compile_kwargs)。由于同構(gòu) block 依次編譯會命中編譯器緩存后續(xù) block 的編譯大幅提速包含重復(fù)塊但本身不是has_repeated_blocks為真按類重建模塊外殼、綁定方法并遞歸處理子模塊僅對包含重復(fù)塊的子層下鉆葉子節(jié)點(diǎn)直接torch.compile(module, **compile_kwargs)。編譯完成后若頂層沒有_orig_mod引用則會把原始模型掛到new_module.__dict__[_orig_mod]上以便后續(xù)解包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直接斷言了區(qū)域編譯的冷啟動時間小于整體編譯這一核心結(jié)論。3.3 如何在 Accelerate 中開啟區(qū)域編譯通過TorchDynamoPlugin的use_regional_compilationTrue開啟# 配置編譯后端 dynamo_plugin TorchDynamoPlugin( use_regional_compilationTrue, ... # 其他參數(shù)如 backend、mode ) # 使用插件初始化 accelerator accelerator Accelerator(dynamo_plugindynamo_plugin) # 這一步會對模型應(yīng)用 compile_regions model accelerator.prepare(model)在Accelerator.prepare中當(dāng)use_regional_compilation為真時常規(guī)分布式路徑調(diào)用compile_regions(model, **dynamo_plugin.to_kwargs())accelerator.py。值得注意的細(xì)節(jié)是TorchDynamoPlugin.to_kwargs()dataclasses.py會把use_regional_compilation從透傳給torch.compile的關(guān)鍵字參數(shù)中剔除保證該開關(guān)只由 Accelerate 消費(fèi)、不會傳給 PyTorch。3.4 直接使用 compile_regions 工具函數(shù)你也可以脫離Accelerator像使用torch.compile一樣直接調(diào)用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] 現(xiàn)在是 OptimizedModule # 內(nèi)部 (_orig_mod) 仍是原始 GPT2Block其余部分單獨(dú)編譯其 docstring 示例other.py展示了編譯后compiled_model.transformer.h[0]成為OptimizedModule、內(nèi)部保留原始GPT2Block結(jié)構(gòu)的結(jié)果。3.5 DeepSpeed 與 FSDP2 場景下的專用變體由于不同分布式封裝對torch.compile的敏感度不同倉庫額外提供了兩個變體compile_regions_deepspeedother.py針對DeepSpeedEngine.module。因為模型被DeepSpeedEngine包裹并掛載了大量 hooks、參數(shù) offload 等直接torch.compile(...)會干擾這些機(jī)制因此改用原地的module.compile()方法逐塊編譯。在 accelerator.py 中當(dāng) DeepSpeed 版本滿足條件且開啟了編譯與區(qū)域編譯時會被調(diào)用compile_regions_fsdp2other.py與compile_regions類似但同樣使用原地的module.compile()。其必要性在于torch.compile(module)返回的OptimizedModule的__call__會繞過nn.Module._call_impl導(dǎo)致之后由fully_shard添加的 forward/pre hooks 不再觸發(fā)、逐層 all-gather/reshard 失效而原地的module.compile()保留_call_impl調(diào)用路徑使 FSDP2 hooks 依然生效。該行為由 tests/fsdp/test_fsdp.py 的test_compile_regions_fsdp2_preserves_block_class測試覆蓋驗證isinstance檢查通過block 類被保留。在 FSDP2 路徑中開啟區(qū)域編譯時還會額外設(shè)置兩個torch._dynamo.config開關(guān)accelerator.pycapture_scalar_outputs True與skip_fwd_side_effects_in_bwd_under_checkpoint True以匹配 per-block 編譯配方適應(yīng) MoE token-choice 分派的動態(tài) shape并保持激活檢查點(diǎn)AC 編譯邊界的一致性。四、完整訓(xùn)練示例將上述要素組合成一個可運(yùn)行的訓(xùn)練片段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 壓低調(diào)度開銷 use_regional_compilationTrue, # 針對重復(fù) block 分區(qū)域編譯 ) 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()首次進(jìn)入訓(xùn)練循環(huán)時模型會按區(qū)域依次編譯從源碼路徑看編譯發(fā)生在prepare階段accelerator.py而實(shí)際觸發(fā) JIT 編譯則是在第一次前向推理時。若后續(xù)需要解包編譯后的模型例如保存權(quán)重可使用accelerator.unwrap_model(model, keep_torch_compileFalse)其底層通過 extract_model_from_parallel 配合_orig_mod引用還原原始模型——has_compiled_regions與is_compiled_moduleother.py正是用于判斷模型是否帶有編譯包裝。五、基準(zhǔn)測試區(qū)域編譯 vs 整體編譯倉庫在 benchmarks/torch.compile/README.md 中給出了系統(tǒng)性對比實(shí)驗基準(zhǔn)腳本為 benchmarks/torch.compile/regional_compilation.py。實(shí)驗設(shè)計要點(diǎn)如下對比三種方案Baseline不編譯純 eager 執(zhí)行、Full compilationtorch.compile()整體編譯、Regional compilationaccelerate.utils.compile_regions()測試模型為 1B13B 參數(shù)的 LLaMA 系列因果語言模型如 Llama-3.2-1B、3B、8B 與 13B 級模型分別在 batch size 1 與 4、序列長度 128 下測量編譯時間COMPILE_ITERS 2次取中位數(shù)語句中會先重置 dynamo 并清空 inductor 緩存與推理時間INFERENCE_ITERS 100次刻意在torch.no_grad()上下文之外運(yùn)行前向以模擬訓(xùn)練環(huán)境需要梯度下的真實(shí)表現(xiàn)。運(yùn)行方式python regional_compilation.py腳本會自動下載模型配置、構(gòu)建模型并輸出基于torch.utils.benchmark的橫向?qū)Ρ缺砗?colorize 高亮。關(guān)鍵結(jié)論性能相當(dāng)區(qū)域編譯的加速效果與整體編譯接近對更大規(guī)模模型尤其如此編譯更快區(qū)域編譯顯著縮短模型編譯耗時對部署場景更友好batch size 的影響隨著 batch size 增大兩種編譯策略的性能差距縮小說明編譯開銷在大 batch 場景下的占比更低模型規(guī)?？剂磕Ｐ驮酱髤^(qū)域編譯在編譯時間上的節(jié)省越可觀實(shí)踐價值對真實(shí)應(yīng)用而言區(qū)域編譯是優(yōu)化訓(xùn)練冷啟動時間的務(wù)實(shí)選擇尤其適合大模型場景。需要說明的是上述結(jié)論來自倉庫內(nèi)部基準(zhǔn)腳本的設(shè)計與文檔描述具體數(shù)值會隨硬件、PyTorch 版本與模型結(jié)構(gòu)而變化建議在自身環(huán)境上運(yùn)行 regional_compilation.py 復(fù)現(xiàn)驗證。此外該倉庫還有一個相關(guān)的區(qū)域編譯示例 benchmarks/torch.compile/regional_compilation.py 及其配套圖片目錄 benchmarks/torch.compile/imgs可供深入分析。六、結(jié)論整體編譯Full Compilation與區(qū)域編譯Regional Compilation都能顯著加速模型執(zhí)行而區(qū)域編譯在編譯時間與運(yùn)行時性能之間提供了更務(wù)實(shí)的平衡尤其適合大模型 較大 batch size的訓(xùn)練冷啟動優(yōu)化。實(shí)踐中追求最大推理/訓(xùn)練吞吐可選用backendinductor, modereduce-overhead或max-autotune追求更快的首次編譯與部署體驗建議疊加use_regional_compilationTrue使用 DeepSpeed 或 FSDP2 時Accelerate 已內(nèi)置對應(yīng)的區(qū)域編譯變體無需手工處理 wrapper 帶來的 hook 兼容問題所有參數(shù)均可通過ACCELERATE_DYNAMO_*環(huán)境變量無侵入切換便于在 CI 或多機(jī)環(huán)境中靈活調(diào)整。贊分享人工智能深度學(xué)習(xí)分布式訓(xùn)練【免費(fèi)下載鏈接】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點(diǎn)擊查看免費(fèi)下載相關(guān)推薦CopilotKit Tool Rendering 實(shí)戰(zhàn)用 useRenderTool 把后端工具調(diào)用渲染成聊天內(nèi)的 React 卡片CopilotKit Tool Rendering 實(shí)戰(zhàn)用 useRenderTool 把后端工具調(diào)用渲染成聊天內(nèi)的 React 卡片 本篇基于 Copilo人工智能深度學(xué)習(xí)分布式訓(xùn)練vLLM-Omni 區(qū)域編譯Regional Compilation實(shí)踐指南配置、動態(tài)形狀與 MiniMax-H3 固定 packed shape 優(yōu)化vLLM Omni 區(qū)域編譯Regional Compilation實(shí)踐指南配置、動態(tài)形狀與 MiniMax H3 固定 packed shape 優(yōu)化人工智能大模型模型推理服務(wù)多模態(tài)語音音頻媒體生成本地部署PyTorch torch.compile 編譯時間優(yōu)化實(shí)戰(zhàn)從避免重編譯到嵌套編譯區(qū)域PyTorch torch.compile 編譯時間優(yōu)化實(shí)戰(zhàn)從避免重編譯到嵌套編譯區(qū)域 torch.compile 是一個即時編譯器JIT首次調(diào)用編譯后人工智能機(jī)器學(xué)習(xí)深度學(xué)習(xí)分布式訓(xùn)練模型編譯上一篇KLayout開源免費(fèi)的終極集成電路版圖設(shè)計解決方案下一篇drizzle-orm-pg 0.13.1 解析node-postgres 命名預(yù)編譯語句與 .prepare(name) 的使用指南創(chuàng)作聲明：本文部分內(nèi)容由AI輔助生成（AIGC），僅供參考