Transformer Engine

这篇文章是 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. 核心加速策略 核心加速策略先按五条线学习： ...

本篇目标：了解 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。 ...