apex

介紹

該存儲庫持有NVIDIA維護的實用程序，以簡化Pytorch的混合精度和分佈式培訓。此處的某些代碼最終將包含在上游Pytorch中。 APEX的目的是盡快為用戶提供最新的實用程序。

完整的API文檔：https：//nvidia.github.io/apex

GTC 2019和Pytorch Devcon 2019幻燈片

內容

1。放大器：自動混合精度

棄用。使用Pytorch放大器

apex.amp是一種工具，可以通過僅更改腳本的3行來實現混合精確培訓。用戶可以通過提供不同的標誌來進行amp.initialize來輕鬆嘗試使用不同的純淨和混合精確訓練模式。

網絡研討會引入AMP（FLAG cast_batchnorm已重命名為keep_batchnorm_fp32 ）。

API文檔

全面的成像網示例

DCGAN示例即將到來...

轉到新的AMP API（適用於未經棄用的“ AMP”和“ FP16_OPTIMIZER” API的用戶）

2。分佈式培訓

apex.parallel.DistributedDataParallel被棄用。使用torch.nn.parallel.DistributedDataParallel

apex.parallel.DistributedDataParallel是一個模塊包裝器，類似於torch.nn.parallel.DistributedDataParallel 。它啟用了方便的多進程分佈式培訓，可針對NVIDIA的NCCL通信庫進行了優化。

API文檔

Python來源

示例/演練

ImageNet示例顯示了apex.parallel.DistributedDataParallel以及apex.amp的使用。

同步分批歸一化

棄用。使用torch.nn.SyncBatchNorm

apex.parallel.SyncBatchNorm擴展了torch.nn.modules.batchnorm._BatchNorm以支持同步BN。它在多進程（分佈式dataparallel）培訓期間跨過程中的統計數據。同步BN已在只有小的局部MiniBatch適合每個GPU的情況下使用。在所有過程中，Allreded Stats將BN層的有效批量大小提高到全局批處理大小（從技術上講，這是正確的公式）。已經觀察到同步BN在我們的一些研究模型中可以提高融合精度。

檢查點

為了正確保存和加載您的amp培訓，我們介紹了amp.state_dict() ，其中包含所有loss_scalers及其相應的不合理步驟，以及amp.load_state_dict()以還原這些屬性。

為了獲得位準確性，我們建議以下工作流程：

 # Initialization
opt_level = 'O1'
model , optimizer = amp . initialize ( model , optimizer , opt_level = opt_level )

# Train your model
...
with amp . scale_loss ( loss , optimizer ) as scaled_loss :
    scaled_loss . backward ()
...

# Save checkpoint
checkpoint = {
    'model' : model . state_dict (),
    'optimizer' : optimizer . state_dict (),
    'amp' : amp . state_dict ()
}
torch . save ( checkpoint , 'amp_checkpoint.pt' )
...

# Restore
model = ...
optimizer = ...
checkpoint = torch . load ( 'amp_checkpoint.pt' )

model , optimizer = amp . initialize ( model , optimizer , opt_level = opt_level )
model . load_state_dict ( checkpoint [ 'model' ])
optimizer . load_state_dict ( checkpoint [ 'optimizer' ])
amp . load_state_dict ( checkpoint [ 'amp' ])

# Continue training
...

請注意，我們建議使用相同的opt_level還原模型。另請注意，我們建議在amp.initialize之後調用load_state_dict方法。

安裝

每個apex.contrib模塊需要一個或多個安裝選項，而不是--cpp_ext和--cuda_ext 。請注意，貢獻模塊不一定支持穩定的Pytorch發行。

容器

NGC：https：//catalog.ngc.nvidia.com/orgs/nvidia/nvidia/containers/pytorch上。目前，容器配備了所有可用的自定義擴展名。

有關詳細信息，請參見NGC文檔：

• 如何拉容器
• 如何運行拉動容器
• 發行說明

來自來源

要從源安裝APEX，我們建議使用可從https://github.com/pytorch/pytorch獲得的夜間pytorch。

