在单张 H100 上预训练 13B 模型：POET-X 如何把正交变换的内存开销砍到 1/3

想象一个场景：你是一个小实验室的博士生，手里只有一张 H100。你想预训练一个 13B 参数的 Llama，但 AdamW 告诉你——想都别想。13B 模型的 FP16 权重 26GB，AdamW 的 optimizer state 还要再吃 52GB（两份 FP32），加上梯度、激活值，单卡 80GB 显存根本不够。标准答案是多卡张量并行，但多卡要钱，钱从哪来？

剑桥、港科大、上交、MPI 合作的这篇论文给出的答案很反直觉：别用 AdamW 了，换优化器。他们叫 POET-X，是 2025 年 POET 的内存优化版本。核心声明一句话：在单张 H100 上预训练 Llama-8B/13B，显存占用和 LoRA 相当，性能比 AdamW 还好。

POET 是什么：不优化权重，优化"正交变换"

要理解 POET-X，必须先理解 POET 在干什么。

POET 的核心思想是谱保持（spectrum preservation）。它不直接优化权重矩阵 W，而是把 W 固定为一个随机矩阵，然后在它左右各乘一个正交矩阵 R 和 P：W_RP = R · W · P。训练时只优化 R 和 P，W 保持不变。

为什么这么做？因为正交变换保持矩阵的奇异值不变。也就是说，不管 R 和 P 怎么变，W_RP 的谱性质始终和初始随机矩阵 W 一样。这个"谱保持"性质带来了两个好处：训练稳定性（不会梯度爆炸/消失）和泛化性（hyperspherical energy 小）。

但 POET 有一个致命问题：R 和 P 都是 m×m 和 n×n 的稠密正交矩阵，训练时需要存它们的参数、梯度和 optimizer state，还要做大量矩阵乘法。论文里说 POET 的内存开销比 AdamW 还大，速度还慢 8 倍。所以 POET 提出后基本没人用——理论再好，跑不动也是白搭。

POET-X 要解决的就是这个"理论好但工程跑不动"的问题。

POET-X 的五板斧

POET-X 的贡献不是提出一个新算法，而是把 POET 的工程实现优化到可用。论文里列了五个优化点，每个都是针对 POET 的具体瓶颈。

第一斧：从"权重中心"到"输入中心"

POET 原始实现的 forward pass 是这样的：先算 R·W·P（两次矩阵-矩阵乘法），再乘输入 x（一次矩阵-向量乘法）。复杂度 O(nm²)，因为 R·W 是 m×m 乘 m×n。

POET-X 的改写很巧妙：把乘法顺序换一下。先算 R^T·x（矩阵-向量），再算 W^T·(R^T·x)（矩阵-向量），最后算 P^T·(...)（矩阵-向量）。三次矩阵-向量乘法，复杂度 O(nm)。

这个改写不是免费的——计算 P 的梯度时仍然需要访问 W。论文花了很大篇幅讲怎么用"置换加速"和"批并行"绕过这个困难。但核心思路就是这一步：把矩阵-矩阵乘法降级为矩阵-向量乘法。

第二斧：置换加速与合并

POET 的正交矩阵 R 和 P 内部包含随机置换矩阵 Ψ。POET 原始实现会显式构造这些 m×m 的置换矩阵，然后做矩阵乘法。POET-X 的优化是：置换矩阵不需要显式构造，只需要存一个索引数组。

论文写了一个自定义 CUDA kernel，把"矩阵乘置换矩阵"变成"按索引重排行/列"。Table 1 显示这个优化带来了 14-20× 的加速。

更聪明的是"置换合并"：forward pass 里需要 4 次置换，其中 2 次可以提前合并到权重矩阵 W 上（因为 W 在 inner loop 里是固定的）。这样实际运行时只需要 2 次置换。

第三斧：块对角的批并行

POET 的正交矩阵不是稠密的，而是块对角的：G = Diag(G̃¹, G̃², ..., G̃^{m/b})。每个块是 b×b 的小矩阵（b=256 或 512）。

POET 原始实现会先把这个块对角矩阵显式构造出来，然后做稀疏矩阵乘法。POET-X 的观察是：既然每个块独立，为什么不直接做 batch 矩阵乘法？ 跳过显式构造，把每个块当成 batch 的一个元素，用 torch.bmm 一次算完。Table 3 显示 2.3× 加速，Table 4 显示额外 9-31% 的内存节省。

第四斧：Cayley-Neumann 参数化的 Triton 重写

POET 用 Cayley-Neumann 参数化（CNP）来保证正交性。CNP 的核心是 Cayley 变换的近似：G ≈ (I+Q)(I+Q+Q²+Q³)，其中 Q 是斜对称矩阵。

POET-X 对 CNP 做了三个优化：

