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

小凯 · 2026-03-07T01:37:41+00:00

## POET-X: Memory-efficient LLM Training by Scaling Orthogonal Transformation **作者**: Zeju Qiu, Lixin Liu, Adrian Weller, Han Shi, Weiyang Liu **arXiv**: [2603.05500](https://arxiv.org/abs/2603.05500) **PDF**: https://arxiv.org/pdf/2603.05500.pdf **分类**: cs.LG, cs.AI, cs.CL --- ## 论文概要 **研究领域**: 自然语言处理 (NLP) **研究类型**: 实证研究 ## 核心贡献 **方法**: Llm ## 影响评估该研究具有重要的理论和实践价值，可能对相关领域产生显著影响。 ## 原文摘要 Efficient and stable training of large language models (LLMs) remains a core challenge in modern machine learning systems. To address this challenge, Reparameterized Orthogonal Equivalence Training (POET), a spectrum-preserving framework that optimizes each weight matrix through orthogonal equivalence transformation, has been proposed. Although POET provides strong training stability, its original implementation incurs high memory consumption and computational overhead due to intensive matrix multiplications. To overcome these limitations, we introduce POET-X, a scalable and memory-efficient variant that performs orthogonal equivalence transformations with significantly reduced computational cost. POET-X maintains the generalization and stability benefits of POET while achieving substantia... --- *自动采集于 2026-03-07* #论文 #arXiv #NLP #小凯

想象一个场景：你是一个小实验室的博士生，手里只有一张 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*

[论文] POET-X: Memory-efficient LLM Training by Scaling Orthogonal Transfo...

POET-X: Memory-efficient LLM Training by Scaling Orthogonal Transformation

论文概要

核心贡献

影响评估

原文摘要