LLM System: Transformer Engine 00 - 学习地图

Thu, 21 May 2026 06:00:00 +0800

这篇文章是 LLM System 系列里 Transformer Engine 子专题的第 0 篇，也是这个主题的学习入口。

我准备用这个系列系统学习 NVIDIA Transformer Engine。目标不是只会调用几个 te.* API，而是把 TE 放到大模型训练系统里理解：它为什么存在、如何利用 FP8 / FP4 和 fused kernel、怎么被 Megatron-LM 接入，以及后续如何用 profiler 分析和改进它。

0. 学习目标

这组笔记会围绕四件事展开：

1. TE 在 AI Infra 技术栈中的位置
2. TE 的核心加速策略：低精度、融合算子、Attention、MoE、通信重叠
3. Megatron-LM / Megatron-Core 如何接入 TE
4. 如何 benchmark、trace 并尝试改进 TE

1. 技术定位

Transformer Engine 不是训练框架，而是 NVIDIA 为 Transformer 训练/推理提供的高性能 building block 库。它大致位于训练框架和底层 CUDA/cuBLAS/cuDNN kernel 之间。

这一阶段要回答的问题：

1. TE 和 PyTorch AMP 的关系是什么？
2. TE 和 Megatron-LM 的边界在哪里？
3. TE 为什么不是一个完整训练框架？
4. TE 为什么不只是 FP8，而是 Transformer 优化库？
5. TE 在 AI Infra 技术栈里更靠近 compiler/runtime/kernel，还是更靠近 model framework？

2. 核心加速策略

核心加速策略先按五条线学习：

FP8 on Echo的技术博客

LLM System: Transformer Engine 00 - 学习地图

0. 学习目标

1. 技术定位

2. 核心加速策略