1. 只存斜对称矩阵的上三角：Q 是 b×b 的斜对称矩阵，只有 b(b-1)/2 个独立参数。POET-X 只存这些，参数量和 optimizer state 直接减半。 2. 重排计算公式：把 CNP 展开式重排为 G ≈ 2(Q+Q²+Q²·Q) + Q²·Q² + I，发现所有项都只依赖 Q 和 Q²。只需要算一次 Q²，后续全靠它组合。 3. Triton kernel fusion：把 Q 和 Q² 加载到 shared memory 一次，在同一个 Triton kernel 里算出 Q³、Q⁴ 和最终结果。避免反复从 global memory 读数据。

Table 5 显示这个优化带来 1.9-3× 的加速。

第五斧：激活值检查点

POET-X 的 forward pass 是三步矩阵乘法：a = G_R^T·x, b = W·a, z = G_P^T·b。PyTorch Autograd 需要存中间激活值 b 来算 G_P 的梯度。

POET-X 的观察是：b 可以从 z 和 G_P 反算回来（b = G_P·z），所以不需要存。这就是检查点（checkpointing）的标准操作——用重算换内存。论文把这个和前面的优化叠加，最终实现了和 LoRA 相当的内存效率。

实验：单卡 H100 跑 13B 预训练

论文的核心实验结果：

单卡 H100 上预训练 Llama-8B：POET-X 跑得动，AdamW OOM。 单卡 H100 上预训练 13B 模型：POET-X 也能跑。 内存占用：和 LoRA 相当，远低于 AdamW。 运行速度：比原始 POET 快 8×，和 AdamW 相当。 性能：在多个 benchmark 上比 AdamW 好——不是"将就能用"，而是"真的更好"。

这个结果的意义不仅是"省了一张卡的钱"。它意味着：小实验室也能预训练大模型了。以前 13B 预训练至少要 8 卡，现在单卡就行。这改变了 LLM 研究的准入门槛。

我的几点观察

第一，这篇论文是"工程驱动理论"的典范。POET 的理论在 2025 年就提出了，但没人用，因为工程实现太烂。POET-X 没有改任何理论，只是把实现优化到可用——但正是这个工程工作，让 POET 从"论文里的漂亮想法"变成了"实际能用的优化器"。这提醒我们：在系统研究中，实现质量比理论优雅更重要。

第二，五个优化点都不是新发明。矩阵-向量替代矩阵-矩阵、索引替代置换矩阵、batch bmm、kernel fusion、checkpointing——每个都是 GPU 编程的常规操作。POET-X 的贡献是系统性地把这些技术应用到 POET 的每一个瓶颈上。这种"把已知技术组合到极致"的研究风格，和 FlashAttention 是一类。

第三，"和 LoRA 相当的内存、比 AdamW 更好的性能"这个组合很有杀伤力。LoRA 省内存但只做微调，POET-X 省内存且做预训练。这意味着 POET-X 实际上填补了一个空白：在 LoRA 的内存预算下做 AdamW 级别的训练。

第四，谱保持的稳定性优势在 scale up 时更明显。论文提到 POET-X 在多节点训练中也表现稳定，这和谱保持的理论预期一致——正交变换不会让梯度爆炸。这个性质在训练超大模型时可能比内存节省更重要。

局限与未解之问

论文没有讨论几个重要问题：

第一，block size b 的选择。论文用 b=256 或 512，但没有讨论这个选择对性能的影响。b 太小则正交矩阵的表达能力不足，b 太大则内存和计算开销上升。这个 trade-off 的最优点在哪？

第二，和其他内存高效优化器的比较。论文只比了 AdamW 和 LoRA，没比 8-bit Adam、Adafactor、Sophia 这些同样主打内存高效的优化器。POET-X 相对它们的优势有多大？

第三，长训练的稳定性。论文的实验是预训练级别的（几百B token），但 POET 的谱保持性质在极长训练（几T token）下是否仍然稳定？正交矩阵的数值误差会不会累积？

第四，推理开销。论文说 R 和 P 训练后可以合并到 W 里，推理零开销。但合并后的 W 和直接训练的 W 在数值精度上有没有差异？这个差异对下游任务的影响有多大？

结语

POET-X 让我想起一个被忽视的事实：很多好算法不是被理论否定的，而是被工程实现杀死的。POET 的理论在 2025 年就成立，但直到 2026 年有人把工程做好，它才真正可用。

这篇论文的深层启示是：在 LLM 训练这个领域，"系统优化"和"算法创新"同样重要。一个算法再优雅，如果实现质量差，就不会有人用。反过来，一个算法虽然"旧"，但如果工程做到极致，也能焕发第二春。POET-X 属于后者。

代码已开源：https://github.com/Sphere-AI-Lab/poetx

---

*论文链接：https://arxiv.org/abs/2603.05500* *项目主页：https://spherelab.ai/poetx* *代码仓库：https://github.com/Sphere-AI-Lab/poetx*