apex

Apex 是 NVIDIA 维护的一组工具集，旨在让 PyTorch 用户快速使用最新的混合精度和分布式训练功能，其中部分代码计划最终合并到 PyTorch 上游。 它提供的内容： - 混合精度训练辅助工具（apex.amp），用于以较低精度运行模型。 - 分布式训练工具，包括 apex.parallel.DistributedDataParallel 和 SyncBatchNorm。 - 优化器的融合内核，如 FusedAdam，以及归一化层，如 FusedLayerNorm 和 FusedRMSNorm。 - 一组可选的 contrib 模块，涵盖稀疏性、组批归一化、xentropy、focal loss、transducer 损失、peer memory、NCCL P2P、瓶颈层、分布式优化器和 GPU direct storage 等领域。 安装说明： - NGC 上的 NVIDIA PyTorch 容器已预装自定义扩展。 - 从源码安装时，推荐使用环境变量如 APEX_CPP_EXT=1 和 APEX_CUDA_EXT=1 配合 pip install --no-build-isolation；旧版命令行标志（--cpp_ext、--cuda_ext）仍然支持。 - 可以仅构建 Python 版本，但会省略融合内核，因此 FusedAdam、FusedLayerNorm/FusedRMSNorm 以及 SyncBatchNorm、DistributedDataParallel 和 amp 的融合路径将不可用或速度较慢。 - 各个 contrib 模块需要各自的安装选项；有些可能仅与 PyTorch nightly 版本兼容，而非稳定版。APEX_ALL_CONTRIB_EXT=1 可一次性构建所有 contrib 扩展。 - 推荐使用 Ninja 以加快编译速度，并针对 CPU 或内存有限的环境提供了并行构建选项的文档。 - Windows 支持被描述为实验性的。 这是一个面向训练深度学习模型的开发者的库，他们希望获得混合精度和分布式训练构建模块，而非最终用户应用程序。