如何使用 NVIDIA Apex（FusedAdam、FusedLayerNorm）和原生 torch.amp 加速 Transformer 训练

本教程详细实现了 NVIDIA Apex，重点关注现代 GPU 训练工作流程中仍然重要的组件。我们首先检查 CUDA 运行环境，从源码构建 Apex（包含 CUDA 和 C++ 扩展），并检测环境中实际可用的融合内核。这对于确保高性能内核可用至关重要，因为纯 Python 安装可能看似成功但实际上缺少关键内核。

构建完成后，我们基准测试 FusedAdam 与 PyTorch AdamW，比较 FusedLayerNorm 和 FusedRMSNorm 与标准归一化层，并运行已弃用的 apex.amp 和现代 torch.amp 示例。最后，在一个小型 Transformer 训练实验中整合所有内容，比较纯 FP32 PyTorch 路径与融合 Apex+AMP 路径，评估对吞吐量的实际影响。

实验设置包括：检测 CUDA GPU 可用性、安装必要依赖、克隆并构建 Apex 源码。通过检查 amp_C 和 fused_layer_norm_cuda 模块判断内核可用性。

基准测试结果：

• FusedAdam 在优化器步骤上比 AdamW 快约 2.67 倍。
• FusedLayerNorm 的前向+反向传播比标准 LayerNorm 快约 1.87 倍。
• FusedRMSNorm 与 FusedLayerNorm 类似，适用于 LLaMA 风格模型。

混合精度部分演示了弃用的 apex.amp 和推荐的 torch.amp 用法。端到端 Transformer 训练（词汇量 2000，隐藏维度 256，4 层）显示，使用 FusedAdam、FusedLayerNorm 和 torch.amp（FP16）相比纯 FP32 基线，吞吐量提升约 1.48 倍，且损失收敛相似。

结论：FusedAdam、FusedLayerNorm 和 FusedRMSNorm 是 Apex 中仍然相关的部分；apex.amp 已弃用，应使用 torch.amp；融合内核与原生 torch.amp 配合良好；对于实际工作负载，建议使用更大模型和 bf16 自动转换。