TransformerEngine

收录理由

Transformer Engine 是 NVIDIA 自家推出的库，专门让低精度训练在其 GPU 上变得切实可用。它直接提供 PyTorch 和 JAX 模块，用来搭建支持 FP8 的 Transformer 层；到了 Blackwell 架构上，还加入 MXFP8 和 NVFP4，从而在训练和推理中降低显存占用、提升速度，省去了混合精度通常需要手动维护缩放因子的麻烦。NVIDIA 自己的许多大语言模型和 MoE 训练任务都在用它，并且它能接入 NeMo 和 MaxText，说明这些优化确实经过大规模实战检验。在 Hopper、Ada 和 Blackwell GPU 上，融合内核以及针对 MoE、张量/序列/上下文并行等特性的支持都直接可用；框架无关的 C++ API 也让其他深度学习库能方便地加上 FP8 支持。不过它并不是开箱即用的训练工具，使用时需要在你现有的 PyTorch 或 JAX 技术栈里集成，而不是引入一套全新的工作流。