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. 核心加速策略

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

LLM System: Transformer Engine 01 - 在 AI Infra 技术栈中的位置

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

本篇目标：了解 Transformer Engine 的技术定位，搞清楚它为什么存在，以及它和 PyTorch、cuBLAS、Megatron-LM 的边界。

基本接口

Layer 类定义接口非常直接，最表层的使用方式就是把 torch.nn 模块替换成 transformer_engine.pytorch 模块。

普通 PyTorch 写法：

self.linear = torch.nn.Linear(hidden_size, 4 * hidden_size)

TE 写法：

import transformer_engine.pytorch as te

self.linear = te.Linear(hidden_size, 4 * hidden_size)

量化上下文：

from transformer_engine.pytorch import fp8_autocast

with fp8_autocast(enabled=True):
 y = module(x)

进入 TE 的 FP8 上下文之后，TE 会围绕量化、反量化、fused path、tensor cache 和 backend 选择做一系列处理。相较于纯 PyTorch 计算图优化，TE 会拿到更多信息 tensor parallel、sequence parallel、FP8 recipe 等。这些额外信息给底层算子优化留下了空间。

这一点目前还是比较 general 层面的认知，后面要继续顺着源码和 profiler trace 去验证。

TE 和 Megatron 的边界

Megatron-Core 负责模型并行、训练 loop、optimizer、activation checkpoint、MoE routing、pipeline schedule 和 config。