[ICLR 2025] POET-X：打破大模型训练的内存围城，实现 13B 模型单卡预训练

TL;DR

在 LLM 预训练领域，内存效率与训练稳定性往往不可兼得。POET-X 通过对正交等价变换（POET）进行底层重构，利用 Input-centric 计算、Triton 算子融合以及块稀疏优化，实现了媲美 LoRA 的极低内存占用。它能让单张 H100 跑起 13B 模型预训练，且收敛速度和质量均超越了传统的 AdamW 优化器。

1. 痛点：为什么原始 POET 叫好不叫座？

正交等价变换（POET）最初被提出是为了解决 LLM 训练中的稳定性问题。它通过保持权重矩阵的奇异值（Spectrum-preserving）来防止梯度消失或爆炸。然而，原始 POET 存在一个致命伤：内存爆炸。

• Weight-centric 弊端：原始方法直接操作巨大的权重矩阵 $W$，计算 $RWP$。这需要大量的中间激活存储，导致内存消耗甚至超过了 AdamW。
• 计算冗余：频繁的排列（Permutation）操作和非融合算子导致 GPU 吞吐量极低，无法在主流集群上扩展。

2. 核心机理：从权重中心到输入中心

POET-X 的核心直觉在于：不需要显式地去变换权重矩阵，而是通过一系列线性映射来变换输入信号。

2.1 以输入为中心的重构 (Input-centric)

作者将 $z=(RPW{)}^{o}px$ 的计算逻辑从“更新权重”转变为“对输入向量 $x$ 进行一系列正交变换”。这种方法避免了存储变换后的中间权重矩阵，直接将内存复杂度从 $O(n{m}^2)$ 降低到线性水平。

2.2 极致的算子加速

为了解决计算开销，POET-X 引入了三项关键改进：

1. 排列加速 (Permutation Acceleration)：不再构建显式的置换矩阵，而是直接通过自定义 CUDA Kernel 进行索引映射。
2. 块并行 CNP：利用 Triton 实现了 Cayley-Neumann 参数化的内核融合，仅存储 skew-symmetric 矩阵的一半，内存占用直接减半。
3. 梯度检查点 (Checkpointing)：提供了 POET-Xmem 模式，通过在后向传播时重新计算中间激活，实现了极致的显存压缩。

上图展示了 POET-X 如何通过 Input-centric 实现相比 AdamW 和原始 POET 更优的显存特征。

3. 实验结果：单卡挑战 13B

在 H100 上的测试显示，POET-X 的表现令人惊艳。

• 显存奇迹：在 Llama-8B 设置下，AdamW 消耗约 81GB 显存，而 POET-Xmem 仅需 26GB 左右，达到了与 LoRA 相当的水平。
• 性能超越：在 C4 数据集的预训练中，POET-X 的验证困惑度（Perplexity）一致优于 AdamW、GaLore 和 APOLLO。

表 6 显示了在 3B 模型规模下，POET-X (b=512) 取得了仅次于 Muon 的收敛性能，但显存占用远低于后者。

4. 分布式扩展性：逃离 FSDP 通信陷阱

由于 POET-X 的显存占用极低，它允许开发者在单个节点内使用 DDP (Distributed Data Parallel) 而非复杂的 FSDP (Fully Sharded Data Parallel)。

• FSDP 因为需要跨卡切分权重和梯度，会产生巨大的 All-gather 通信开销。
• POET-X 能够将整个模型参数放入单卡，仅通过带宽占用极小的梯度规约即可完成同步。这使得它在 64 GPU 上的线性扩展比率（Scaling Ratio）远高于 AdamW 相关变体。

5. 总结与洞察

POET-X 不仅仅是一个优化器插件，它代表了 LLM 训练的一种新范式：通过数学上的结构化约束（正交性）来换取稳定，再通过工程上的极致优化（内核融合）来换取效率。

局限性：虽然内存效率极高，但由于引入了额外的正交变换步骤，其单步迭代时间（Raw latency）仍比简单的 Linear 层略高。然而，考虑到其带来的稳定性增益和减少的通信开销，这在超大规模预训练中是一个非常划算的交易。

未来启示：POET-X 的成功暗示，未来的大模型预训练可能会越来越多地采用“由于数学结构导致的稀疏性”，而非单纯的硬剪枝或低秩近似。