練實戰(zhàn)：從 opt_level 到統(tǒng)一 API 的完整指南)
人工智能大模型音樂生成音頻預(yù)訓(xùn)練【免費下載鏈接】jukeboxCode for the paper Jukebox: A Generative Model for Music項目地址https://gitcode.com/gh_mirrors/ju/jukebox點擊查看免費下載本文以 NVIDIA Apex 倉庫中 amp.rst 文檔為主線系統(tǒng)講解apex.ampAutomatic Mixed Precision這一自動混合精度工具如何在只改動 3 行代碼的前提下啟用 Tensor Core 加速訓(xùn)練如何理解 O0–O3 四個優(yōu)化級別及其背后的六個核心屬性以及新舊 API 的遷移要點。讀完本文你將能夠為自己的 PyTorch 訓(xùn)練腳本正確接入amp.initialize/amp.scale_loss并針對不同模型挑選最合適的精度策略。背景Amp 解決什么問題混合精度訓(xùn)練的核心矛盾在于FP16 能利用 NVIDIA Tensor Core 成倍加速 GEMM 與卷積等算子但 FP16 的動態(tài)范圍約 5 個數(shù)量級遠小于 FP32直接以 FP16 訓(xùn)練容易導(dǎo)致梯度下溢underflow、精度損失甚至發(fā)散。手動在腳本中穿插.half()轉(zhuǎn)換不僅繁瑣還極易出錯。apex.amp的思路是讓用戶以默認的 FP32 方式構(gòu)建模型與數(shù)據(jù)由 Amp 在后臺自動完成何時轉(zhuǎn) FP16、何時保留 FP32、如何縮放損失的決策。正如倉庫 README 所描述的apex.amp是通過只修改腳本中的 3 行代碼即可啟用混合精度訓(xùn)練的工具。在 amp.rst 中Amp 給出了一個完整的最小示例# 以默認 FP32 精度聲明模型和優(yōu)化器 model torch.nn.Linear(D_in, D_out).cuda() optimizer torch.optim.SGD(model.parameters(), lr1e-3) # 讓 Amp 按 opt_level 的要求執(zhí)行類型轉(zhuǎn)換 model, optimizer amp.initialize(model, optimizer, opt_levelO1) ... # loss.backward() 改為 with amp.scale_loss(loss, optimizer) as scaled_loss: scaled_loss.backward()關(guān)鍵原則是無論選擇哪個opt_level用戶都不應(yīng)手動對模型或數(shù)據(jù)調(diào)用.half()。Amp 希望用戶從一段現(xiàn)有的 FP32 腳本出發(fā)加入上述三行代碼即可開始混合精度訓(xùn)練當(dāng)需要關(guān)閉 Amp 時enabledFalse原腳本的行為與未接入 Amp 時完全一致。因此遵循 Amp API 沒有風(fēng)險卻可能帶來可觀的性能收益。由于無需手動轉(zhuǎn)換模型與輸入數(shù)據(jù)一段遵循新 API 的腳本可以僅通過更換opt_level在多種純精度/混合精度模式間切換而不需要改動任何其他代碼。opt_level 與屬性理解 Amp 的配置模型Amp 允許用戶方便地實驗不同的純精度與混合精度模式。常用默認模式通過選擇優(yōu)化級別opt_level來指定每個opt_level定義了一組控制 Amp 實現(xiàn)純精度或混合精度訓(xùn)練的屬性properties。你也可以向amp.initialize直接傳入特定屬性的值實現(xiàn)比默認配置更細粒度的控制——手動指定的值會覆蓋opt_level建立的默認值。從源碼看這六個屬性在 frontend.py 的Properties類中統(tǒng)一管理并注冊為一個options字典默認值分別為enabledFalse、opt_levelNone、cast_model_typeNone、patch_torch_functionsFalse、keep_batchnorm_fp32None、master_weightsNone、loss_scale1.0。屬性及其語義如下屬性作用cast_model_type將模型的參數(shù)與緩沖區(qū)buffers轉(zhuǎn)換為目標類型。patch_torch_functions對所有 Torch 函數(shù)與 Tensor 方法進行打補丁patch將適合 Tensor Core 的算子如 GEMM、卷積以 FP16 執(zhí)行將受益于 FP32 精度的算子以 FP32 執(zhí)行。keep_batchnorm_fp32即便模型其余部分為 FP16也建議將 BatchNorm 權(quán)重保留在 FP32——這既提升精度又能啟用 cudnn BatchNorm改善性能。master_weights為 FP16 模型權(quán)重維護一份 FP32 主權(quán)重master weights。優(yōu)化器對 FP32 主權(quán)重執(zhí)行 step以提升精度并捕獲小梯度。loss_scale若為浮點數(shù)則作為靜態(tài)固定損失縮放值若為字符串dynamic則由 Amp 自動隨時間自適應(yīng)調(diào)整縮放值。屬性覆蓋與一致性檢查選中某個opt_level后你可以通過屬性關(guān)鍵字參數(shù)做手動覆蓋。但如果你試圖覆蓋一個對所選opt_level而言沒有意義的屬性Amp 會拋出帶解釋的錯誤。例如為opt_levelO1指定master_weightsTrue是不合理的O1是在 Torch 函數(shù)周圍插入轉(zhuǎn)換cast而不是轉(zhuǎn)換模型權(quán)重數(shù)據(jù)、激活和權(quán)重在流經(jīng)被 patch 的函數(shù)時被就地動態(tài)重新轉(zhuǎn)換out-of-place recast on the fly。因此 O1 下模型權(quán)重本身可以也應(yīng)該保持 FP32無需維護獨立的 FP32 主權(quán)重。這一檢查在源碼中有明確體現(xiàn)。frontend.py 的Properties.__setattr__對每個屬性做了守衛(wèi)例如設(shè)置master_weights時若當(dāng)前opt_level O1且值為非 None會調(diào)用warn_or_err報錯設(shè)置cast_model_type時若為 O1 且值既非False也非torch.float32同樣會警告。此外keep_batchnorm_fp32被限定為布爾值、字符串True/False或None否則觸發(fā)斷言。四個優(yōu)化級別 O0–O3 詳解Amp 識別的opt_level為O0、O1、O2、O3四種。其中O0與O3并非真正的混合精度但分別對建立精度基線與速度基線很有價值O1與O2是混合精度的兩種不同實現(xiàn)建議兩種都嘗試看哪種對具體模型帶來最好的加速與精度。注意O0/O1中的前綴是大寫字母 Othe letter O不是數(shù)字 0。源碼 frontend.py 的initialize中對非法的opt_level會拋出RuntimeError并明確提示這一點。在 frontend.py 中四個級別分別由O0、O1、O2、O3類實現(xiàn)它們的__call__方法直接對Properties對象寫入默認值注冊在opt_levels字典中供initialize調(diào)用。O0FP32 訓(xùn)練精度基線O0對你的模型而言通常是 no-op你的模型本來就是 FP32可用于建立精度基線。O0默認屬性值cast_model_typetorch.float32patch_torch_functionsFalsekeep_batchnorm_fp32None不適用一切皆為 FP32master_weightsFalseloss_scale1.0O1混合精度常規(guī)使用推薦O1對所有 Torch 函數(shù)與 Tensor 方法打補丁按白名單-黑名單whitelist-blacklist模型轉(zhuǎn)換其輸入白名單算子如適合 Tensor Core 的 GEMM、卷積以 FP16 執(zhí)行黑名單中受益于 FP32 精度的算子如 softmax以 FP32 執(zhí)行。O1默認使用動態(tài)損失縮放除非被覆蓋。O1默認屬性值cast_model_typeNone不適用patch_torch_functionsTruekeep_batchnorm_fp32None不適用所有模型權(quán)重保持 FP32master_weightsNone不適用模型權(quán)重保持 FP32loss_scaledynamicO1 的白名單/黑名單機制源碼級白名單與黑名單列表定義在 lists 目錄 下的三個文件中torch_overrides.pyFP16_FUNCS包含conv1d/2d/3d、conv_transpose*、prelu以及 BLAS 類算子addmm、addmv、addr、matmul、mm、mvFP32_FUNCS包含exp、log、mean、sum、std、var、norm、pow等CASTS需要類型提升的多 Tensor 運算包含add、div、mul、比較類算子等SEQUENCE_CASTS包含cat、stack。值得注意的是bmm/addbmm/baddbmm會根據(jù) CUDA 版本動態(tài)歸類CUDA ≥ 9.1 才具備快速 FP16 batched matmul 內(nèi)核因此版本足夠新時加入 FP16 列表否則歸入 FP32 列表。functional_overrides.pyFP16_FUNCS為torch.nn.functional下的conv1d/2d/3d、conv_transpose*、linearFP32_FUNCS包括softmax、log_softmax、softmin、softplus、layer_norm、group_norm、各類 losscross_entropy、mse_loss、nll_loss等以及interpolate。tensor_overrides.pyTensor 方法層面的對應(yīng)列表。此外amp.py 中的init()會遍歷這些列表對白名單函數(shù)用cached_cast包裝maybe_halfFP16 轉(zhuǎn)換帶緩存、對黑名單函數(shù)用cached_cast包裝maybe_float不緩存并處理CASTS/SEQUENCE_CASTS的類型提升以及對 RNN 及其 cell 的專門白名單處理。對于BANNED_FUNCS如F.binary_cross_entropyAmp 默認在遇到 FP16 輸入時報錯并給出替換建議改用BCEWithLogitsLoss或注冊amp.register_float_function(torch, sigmoid)。O2幾乎全 FP16 混合精度O2將模型權(quán)重轉(zhuǎn)換為 FP16patch 模型的forward方法以將輸入數(shù)據(jù)轉(zhuǎn)換為 FP16將 BatchNorm 保留在 FP32維護 FP32 主權(quán)重并更新優(yōu)化器的param_groups使optimizer.step()直接作用于 FP32 權(quán)重必要時隨后執(zhí)行 FP32 主權(quán)重→FP16 模型權(quán)重的拷貝同時實現(xiàn)動態(tài)損失縮放除非被覆蓋。與O1不同O2不patch Torch 函數(shù)或 Tensor 方法。O2默認屬性值cast_model_typetorch.float16patch_torch_functionsFalsekeep_batchnorm_fp32Truemaster_weightsTrueloss_scaledynamicO2 的 master weights 機制源碼級在 _process_optimizer.py 的lazy_init_with_master_weights中每個torch.cuda.HalfTensor參數(shù)都會被替換為一份 FP32 主參數(shù)param.detach().clone().float()且requires_gradTrue并用load_state_dict(state_dict())技巧將原有的 per-param 狀態(tài)張量重鑄為 FP32。optimizer.step被包裝為先執(zhí)行原old_step()再調(diào)用_master_params_to_model_params()有 C 擴展時通過multi_tensor_applier調(diào)用amp_C.multi_tensor_scale一次完成全部參數(shù)拷貝最后清空主梯度。zero_grad也被重寫為只清零模型梯度而置空主梯度。O3FP16 訓(xùn)練速度基線O3可能無法達到真正混合精度選項O1/O2的穩(wěn)定性但它可用于建立模型的速度基線供O1/O2的性能對比參考。若模型使用 BatchNorm為建立光速speed of light基線可嘗試在O3上額外覆蓋keep_batchnorm_fp32True如前所述這會啟用 cudnn BatchNorm。O3默認屬性值cast_model_typetorch.float16patch_torch_functionsFalsekeep_batchnorm_fp32Falsemaster_weightsFalseloss_scale1.0從 frontend.py 中O1、O2、O3類的brief/more文檔字符串還可以看到官方的推薦口徑O1是首次嘗試混合精度訓(xùn)練時最安全的方式O2的 FP32 主權(quán)重還可以改善收斂與穩(wěn)定性O(shè)3用于建立性能上限。實戰(zhàn)中通常的做法是先用O0確認精度基線用O3觀察速度上限然后在O1/O2中選擇精度與速度平衡最佳者。統(tǒng)一 APIinitialize、scale_loss 與 master_params新版 Amp 將所有功能收斂到apex.amp模塊下核心入口有三個amp.initialize、amp.scale_loss、amp.master_params。它們由init.py 統(tǒng)一導(dǎo)出initialize來自 frontend.pyscale_loss/disable_casts來自 handle.pymaster_params來自 _amp_state.py。amp.initialize 的完整簽名與參數(shù)amp.initialize(models, optimizersNone, enabledTrue, opt_levelO1, cast_model_typeNone, patch_torch_functionsNone, keep_batchnorm_fp32None, master_weightsNone, loss_scaleNone, cast_model_outputsNone, num_losses1, verbosity1, min_loss_scaleNone, max_loss_scale2.**24)各參數(shù)含義依據(jù) frontend.py 的 docstring 與實現(xiàn)modelstorch.nn.Module或模塊列表即待修改/轉(zhuǎn)換的模型。optimizers可選優(yōu)化器或優(yōu)化器列表訓(xùn)練時必需推理時可選。enabled默認True若為False所有 Amp 調(diào)用變?yōu)?no-op腳本如同未接入 Amp 一樣運行。opt_level默認O1接受O0/O1/O2/O3。cast_model_type/patch_torch_functions/keep_batchnorm_fp32/master_weights/loss_scale均為可選屬性覆蓋任何非None的關(guān)鍵字參數(shù)都會被解釋為手動覆蓋。cast_model_outputs默認None確保模型輸出始終轉(zhuǎn)換為指定類型無論opt_level。num_losses默認1預(yù)先告知 Amp 將使用多少個 loss/backward 次數(shù)配合amp.scale_loss的loss_id參數(shù)可為每個 loss/backward 使用不同的損失縮放值提升穩(wěn)定性。verbosity默認1設(shè)為 0 可抑制 Amp 相關(guān)輸出。min_loss_scale默認None動態(tài)損失縮放可選值的地板不用動態(tài)縮放時忽略。max_loss_scale默認2.**24動態(tài)損失縮放可選值的上限不用動態(tài)縮放時忽略。合法調(diào)用形式模型與優(yōu)化器均可為單個對象或列表返回形式與傳入形式保持一致model, optim amp.initialize(model, optim, ...) model, [optim1, optim2] amp.initialize(model, [optim1, optim2], ...) [model1, model2], optim amp.initialize([model1, model2], optim, ...) [model1, model2], [optim1, optim2] amp.initialize([model1, model2], [optim1, optim2], ...)一個很實用的特性是loss_scale與keep_batchnorm_fp32都接受數(shù)值或字符串如128.0、dynamic、True/False便于與 argparse 直接互操作——這一點在 Imagenet 示例 中體現(xiàn)得很清楚它通過--opt-level、--keep-batchnorm-fp32、--loss-scale三個命令行參數(shù)接收字符串然后原樣傳給amp.initialize。initialize 內(nèi)部做了什么frontend.py 的initialize依次完成檢查torch.backends.cudnn.enabled若被關(guān)閉則拋出RuntimeErrorAmp 依賴 cudnn校驗opt_level合法性并從opt_levels字典取回對應(yīng)默認屬性應(yīng)用用戶手動覆蓋調(diào)用 _initialize.py 的_initialize執(zhí)行真正的初始化。_initialize內(nèi)部的關(guān)鍵步驟源碼證據(jù)前置檢查check_models拒絕傳入已被torch.nn.parallel.DistributedDataParallel、apex.parallel.DistributedDataParallel或torch.nn.parallel.DataParallel包裝的模型因為并行包裝必須在amp.initialize返回之后再加如示例代碼注釋所強調(diào)如果先 DDP(model) 再 amp.initialize可能破壞 DDP 的 allreduce hookscheck_params_fp32檢查模型參數(shù)/緩沖區(qū)必須位于 CUDA 且為 FP32若為 Half 或非 CUDA 會warn_or_err提醒不需要手動 .half()check_optimizers拒絕傳入已被FP16_Optimizer包裝的優(yōu)化器。模型轉(zhuǎn)換當(dāng)cast_model_type存在時若keep_batchnorm_fp32True則調(diào)用apex.fp16_utils.convert_network跳過 BatchNorm完成轉(zhuǎn)換否則直接model.to(cast_model_type)隨后用patch_forward包裝每個模型的forward進入的輸入經(jīng)applier遞歸轉(zhuǎn)換為目標類型applier支持 Tensor、字符串、ndarray、自定義 batch 類的.to()方法、Mapping 與 Iterable 容器輸出的結(jié)果默認轉(zhuǎn)回torch.float32——這正是用戶永遠不需要調(diào)用 .half()的實現(xiàn)基礎(chǔ)。優(yōu)化器處理FusedAdam走wrap_fused_adam目前限制為opt_levelO2、keep_batchnorm_fp32False、loss_scale為數(shù)值或dynamic其余優(yōu)化器走 _process_optimizer.py 的_process_optimizer按master_weights分支安裝 lazy-init /_prepare_amp_backward/_post_amp_backward/_master_params_to_model_params等鉤子方法。損失縮放器為每個num_losses創(chuàng)建LossScaler見 scaler.py。O1 補丁當(dāng)patch_torch_functionsTrue時調(diào)用amp_init完成 Torch 命名空間的 patch見下文并把optimizer.step包裝為在disable_casts()上下文內(nèi)執(zhí)行——因為 step 只應(yīng)作用于 FP32 主參數(shù)。amp.scale_loss統(tǒng)一的反向傳播入口amp.scale_loss是一個上下文管理器context manager用法為with amp.scale_loss(loss, optimizer) as scaled_loss: scaled_loss.backward()其語義依據(jù) handle.py 的實現(xiàn)與 docstring進入上下文時創(chuàng)建scaled_loss loss.float() * 當(dāng)前損失縮放值yield 給用戶調(diào)用backward()退出上下文時若delay_unscaleFalse梯度會被檢查是否含 inf/NaN 并進行反縮放unscale隨后即可調(diào)用optimizer.step()。完整簽名amp.scale_loss(loss, optimizers, loss_id0, modelNone, delay_unscaleFalse, delay_overflow_checkFalse)。其中l(wèi)oss_id配合initialize的num_losses實現(xiàn)每個 loss 獨立的縮放器delay_unscaleTrue是小眾的忍者級性能優(yōu)化通常強烈建議保持默認False因為它會引入多模型/優(yōu)化器/loss 場景下的怪異陷阱。兩個重要警告來自 docstring若 Amp 使用顯式 FP32 主參數(shù)O2默認如此或手動master_weightsTrue任何 FP16 梯度都會先拷貝為 FP32 主梯度再反縮放optimizer.step()應(yīng)用反縮放后的主梯度到主參數(shù)上。此時只有 FP32 主梯度會被反縮放FP16 模型參數(shù)的直接.grad屬性在退出上下文后仍是縮放過的——這一細節(jié)影響梯度裁剪裁剪應(yīng)基于主參數(shù)進行。梯度溢出時的處理源碼級scaler.py 的LossScaler承擔(dān)縮放/反縮放與溢出檢測。當(dāng)檢測到溢出inf/NaN且為動態(tài)縮放時update_scale()將損失縮放減半受min_loss_scale地板約束并返回should_skipTrue此時 handle.py 的scale_loss會把optimizer.step臨時替換為skip_step打印 Gradient overflow. Skipping step... 并跳過本次更新同時清空不會由zero_grad清零的主梯度。若連續(xù)scale_window默認 2000次迭代無溢出則損失縮放加倍上限max_loss_scale默認2**24。此外LossScaler在有 C 擴展--cuda_ext --cpp_ext安裝時使用amp_C.multi_tensor_scale/multi_tensor_axpby融合內(nèi)核批量完成 unscale 與梯度拷貝通過multi_tensor_applier否則回退到 Python 實現(xiàn)并打印警告。amp.master_params遍歷主參數(shù)def master_params(optimizer): for group in optimizer.param_groups: for p in group[params]: yield p這是 _amp_state.py 中定義的一個生成器迭代amp.initialize返回的優(yōu)化器所擁有的參數(shù)。在 O2master weights模式下它產(chǎn)出的是 FP32 主參數(shù)——當(dāng)需要對這些參數(shù)做梯度裁剪時見下文高級用法應(yīng)使用master_params(optimizer)而不是model.parameters()。高級用法新版統(tǒng)一 Amp API 支持跨迭代的梯度累積、單迭代多次反向傳播、多模型/多優(yōu)化器、自定義用戶定義autograd 函數(shù)以及自定義數(shù)據(jù) batch 類。梯度裁剪與 GAN 也需要特殊處理但這些處理方式不隨opt_level改變。文檔中對應(yīng)的高級主題詳見倉庫內(nèi)的 advanced.rst。梯度裁剪基于主參數(shù)由于 O2 模式下 FP16 模型參數(shù)的.grad仍是縮放過的正確做法是使用amp.master_params(optimizer)拿到 FP32 主參數(shù)再對主參數(shù)做裁剪# 推薦對主參數(shù)進行梯度裁剪 torch.nn.utils.clip_grad_norm_(amp.master_params(optimizer), max_norm)梯度累積與多次 backwardamp.scale_loss支持跨迭代梯度累積與單次迭代多次反向傳播。若使用delay_unscaleTrue則本次 backward 退出后還不能調(diào)用optimizer.step()必須等后續(xù)某次 backward 以delay_unscaleFalse默認結(jié)束后再 step。多模型 / 多優(yōu)化器 / 多 lossamp.initialize接受模型列表與優(yōu)化器列表amp.scale_loss的optimizers參數(shù)可傳單個優(yōu)化器或優(yōu)化器列表。當(dāng)存在多個 loss 時通過num_lossesN與loss_idi組合Amp 為每個 loss 維護獨立的損失縮放器model1, model2 ... optim1, optim2 ... [model1, model2], [optim1, optim2] amp.initialize( [model1, model2], [optim1, optim2], opt_levelO1, num_losses2) with amp.scale_loss(loss1, [optim1, optim2], loss_id0) as scaled_loss1: scaled_loss1.backward() with amp.scale_loss(loss2, [optim1, optim2], loss_id1) as scaled_loss2: scaled_loss2.backward() optimizer1.step(); optimizer2.step()即便num_losses保持默認 1Amp 也支持多個 loss/backward只是所有 loss 共享同一個全局損失縮放值。GAN 與自定義 batch 類GAN生成器與判別器可分別作為獨立模型傳入amp.initialize當(dāng)需要生成器不更新只做判別之類的凍結(jié)場景時_initialize.py的to_type對requires_grad的輸入 Tensor 有專門注釋GANs require this說明其對輸入數(shù)據(jù)的處理已考慮 GAN 場景。自定義數(shù)據(jù) batch 類_initialize.py的applier函數(shù)會檢查對象是否有.to()方法并調(diào)用它來轉(zhuǎn)換 batch 內(nèi)的浮點 Tensor——因此自定義 batch 類只需實現(xiàn)to(dtype)方法即可被patch_forward正確轉(zhuǎn)換。從舊 API 遷移Transition GuideAmp 強烈建議遷移到新 API因為它更通用、更易用、面向未來原有的FP16_Optimizer類與舊版 Amp API 已被棄用deprecated隨時可能被移除。舊 Amp API 用戶在新 API 中opt_levelO1執(zhí)行與舊 Amp 相同的 Torch 命名空間 patch。區(qū)別在于新 API 支持靜態(tài)或動態(tài)損失縮放而舊 API 只支持動態(tài)損失縮放。遷移要點刪除舊的amp_handle amp.init()調(diào)用及返回的amp_handle——amp.initialize()已承擔(dān)并超越了amp.init()的職責(zé)原來通過amp_handle暴露的函數(shù)現(xiàn)在是amp模塊下的自由函數(shù)反向傳播上下文管理器需相應(yīng)修改# 舊 API with amp_handle.scale_loss(loss, optimizer) as scaled_loss: scaled_loss.backward() # - # 新 API with amp.scale_loss(loss, optimizer) as scaled_loss: scaled_loss.backward()值得注意的是舊 API 中為用戶函數(shù)標注特定精度的調(diào)用如amp.half_function、amp.float_function、amp.promote_function以及注冊形式的amp.register_half_function/register_float_function/register_promote_function實現(xiàn)見 amp.py在新 API 中仍然被支持。舊 FP16_Optimizer 用戶opt_levelO2等價于FP16_Optimizer配合dynamic_loss_scaleTrue。反向傳播同樣需要改為統(tǒng)一版本optimizer.backward(loss) # - with amp.scale_loss(loss, optimizer) as scaled_loss: scaled_loss.backward()舊FP16_Optimizer的一個痛點在于用戶必須手動將模型轉(zhuǎn)為 half調(diào)用.half()或使用apex.fp16_utils的包裝函數(shù)還必須手動對輸入數(shù)據(jù)調(diào)用.half()。新 API 中這兩者都不再需要——無論選擇哪個--opt-level你都應(yīng)該用默認 FP32 格式構(gòu)建模型和傳遞輸入數(shù)據(jù)Amp 會在model, optimizer amp.initialize(...)期間根據(jù)opt_level及覆蓋標志自動完成正確的轉(zhuǎn)換。浮點輸入數(shù)據(jù)可以是 FP32 或 FP16但更簡單的做法是直接保持 FP32——amp.initialize返回的模型其forward已被 patch會自動將輸入數(shù)據(jù)轉(zhuǎn)換到合適類型。一個完整的實踐示例結(jié)合倉庫中的 Imagenet 示例一個接入 Amp 的典型訓(xùn)練腳本骨架如下省略數(shù)據(jù)加載等細節(jié)import torch from apex import amp # 1. 以 FP32 構(gòu)建模型與優(yōu)化器 model model.cuda() optimizer torch.optim.SGD(model.parameters(), lrargs.lr, momentumargs.momentum, weight_decayargs.weight_decay) # 2. 初始化 Amp先于任何 DDP 包裝 model, optimizer amp.initialize(model, optimizer, opt_levelargs.opt_level, # 如 O1 keep_batchnorm_fp32args.keep_batchnorm_fp32, loss_scaleargs.loss_scale) # 如 dynamic # 3. 分布式場景包裝 DDP 必須在 amp.initialize 之后 if args.distributed: model DDP(model, delay_allreduceTrue) # 4. 訓(xùn)練循環(huán) for epoch in range(args.start_epoch, args.epochs): for i, (input, target) in enumerate(train_loader): output model(input) loss criterion(output, target) with amp.scale_loss(loss, optimizer) as scaled_loss: scaled_loss.backward() optimizer.step() optimizer.zero_grad()命令行調(diào)用示例對應(yīng)示例倉庫中由 argparse 接收的參數(shù)python main_amp.py --arch resnet50 --opt-level O1 --loss-scale dynamic python main_amp.py --arch resnet50 --opt-level O2 --keep-batchnorm-fp32 TrueAmp 的安裝方式參見倉庫 README推薦帶 CUDA/C 擴展安裝pip install -v --no-cache-dir --global-option--cpp_ext --global-option--cuda_ext .這樣可啟用multi_tensor_apply融合內(nèi)核unscale、梯度拷貝、master weight 拷貝均受益純 Python 安裝PyTorch 0.4 時代的需求下amp仍然可用但會退化為 Python 回退實現(xiàn)并可能更慢。小結(jié)apex.amp的設(shè)計哲學(xué)是讓用戶始終以 FP32 視角編寫訓(xùn)練腳本把精度策略交給 opt_level。本文梳理了文檔 amp.rst 的核心內(nèi)容六個底層屬性cast_model_type、patch_torch_functions、keep_batchnorm_fp32、master_weights、loss_scale構(gòu)成了精度策略的原子單元O0–O3 四個級別提供了開箱即用的默認組合amp.initialize/amp.scale_loss/amp.master_params三個入口覆蓋了訓(xùn)練的全部關(guān)鍵環(huán)節(jié)而舊 API 的遷移路徑清晰且無損舊式函數(shù)標注仍然有效。結(jié)合 frontend.py、_initialize.py、_process_optimizer.py、scaler.py 與 handle.py 等源碼可以看到每個 API 承諾背后都有完整的實現(xiàn)支撐。初次嘗試時從O1最安全出發(fā)用O0/O3建立精度與速度基線再探索O2的 master weights 收益是一個穩(wěn)妥的實踐路徑。贊分享人工智能大模型音樂生成音頻預(yù)訓(xùn)練【免費下載鏈接】jukeboxCode for the paper Jukebox: A Generative Model for Music項目地址https://gitcode.com/gh_mirrors/ju/jukebox點擊查看免費下載相關(guān)推薦Ignite混合精度訓(xùn)練實戰(zhàn)AMP與Apex的完整對比Ignite混合精度訓(xùn)練實戰(zhàn)AMP與Apex的完整對比 PyTorch Ignite是一個高級庫用于幫助在PyTorch中靈活透明地訓(xùn)練和評估神經(jīng)網(wǎng)絡(luò)。在深深度學(xué)習(xí)模型訓(xùn)練使用 Apex AMP 混合精度訓(xùn)練 DCGANexamples/dcgan/main_amp.py 實戰(zhàn)指南使用 Apex AMP 混合精度訓(xùn)練 DCGAN examples/dcgan/main_amp.py 實戰(zhàn)指南 本指南圍繞 Apex 倉庫中 example人工智能深度學(xué)習(xí)分布式訓(xùn)練模型優(yōu)化Colossal-AI 混合精度訓(xùn)練完全指南基于 Booster 的 AMP 配置與實戰(zhàn)Torch/Apex/Naive AMPColossal AI 混合精度訓(xùn)練完全指南基于 Booster 的 AMP 配置與實戰(zhàn)Torch/Apex/Naive AMP 導(dǎo)讀 本文以 Colos人工智能大模型分布式訓(xùn)練深度學(xué)習(xí)模型優(yōu)化高性能計算上一篇PageMenu菜單項居中布局打造優(yōu)雅視覺平衡的終極指南下一篇嵌入式數(shù)據(jù)庫sled高級應(yīng)用合并操作符與觸發(fā)器實現(xiàn)創(chuàng)作聲明：本文部分內(nèi)容由AI輔助生成（AIGC），僅供參考 Apex Amp 混合精度訓(xùn)練實戰(zhàn)：從 opt_level 到統(tǒng)一 API 的完整指南
人工智能大模型音樂生成音頻預(yù)訓(xùn)練【免費下載鏈接】jukeboxCode for the paper Jukebox: A Generative Model for Music項目地址https://gitcode.com/gh_mirrors/ju/jukebox點擊查看免費下載本文以 NVIDIA Apex 倉庫中 amp.rst 文檔為主線系統(tǒng)講解apex.ampAutomatic Mixed Precision這一自動混合精度工具如何在只改動 3 行代碼的前提下啟用 Tensor Core 加速訓(xùn)練如何理解 O0–O3 四個優(yōu)化級別及其背后的六個核心屬性以及新舊 API 的遷移要點。讀完本文你將能夠為自己的 PyTorch 訓(xùn)練腳本正確接入amp.initialize/amp.scale_loss并針對不同模型挑選最合適的精度策略。背景Amp 解決什么問題混合精度訓(xùn)練的核心矛盾在于FP16 能利用 NVIDIA Tensor Core 成倍加速 GEMM 與卷積等算子但 FP16 的動態(tài)范圍約 5 個數(shù)量級遠小于 FP32直接以 FP16 訓(xùn)練容易導(dǎo)致梯度下溢underflow、精度損失甚至發(fā)散。手動在腳本中穿插.half()轉(zhuǎn)換不僅繁瑣還極易出錯。apex.amp的思路是讓用戶以默認的 FP32 方式構(gòu)建模型與數(shù)據(jù)由 Amp 在后臺自動完成何時轉(zhuǎn) FP16、何時保留 FP32、如何縮放損失的決策。正如倉庫 README 所描述的apex.amp是通過只修改腳本中的 3 行代碼即可啟用混合精度訓(xùn)練的工具。在 amp.rst 中Amp 給出了一個完整的最小示例# 以默認 FP32 精度聲明模型和優(yōu)化器 model torch.nn.Linear(D_in, D_out).cuda() optimizer torch.optim.SGD(model.parameters(), lr1e-3) # 讓 Amp 按 opt_level 的要求執(zhí)行類型轉(zhuǎn)換 model, optimizer amp.initialize(model, optimizer, opt_levelO1) ... # loss.backward() 改為 with amp.scale_loss(loss, optimizer) as scaled_loss: scaled_loss.backward()關(guān)鍵原則是無論選擇哪個opt_level用戶都不應(yīng)手動對模型或數(shù)據(jù)調(diào)用.half()。Amp 希望用戶從一段現(xiàn)有的 FP32 腳本出發(fā)加入上述三行代碼即可開始混合精度訓(xùn)練當(dāng)需要關(guān)閉 Amp 時enabledFalse原腳本的行為與未接入 Amp 時完全一致。因此遵循 Amp API 沒有風(fēng)險卻可能帶來可觀的性能收益。由于無需手動轉(zhuǎn)換模型與輸入數(shù)據(jù)一段遵循新 API 的腳本可以僅通過更換opt_level在多種純精度/混合精度模式間切換而不需要改動任何其他代碼。opt_level 與屬性理解 Amp 的配置模型Amp 允許用戶方便地實驗不同的純精度與混合精度模式。常用默認模式通過選擇優(yōu)化級別opt_level來指定每個opt_level定義了一組控制 Amp 實現(xiàn)純精度或混合精度訓(xùn)練的屬性properties。你也可以向amp.initialize直接傳入特定屬性的值實現(xiàn)比默認配置更細粒度的控制——手動指定的值會覆蓋opt_level建立的默認值。從源碼看這六個屬性在 frontend.py 的Properties類中統(tǒng)一管理并注冊為一個options字典默認值分別為enabledFalse、opt_levelNone、cast_model_typeNone、patch_torch_functionsFalse、keep_batchnorm_fp32None、master_weightsNone、loss_scale1.0。屬性及其語義如下屬性作用cast_model_type將模型的參數(shù)與緩沖區(qū)buffers轉(zhuǎn)換為目標類型。patch_torch_functions對所有 Torch 函數(shù)與 Tensor 方法進行打補丁patch將適合 Tensor Core 的算子如 GEMM、卷積以 FP16 執(zhí)行將受益于 FP32 精度的算子以 FP32 執(zhí)行。keep_batchnorm_fp32即便模型其余部分為 FP16也建議將 BatchNorm 權(quán)重保留在 FP32——這既提升精度又能啟用 cudnn BatchNorm改善性能。master_weights為 FP16 模型權(quán)重維護一份 FP32 主權(quán)重master weights。優(yōu)化器對 FP32 主權(quán)重執(zhí)行 step以提升精度并捕獲小梯度。loss_scale若為浮點數(shù)則作為靜態(tài)固定損失縮放值若為字符串dynamic則由 Amp 自動隨時間自適應(yīng)調(diào)整縮放值。屬性覆蓋與一致性檢查選中某個opt_level后你可以通過屬性關(guān)鍵字參數(shù)做手動覆蓋。但如果你試圖覆蓋一個對所選opt_level而言沒有意義的屬性Amp 會拋出帶解釋的錯誤。例如為opt_levelO1指定master_weightsTrue是不合理的O1是在 Torch 函數(shù)周圍插入轉(zhuǎn)換cast而不是轉(zhuǎn)換模型權(quán)重數(shù)據(jù)、激活和權(quán)重在流經(jīng)被 patch 的函數(shù)時被就地動態(tài)重新轉(zhuǎn)換out-of-place recast on the fly。因此 O1 下模型權(quán)重本身可以也應(yīng)該保持 FP32無需維護獨立的 FP32 主權(quán)重。這一檢查在源碼中有明確體現(xiàn)。frontend.py 的Properties.__setattr__對每個屬性做了守衛(wèi)例如設(shè)置master_weights時若當(dāng)前opt_level O1且值為非 None會調(diào)用warn_or_err報錯設(shè)置cast_model_type時若為 O1 且值既非False也非torch.float32同樣會警告。此外keep_batchnorm_fp32被限定為布爾值、字符串True/False或None否則觸發(fā)斷言。四個優(yōu)化級別 O0–O3 詳解Amp 識別的opt_level為O0、O1、O2、O3四種。其中O0與O3并非真正的混合精度但分別對建立精度基線與速度基線很有價值O1與O2是混合精度的兩種不同實現(xiàn)建議兩種都嘗試看哪種對具體模型帶來最好的加速與精度。注意O0/O1中的前綴是大寫字母 Othe letter O不是數(shù)字 0。源碼 frontend.py 的initialize中對非法的opt_level會拋出RuntimeError并明確提示這一點。在 frontend.py 中四個級別分別由O0、O1、O2、O3類實現(xiàn)它們的__call__方法直接對Properties對象寫入默認值注冊在opt_levels字典中供initialize調(diào)用。O0FP32 訓(xùn)練精度基線O0對你的模型而言通常是 no-op你的模型本來就是 FP32可用于建立精度基線。O0默認屬性值cast_model_typetorch.float32patch_torch_functionsFalsekeep_batchnorm_fp32None不適用一切皆為 FP32master_weightsFalseloss_scale1.0O1混合精度常規(guī)使用推薦O1對所有 Torch 函數(shù)與 Tensor 方法打補丁按白名單-黑名單whitelist-blacklist模型轉(zhuǎn)換其輸入白名單算子如適合 Tensor Core 的 GEMM、卷積以 FP16 執(zhí)行黑名單中受益于 FP32 精度的算子如 softmax以 FP32 執(zhí)行。O1默認使用動態(tài)損失縮放除非被覆蓋。O1默認屬性值cast_model_typeNone不適用patch_torch_functionsTruekeep_batchnorm_fp32None不適用所有模型權(quán)重保持 FP32master_weightsNone不適用模型權(quán)重保持 FP32loss_scaledynamicO1 的白名單/黑名單機制源碼級白名單與黑名單列表定義在 lists 目錄 下的三個文件中torch_overrides.pyFP16_FUNCS包含conv1d/2d/3d、conv_transpose*、prelu以及 BLAS 類算子addmm、addmv、addr、matmul、mm、mvFP32_FUNCS包含exp、log、mean、sum、std、var、norm、pow等CASTS需要類型提升的多 Tensor 運算包含add、div、mul、比較類算子等SEQUENCE_CASTS包含cat、stack。值得注意的是bmm/addbmm/baddbmm會根據(jù) CUDA 版本動態(tài)歸類CUDA ≥ 9.1 才具備快速 FP16 batched matmul 內(nèi)核因此版本足夠新時加入 FP16 列表否則歸入 FP32 列表。functional_overrides.pyFP16_FUNCS為torch.nn.functional下的conv1d/2d/3d、conv_transpose*、linearFP32_FUNCS包括softmax、log_softmax、softmin、softplus、layer_norm、group_norm、各類 losscross_entropy、mse_loss、nll_loss等以及interpolate。tensor_overrides.pyTensor 方法層面的對應(yīng)列表。此外amp.py 中的init()會遍歷這些列表對白名單函數(shù)用cached_cast包裝maybe_halfFP16 轉(zhuǎn)換帶緩存、對黑名單函數(shù)用cached_cast包裝maybe_float不緩存并處理CASTS/SEQUENCE_CASTS的類型提升以及對 RNN 及其 cell 的專門白名單處理。對于BANNED_FUNCS如F.binary_cross_entropyAmp 默認在遇到 FP16 輸入時報錯并給出替換建議改用BCEWithLogitsLoss或注冊amp.register_float_function(torch, sigmoid)。O2幾乎全 FP16 混合精度O2將模型權(quán)重轉(zhuǎn)換為 FP16patch 模型的forward方法以將輸入數(shù)據(jù)轉(zhuǎn)換為 FP16將 BatchNorm 保留在 FP32維護 FP32 主權(quán)重并更新優(yōu)化器的param_groups使optimizer.step()直接作用于 FP32 權(quán)重必要時隨后執(zhí)行 FP32 主權(quán)重→FP16 模型權(quán)重的拷貝同時實現(xiàn)動態(tài)損失縮放除非被覆蓋。與O1不同O2不patch Torch 函數(shù)或 Tensor 方法。O2默認屬性值cast_model_typetorch.float16patch_torch_functionsFalsekeep_batchnorm_fp32Truemaster_weightsTrueloss_scaledynamicO2 的 master weights 機制源碼級在 _process_optimizer.py 的lazy_init_with_master_weights中每個torch.cuda.HalfTensor參數(shù)都會被替換為一份 FP32 主參數(shù)param.detach().clone().float()且requires_gradTrue并用load_state_dict(state_dict())技巧將原有的 per-param 狀態(tài)張量重鑄為 FP32。optimizer.step被包裝為先執(zhí)行原old_step()再調(diào)用_master_params_to_model_params()有 C 擴展時通過multi_tensor_applier調(diào)用amp_C.multi_tensor_scale一次完成全部參數(shù)拷貝最后清空主梯度。zero_grad也被重寫為只清零模型梯度而置空主梯度。O3FP16 訓(xùn)練速度基線O3可能無法達到真正混合精度選項O1/O2的穩(wěn)定性但它可用于建立模型的速度基線供O1/O2的性能對比參考。若模型使用 BatchNorm為建立光速speed of light基線可嘗試在O3上額外覆蓋keep_batchnorm_fp32True如前所述這會啟用 cudnn BatchNorm。O3默認屬性值cast_model_typetorch.float16patch_torch_functionsFalsekeep_batchnorm_fp32Falsemaster_weightsFalseloss_scale1.0從 frontend.py 中O1、O2、O3類的brief/more文檔字符串還可以看到官方的推薦口徑O1是首次嘗試混合精度訓(xùn)練時最安全的方式O2的 FP32 主權(quán)重還可以改善收斂與穩(wěn)定性O(shè)3用于建立性能上限。實戰(zhàn)中通常的做法是先用O0確認精度基線用O3觀察速度上限然后在O1/O2中選擇精度與速度平衡最佳者。統(tǒng)一 APIinitialize、scale_loss 與 master_params新版 Amp 將所有功能收斂到apex.amp模塊下核心入口有三個amp.initialize、amp.scale_loss、amp.master_params。它們由init.py 統(tǒng)一導(dǎo)出initialize來自 frontend.pyscale_loss/disable_casts來自 handle.pymaster_params來自 _amp_state.py。amp.initialize 的完整簽名與參數(shù)amp.initialize(models, optimizersNone, enabledTrue, opt_levelO1, cast_model_typeNone, patch_torch_functionsNone, keep_batchnorm_fp32None, master_weightsNone, loss_scaleNone, cast_model_outputsNone, num_losses1, verbosity1, min_loss_scaleNone, max_loss_scale2.**24)各參數(shù)含義依據(jù) frontend.py 的 docstring 與實現(xiàn)modelstorch.nn.Module或模塊列表即待修改/轉(zhuǎn)換的模型。optimizers可選優(yōu)化器或優(yōu)化器列表訓(xùn)練時必需推理時可選。enabled默認True若為False所有 Amp 調(diào)用變?yōu)?no-op腳本如同未接入 Amp 一樣運行。opt_level默認O1接受O0/O1/O2/O3。cast_model_type/patch_torch_functions/keep_batchnorm_fp32/master_weights/loss_scale均為可選屬性覆蓋任何非None的關(guān)鍵字參數(shù)都會被解釋為手動覆蓋。cast_model_outputs默認None確保模型輸出始終轉(zhuǎn)換為指定類型無論opt_level。num_losses默認1預(yù)先告知 Amp 將使用多少個 loss/backward 次數(shù)配合amp.scale_loss的loss_id參數(shù)可為每個 loss/backward 使用不同的損失縮放值提升穩(wěn)定性。verbosity默認1設(shè)為 0 可抑制 Amp 相關(guān)輸出。min_loss_scale默認None動態(tài)損失縮放可選值的地板不用動態(tài)縮放時忽略。max_loss_scale默認2.**24動態(tài)損失縮放可選值的上限不用動態(tài)縮放時忽略。合法調(diào)用形式模型與優(yōu)化器均可為單個對象或列表返回形式與傳入形式保持一致model, optim amp.initialize(model, optim, ...) model, [optim1, optim2] amp.initialize(model, [optim1, optim2], ...) [model1, model2], optim amp.initialize([model1, model2], optim, ...) [model1, model2], [optim1, optim2] amp.initialize([model1, model2], [optim1, optim2], ...)一個很實用的特性是loss_scale與keep_batchnorm_fp32都接受數(shù)值或字符串如128.0、dynamic、True/False便于與 argparse 直接互操作——這一點在 Imagenet 示例 中體現(xiàn)得很清楚它通過--opt-level、--keep-batchnorm-fp32、--loss-scale三個命令行參數(shù)接收字符串然后原樣傳給amp.initialize。initialize 內(nèi)部做了什么frontend.py 的initialize依次完成檢查torch.backends.cudnn.enabled若被關(guān)閉則拋出RuntimeErrorAmp 依賴 cudnn校驗opt_level合法性并從opt_levels字典取回對應(yīng)默認屬性應(yīng)用用戶手動覆蓋調(diào)用 _initialize.py 的_initialize執(zhí)行真正的初始化。_initialize內(nèi)部的關(guān)鍵步驟源碼證據(jù)前置檢查check_models拒絕傳入已被torch.nn.parallel.DistributedDataParallel、apex.parallel.DistributedDataParallel或torch.nn.parallel.DataParallel包裝的模型因為并行包裝必須在amp.initialize返回之后再加如示例代碼注釋所強調(diào)如果先 DDP(model) 再 amp.initialize可能破壞 DDP 的 allreduce hookscheck_params_fp32檢查模型參數(shù)/緩沖區(qū)必須位于 CUDA 且為 FP32若為 Half 或非 CUDA 會warn_or_err提醒不需要手動 .half()check_optimizers拒絕傳入已被FP16_Optimizer包裝的優(yōu)化器。模型轉(zhuǎn)換當(dāng)cast_model_type存在時若keep_batchnorm_fp32True則調(diào)用apex.fp16_utils.convert_network跳過 BatchNorm完成轉(zhuǎn)換否則直接model.to(cast_model_type)隨后用patch_forward包裝每個模型的forward進入的輸入經(jīng)applier遞歸轉(zhuǎn)換為目標類型applier支持 Tensor、字符串、ndarray、自定義 batch 類的.to()方法、Mapping 與 Iterable 容器輸出的結(jié)果默認轉(zhuǎn)回torch.float32——這正是用戶永遠不需要調(diào)用 .half()的實現(xiàn)基礎(chǔ)。優(yōu)化器處理FusedAdam走wrap_fused_adam目前限制為opt_levelO2、keep_batchnorm_fp32False、loss_scale為數(shù)值或dynamic其余優(yōu)化器走 _process_optimizer.py 的_process_optimizer按master_weights分支安裝 lazy-init /_prepare_amp_backward/_post_amp_backward/_master_params_to_model_params等鉤子方法。損失縮放器為每個num_losses創(chuàng)建LossScaler見 scaler.py。O1 補丁當(dāng)patch_torch_functionsTrue時調(diào)用amp_init完成 Torch 命名空間的 patch見下文并把optimizer.step包裝為在disable_casts()上下文內(nèi)執(zhí)行——因為 step 只應(yīng)作用于 FP32 主參數(shù)。amp.scale_loss統(tǒng)一的反向傳播入口amp.scale_loss是一個上下文管理器context manager用法為with amp.scale_loss(loss, optimizer) as scaled_loss: scaled_loss.backward()其語義依據(jù) handle.py 的實現(xiàn)與 docstring進入上下文時創(chuàng)建scaled_loss loss.float() * 當(dāng)前損失縮放值yield 給用戶調(diào)用backward()退出上下文時若delay_unscaleFalse梯度會被檢查是否含 inf/NaN 并進行反縮放unscale隨后即可調(diào)用optimizer.step()。完整簽名amp.scale_loss(loss, optimizers, loss_id0, modelNone, delay_unscaleFalse, delay_overflow_checkFalse)。其中l(wèi)oss_id配合initialize的num_losses實現(xiàn)每個 loss 獨立的縮放器delay_unscaleTrue是小眾的忍者級性能優(yōu)化通常強烈建議保持默認False因為它會引入多模型/優(yōu)化器/loss 場景下的怪異陷阱。兩個重要警告來自 docstring若 Amp 使用顯式 FP32 主參數(shù)O2默認如此或手動master_weightsTrue任何 FP16 梯度都會先拷貝為 FP32 主梯度再反縮放optimizer.step()應(yīng)用反縮放后的主梯度到主參數(shù)上。此時只有 FP32 主梯度會被反縮放FP16 模型參數(shù)的直接.grad屬性在退出上下文后仍是縮放過的——這一細節(jié)影響梯度裁剪裁剪應(yīng)基于主參數(shù)進行。梯度溢出時的處理源碼級scaler.py 的LossScaler承擔(dān)縮放/反縮放與溢出檢測。當(dāng)檢測到溢出inf/NaN且為動態(tài)縮放時update_scale()將損失縮放減半受min_loss_scale地板約束并返回should_skipTrue此時 handle.py 的scale_loss會把optimizer.step臨時替換為skip_step打印 Gradient overflow. Skipping step... 并跳過本次更新同時清空不會由zero_grad清零的主梯度。若連續(xù)scale_window默認 2000次迭代無溢出則損失縮放加倍上限max_loss_scale默認2**24。此外LossScaler在有 C 擴展--cuda_ext --cpp_ext安裝時使用amp_C.multi_tensor_scale/multi_tensor_axpby融合內(nèi)核批量完成 unscale 與梯度拷貝通過multi_tensor_applier否則回退到 Python 實現(xiàn)并打印警告。amp.master_params遍歷主參數(shù)def master_params(optimizer): for group in optimizer.param_groups: for p in group[params]: yield p這是 _amp_state.py 中定義的一個生成器迭代amp.initialize返回的優(yōu)化器所擁有的參數(shù)。在 O2master weights模式下它產(chǎn)出的是 FP32 主參數(shù)——當(dāng)需要對這些參數(shù)做梯度裁剪時見下文高級用法應(yīng)使用master_params(optimizer)而不是model.parameters()。高級用法新版統(tǒng)一 Amp API 支持跨迭代的梯度累積、單迭代多次反向傳播、多模型/多優(yōu)化器、自定義用戶定義autograd 函數(shù)以及自定義數(shù)據(jù) batch 類。梯度裁剪與 GAN 也需要特殊處理但這些處理方式不隨opt_level改變。文檔中對應(yīng)的高級主題詳見倉庫內(nèi)的 advanced.rst。梯度裁剪基于主參數(shù)由于 O2 模式下 FP16 模型參數(shù)的.grad仍是縮放過的正確做法是使用amp.master_params(optimizer)拿到 FP32 主參數(shù)再對主參數(shù)做裁剪# 推薦對主參數(shù)進行梯度裁剪 torch.nn.utils.clip_grad_norm_(amp.master_params(optimizer), max_norm)梯度累積與多次 backwardamp.scale_loss支持跨迭代梯度累積與單次迭代多次反向傳播。若使用delay_unscaleTrue則本次 backward 退出后還不能調(diào)用optimizer.step()必須等后續(xù)某次 backward 以delay_unscaleFalse默認結(jié)束后再 step。多模型 / 多優(yōu)化器 / 多 lossamp.initialize接受模型列表與優(yōu)化器列表amp.scale_loss的optimizers參數(shù)可傳單個優(yōu)化器或優(yōu)化器列表。當(dāng)存在多個 loss 時通過num_lossesN與loss_idi組合Amp 為每個 loss 維護獨立的損失縮放器model1, model2 ... optim1, optim2 ... [model1, model2], [optim1, optim2] amp.initialize( [model1, model2], [optim1, optim2], opt_levelO1, num_losses2) with amp.scale_loss(loss1, [optim1, optim2], loss_id0) as scaled_loss1: scaled_loss1.backward() with amp.scale_loss(loss2, [optim1, optim2], loss_id1) as scaled_loss2: scaled_loss2.backward() optimizer1.step(); optimizer2.step()即便num_losses保持默認 1Amp 也支持多個 loss/backward只是所有 loss 共享同一個全局損失縮放值。GAN 與自定義 batch 類GAN生成器與判別器可分別作為獨立模型傳入amp.initialize當(dāng)需要生成器不更新只做判別之類的凍結(jié)場景時_initialize.py的to_type對requires_grad的輸入 Tensor 有專門注釋GANs require this說明其對輸入數(shù)據(jù)的處理已考慮 GAN 場景。自定義數(shù)據(jù) batch 類_initialize.py的applier函數(shù)會檢查對象是否有.to()方法并調(diào)用它來轉(zhuǎn)換 batch 內(nèi)的浮點 Tensor——因此自定義 batch 類只需實現(xiàn)to(dtype)方法即可被patch_forward正確轉(zhuǎn)換。從舊 API 遷移Transition GuideAmp 強烈建議遷移到新 API因為它更通用、更易用、面向未來原有的FP16_Optimizer類與舊版 Amp API 已被棄用deprecated隨時可能被移除。舊 Amp API 用戶在新 API 中opt_levelO1執(zhí)行與舊 Amp 相同的 Torch 命名空間 patch。區(qū)別在于新 API 支持靜態(tài)或動態(tài)損失縮放而舊 API 只支持動態(tài)損失縮放。遷移要點刪除舊的amp_handle amp.init()調(diào)用及返回的amp_handle——amp.initialize()已承擔(dān)并超越了amp.init()的職責(zé)原來通過amp_handle暴露的函數(shù)現(xiàn)在是amp模塊下的自由函數(shù)反向傳播上下文管理器需相應(yīng)修改# 舊 API with amp_handle.scale_loss(loss, optimizer) as scaled_loss: scaled_loss.backward() # - # 新 API with amp.scale_loss(loss, optimizer) as scaled_loss: scaled_loss.backward()值得注意的是舊 API 中為用戶函數(shù)標注特定精度的調(diào)用如amp.half_function、amp.float_function、amp.promote_function以及注冊形式的amp.register_half_function/register_float_function/register_promote_function實現(xiàn)見 amp.py在新 API 中仍然被支持。舊 FP16_Optimizer 用戶opt_levelO2等價于FP16_Optimizer配合dynamic_loss_scaleTrue。反向傳播同樣需要改為統(tǒng)一版本optimizer.backward(loss) # - with amp.scale_loss(loss, optimizer) as scaled_loss: scaled_loss.backward()舊FP16_Optimizer的一個痛點在于用戶必須手動將模型轉(zhuǎn)為 half調(diào)用.half()或使用apex.fp16_utils的包裝函數(shù)還必須手動對輸入數(shù)據(jù)調(diào)用.half()。新 API 中這兩者都不再需要——無論選擇哪個--opt-level你都應(yīng)該用默認 FP32 格式構(gòu)建模型和傳遞輸入數(shù)據(jù)Amp 會在model, optimizer amp.initialize(...)期間根據(jù)opt_level及覆蓋標志自動完成正確的轉(zhuǎn)換。浮點輸入數(shù)據(jù)可以是 FP32 或 FP16但更簡單的做法是直接保持 FP32——amp.initialize返回的模型其forward已被 patch會自動將輸入數(shù)據(jù)轉(zhuǎn)換到合適類型。一個完整的實踐示例結(jié)合倉庫中的 Imagenet 示例一個接入 Amp 的典型訓(xùn)練腳本骨架如下省略數(shù)據(jù)加載等細節(jié)import torch from apex import amp # 1. 以 FP32 構(gòu)建模型與優(yōu)化器 model model.cuda() optimizer torch.optim.SGD(model.parameters(), lrargs.lr, momentumargs.momentum, weight_decayargs.weight_decay) # 2. 初始化 Amp先于任何 DDP 包裝 model, optimizer amp.initialize(model, optimizer, opt_levelargs.opt_level, # 如 O1 keep_batchnorm_fp32args.keep_batchnorm_fp32, loss_scaleargs.loss_scale) # 如 dynamic # 3. 分布式場景包裝 DDP 必須在 amp.initialize 之后 if args.distributed: model DDP(model, delay_allreduceTrue) # 4. 訓(xùn)練循環(huán) for epoch in range(args.start_epoch, args.epochs): for i, (input, target) in enumerate(train_loader): output model(input) loss criterion(output, target) with amp.scale_loss(loss, optimizer) as scaled_loss: scaled_loss.backward() optimizer.step() optimizer.zero_grad()命令行調(diào)用示例對應(yīng)示例倉庫中由 argparse 接收的參數(shù)python main_amp.py --arch resnet50 --opt-level O1 --loss-scale dynamic python main_amp.py --arch resnet50 --opt-level O2 --keep-batchnorm-fp32 TrueAmp 的安裝方式參見倉庫 README推薦帶 CUDA/C 擴展安裝pip install -v --no-cache-dir --global-option--cpp_ext --global-option--cuda_ext .這樣可啟用multi_tensor_apply融合內(nèi)核unscale、梯度拷貝、master weight 拷貝均受益純 Python 安裝PyTorch 0.4 時代的需求下amp仍然可用但會退化為 Python 回退實現(xiàn)并可能更慢。小結(jié)apex.amp的設(shè)計哲學(xué)是讓用戶始終以 FP32 視角編寫訓(xùn)練腳本把精度策略交給 opt_level。本文梳理了文檔 amp.rst 的核心內(nèi)容六個底層屬性cast_model_type、patch_torch_functions、keep_batchnorm_fp32、master_weights、loss_scale構(gòu)成了精度策略的原子單元O0–O3 四個級別提供了開箱即用的默認組合amp.initialize/amp.scale_loss/amp.master_params三個入口覆蓋了訓(xùn)練的全部關(guān)鍵環(huán)節(jié)而舊 API 的遷移路徑清晰且無損舊式函數(shù)標注仍然有效。結(jié)合 frontend.py、_initialize.py、_process_optimizer.py、scaler.py 與 handle.py 等源碼可以看到每個 API 承諾背后都有完整的實現(xiàn)支撐。初次嘗試時從O1最安全出發(fā)用O0/O3建立精度與速度基線再探索O2的 master weights 收益是一個穩(wěn)妥的實踐路徑。贊分享人工智能大模型音樂生成音頻預(yù)訓(xùn)練【免費下載鏈接】jukeboxCode for the paper Jukebox: A Generative Model for Music項目地址https://gitcode.com/gh_mirrors/ju/jukebox點擊查看免費下載相關(guān)推薦Ignite混合精度訓(xùn)練實戰(zhàn)AMP與Apex的完整對比Ignite混合精度訓(xùn)練實戰(zhàn)AMP與Apex的完整對比 PyTorch Ignite是一個高級庫用于幫助在PyTorch中靈活透明地訓(xùn)練和評估神經(jīng)網(wǎng)絡(luò)。在深深度學(xué)習(xí)模型訓(xùn)練使用 Apex AMP 混合精度訓(xùn)練 DCGANexamples/dcgan/main_amp.py 實戰(zhàn)指南使用 Apex AMP 混合精度訓(xùn)練 DCGAN examples/dcgan/main_amp.py 實戰(zhàn)指南 本指南圍繞 Apex 倉庫中 example人工智能深度學(xué)習(xí)分布式訓(xùn)練模型優(yōu)化Colossal-AI 混合精度訓(xùn)練完全指南基于 Booster 的 AMP 配置與實戰(zhàn)Torch/Apex/Naive AMPColossal AI 混合精度訓(xùn)練完全指南基于 Booster 的 AMP 配置與實戰(zhàn)Torch/Apex/Naive AMP 導(dǎo)讀 本文以 Colos人工智能大模型分布式訓(xùn)練深度學(xué)習(xí)模型優(yōu)化高性能計算上一篇PageMenu菜單項居中布局打造優(yōu)雅視覺平衡的終極指南下一篇嵌入式數(shù)據(jù)庫sled高級應(yīng)用合并操作符與觸發(fā)器實現(xiàn)創(chuàng)作聲明：本文部分內(nèi)容由AI輔助生成（AIGC），僅供參考