超级会员免费看
TransformerEngine（TE）是NVIDIA为在GPU上加速Transformer模型而设计的库，尤其在Hopper GPU上使用FP8精度进行训练和推理。本文通过对源码的阅读，探讨TE如何维护FP8训练中的scale，实现层融合策略。代码示例展示了与FP32训练的不同之处，主要涉及te.Linear、fp8_recipe和te.fp8_autocast。文章还概述了TE的内部工作流程，包括Linear类、amax_history管理和硬件相关的cublas库调用。尽管尚未深入研究所有细节，但已揭示了TE的核心功能和结构。
在几个月前nv就发表过关于fp8数制训练和推理的白皮书，大概三四个月前公开了他们针对transformer模型的fp8训练的软件库TransformerEngine，由于最近在关注fp8训练，因此想了解一下他们的实现方案，但是由于没有H100的卡，目前还不能跑通te的代码，很多细节没办法验证，只能通过源码的阅读来观察它插入了什么功能实现fp8量化训练，特此记录一下。
“Transformer Engine (TE) 是一个用于在 NVIDIA GPU 上加速 Transformer 模型的库，包括在 Hopper GPU 上使用 8 位浮点 (FP8) 精度，以在训练和推理中以较低的内存利用率提供更好的性能。 TE 为流行的 Transformer 架构提供了一系列高度优化的构建块，以及可与您自己的特定于框架的代码无缝使用的类似自动混合精度的 API。 TE 还包括一个与框架无关的 C++ API，它可以与其他深度学习库集成，以实现对 Transformers 的 FP8 支持。”抄自te官方文档https://github.com/NVIDIA/TransformerEngine的一段话。总体来说，te可以实现fp8数制在训练中需要的scale的维护&#
订阅专栏 解锁全文
1314
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