可以從https://pytorch.org獲得的最新穩定版本也應起作用。

我們建議安裝Ninja以更快地進行編譯。

Linux

對於性能和完整功能，我們建議通過CUDA和C ++擴展安裝APEX

git clone https://github.com/NVIDIA/apex
cd apex
# if pip >= 23.1 (ref: https://pip.pypa.io/en/stable/news/#v23-1) which supports multiple `--config-settings` with the same key... 
pip install -v --disable-pip-version-check --no-cache-dir --no-build-isolation --config-settings " --build-option=--cpp_ext " --config-settings " --build-option=--cuda_ext " ./
# otherwise
pip install -v --disable-pip-version-check --no-cache-dir --no-build-isolation --global-option= " --cpp_ext " --global-option= " --cuda_ext " ./

Apex還支持僅通過Python構建

pip install -v --disable-pip-version-check --no-build-isolation --no-cache-dir ./

僅Python構建的省略：

• 使用apex.optimizers.FusedAdam所需的融合內核。
• 使用apex.normalization.FusedLayerNorm和apex.normalization.FusedRMSNorm所需的融合內核。
• 融合的內核可以改善apex.parallel.SyncBatchNorm的性能和數值穩定性。
• 融合的內核可以改善apex.parallel.DistributedDataParallel和apex.amp的性能。 DistributedDataParallel ， amp和SyncBatchNorm仍然是可用的，但它們可能較慢。

[實驗]窗口

pip install -v --disable-pip-version-check --no-cache-dir --no-build-isolation --config-settings "--build-option=--cpp_ext" --config-settings "--build-option=--cuda_ext" .如果您能夠從系統上的來源構建Pytorch，則可以工作。通過pip install -v --no-cache-dir .更有可能工作。
如果您在Conda環境中安裝了Pytorch，請確保在同一環境中安裝頂點。

自定義C ++/CUDA擴展和安裝選項

如果未滿足模塊的要求，則不會構建它。

模塊名稱	安裝選項	雜項
apex_C	--cpp_ext
amp_C	--cuda_ext
syncbn	--cuda_ext
fused_layer_norm_cuda	--cuda_ext	apex.normalization
mlp_cuda	--cuda_ext
scaled_upper_triang_masked_softmax_cuda	--cuda_ext
generic_scaled_masked_softmax_cuda	--cuda_ext
scaled_masked_softmax_cuda	--cuda_ext
fused_weight_gradient_mlp_cuda	--cuda_ext	需要cuda> = 11
permutation_search_cuda	--permutation_search	apex.contrib.sparsity
bnp	--bnp	apex.contrib.groupbn
xentropy	--xentropy	apex.contrib.xentropy
focal_loss_cuda	--focal_loss	apex.contrib.focal_loss
fused_index_mul_2d	--index_mul_2d	apex.contrib.index_mul_2d
fused_adam_cuda	--deprecated_fused_adam	apex.contrib.optimizers
fused_lamb_cuda	--deprecated_fused_lamb	apex.contrib.optimizers
fast_layer_norm	--fast_layer_norm	apex.contrib.layer_norm 。不同於fused_layer_norm
fmhalib	--fmha	apex.contrib.fmha
fast_multihead_attn	--fast_multihead_attn	apex.contrib.multihead_attn
transducer_joint_cuda	--transducer	apex.contrib.transducer
transducer_loss_cuda	--transducer	apex.contrib.transducer
cudnn_gbn_lib	--cudnn_gbn	需要cudnn> = 8.5， apex.contrib.cudnn_gbn
peer_memory_cuda	--peer_memory	apex.contrib.peer_memory
nccl_p2p_cuda	--nccl_p2p	需要nccl> = 2.10， apex.contrib.nccl_p2p
fast_bottleneck	--fast_bottleneck	需要peer_memory_cuda和nccl_p2p_cuda ， apex.contrib.bottleneck
fused_conv_bias_relu	--fused_conv_bias_relu	需要cudnn> = 8.4， apex.contrib.conv_bias_relu