POET-X：通过扩展正交变换实现内存高效的大语言模型训练

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

arXiv: 2603.05500v1

论文信息

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

作者: Zeju Qiu, Lixin Liu, Adrian Weller, et al.

发布日期: 2026-03-05

arXiv ID: 2603.05500v1

PDF 链接: 下载 PDF

3 分钟速览

  • 研究问题：POET（Reparameterized Orthogonal Equivalence Training）通过正交等效变换为 LLM 训练提供了极强的稳定性，但其原始实现因大规模矩阵乘法导致 GPU 内存占用和计算开销过高，难以扩展到十亿参数级模型。
  • 核心方法：将 POET 的权重中心式计算重构为输入中心式矩阵‑向量乘法序列，并针对块对角正交矩阵设计批量并行乘法、定制 CUDA 置换算子、融合 Cayley‑Neumann 参数化中的计算与内存访问，最后结合梯度检查点技术大幅压缩中间激活。
  • 关键结果：POET‑X 将 GPU 内存占用降低到原 POET 的 1/3、运行速度提升 8 倍（见图 3），可在单张 Nvidia H100 上预训练 13B 参数的 Llama 模型，同时验证困惑度优于 AdamW（表 6），内存效率与 LoRA 相当（表 9）。
  • 主要局限：POET‑X 通过 fast 和 mem 两个版本在速度和内存之间权衡；量化训练（POET‑XQ）必须使用内存优化版本，无法直接与快速版本结合；块大小增大时会提升可训练参数量和内存占用。
  • 适合读者：从事大模型预训练系统优化、内存高效训练、稀疏训练、正交参数化以及分布式训练加速的研究者和工程师。

论文背景和研究动机

近年来，大型语言模型的训练极度消耗计算资源，且训练过程常常不稳定。POET 算法（Qiu et al., 2025a）通过谱保持的正交等效变换提供了显著的训练稳定性，其核心是将每个线性层的权重矩阵重新参数化为 WRP=RW0PW_{RP} = R W_0 P，其中 W0W_0 是固定的随机权重，RR 和 PP 是待学习的正交矩阵。这一形式本质上是在保持谱（奇异值）不变的前提下，只更新左右奇异向量，从而保证初始化的超球面能量下界不被破坏，稳定训练。然而，POET 的原始实现需要反复执行大规模的 RWPR W P 这类矩阵‑矩阵乘法，导致内存占用远高于 AdamW，且运行速度慢，无法在单卡或小规模集群上预训练十亿参数模型。

这正是 POET‑X 的出发点：POET 天然的块稀疏训练特性蕴含着极高的参数效率，但原始实现并未将这种参数效率转化为内存和计算效率。作者的目标是 “弥合参数效率和内存效率之间的鸿沟”，使正交等效变换真正可规模化，最终能够让单张 H100 预训练 13B 模型（图 4，表 9）。

核心方法和技术细节

POET‑X 保留了 POET 的块随机正交结构，即 Ri=Ψi⊤⋅Diag⁡(G~i1,…,G~i⌈m/b⌉)⋅ΨiR_i = \Psi_i^{\top} \cdot \operatorname{Diag}(\tilde{G}_i^1, \dots, \tilde{G}_i^{\lceil m/b\rceil}) \cdot \Psi_i，其中每个 G~ij∈Rb×b\tilde{G}_i^j \in \mathbb{R}^{b\times b} 是正交矩阵，Ψi\Psi_i 是随机置换矩阵。POET‑X 的改造主要集中在四个维度。

1. 输入中心式实现

原始 POET 的权重更新是权重中心式的：先计算 RWPR W P 再与输入 xx 相乘。POET‑X 受到矩阵自由方法的启发，将整个计算链改写为输入中心式：z=P⊤(W(R⊤x))z = P^{\top} (W (R^{\top} x))，即连续三次矩阵‑向量乘法（式 3）。这样，不再需要显式存储形如 RWPR W P 的大矩阵，大幅降低了中间激活的内存。由于 WW 本身不需要梯度，只有 RR 和 PP 对应的块对角子矩阵需要优化，梯度计算中的额外激活也被压缩到最小。

2. 置换的加速与合并

置换矩阵 Ψm\Psi_m 和 Ψn\Psi_n 原本作为密集矩阵参与乘法，POET‑X 用自定义 CUDA 算子实现索引重映射，避免构造置换矩阵，在隐藏层维度 2048 时获得近 19 倍加速（表 1）。更重要的是，作者发现可以将两次置换预先合并到权重矩阵 WW 上：在内层优化 GPG_P 和 GRG_R 时，WW 是固定的，因此可以提前计算 Φn⊤WΦm\Phi_n^{\top} W \Phi_m，使每一次训练步的置换次数从 4 次降为 2 次，进一步降低运行时开销（表 2）。

3. 块对角矩阵的批量并行乘法

由于 GRG_R 和 GPG_P 都是块对角正交矩阵，原始的 POET 需要先构建一个巨大的稀疏矩阵再执行乘法。POET‑X 将其分解为独立的块级矩阵乘法，每个块视为一个 batch 成员并行计算。这一方式避免了全矩阵的构建，既节省内存又提升速度（表 3、4），特别适合 GPU 的批量计算特性。

4. 高效的 Cayley‑Neumann 参数化（CNP）

正交性的保持采用 Cayley‑Neumann 参数化：G≈(I+Q)(I+∑i=13Qi)G \approx (I+Q)(I+\sum_{i=1}^{3} Q^i)，其中 QQ 是反对称矩阵。POET‑X 做了两个关键优化：

  • 只存储 QQ 的上三角部分，将参数数量从 b2b^2 降至 b(b−1)/2b(b-1)/2，直接令优化器状态和梯度减半。
  • 计算融合：将 GG 重写为 2(Q+Q2+Q2⋅Q)+Q2⋅Q2+I2(Q + Q^2 + Q^2 \cdot Q) + Q^2 \cdot Q^2 + I（式 9），发现前向反向传播仅依赖 QQ 和 Q2Q^2。作者使用自定义 Triton 内核，将 QQ 和 Q2Q^2 一次性加载到共享内存，就地计算 Q3Q^3, Q4Q^4 和梯度，避免了反复从全局内存读取数据，在块大小为 256 时获得约 3 倍加速（表 5）。同样的融合策略也应用于反向传播（式 10 及梯度推导），使整个 CNP 过程得到系统性的加速。

5. 梯度检查点与两个内存变体

在输入中心式的计算图中（a=GR⊤xa = G_R^{\top} x, b=Wab = W a, z=GP⊤bz = G_P^{\top} b），PyTorch 的自动求导需要保存 bb 作为中间激活，这依然会造成 N×mN \times m 级别的额外内存。作者提供了两个方案：

  • POET‑X_fast：保留 bb 的存储，依照默认自动求导逻辑，获得最快的端到端速度。
  • POET‑X_mem：对 bb 应用梯度检查点，在反向传播时重新计算 bb，进一步压缩内存，代价是额外的前向重计算开销。

这一设计允许用户根据硬件条件在速度和内存间灵活取舍（后续实验中详见表 9 和表 10）。

6. 量化训练变体 POET‑XQ

基于上述内存优化版，POET‑X 能够天然支持在‑飞‑解量化的量化训练。它只存储低比特的 WW，并在需要时解量化后执行矩阵乘法，且无需保存高精度激活。相比之下，AdamW、GaLore 等量化方案仍需额外处理高精度权重存储与计算，POET‑XQ 在内存和速度上皆具优势（表 7、8）。

创新点和贡献

  1. 首次将正交等效变换扩展到十亿参数级：原始的 POET 根本不能在 8B/13B 模型上运行（表 9 中显示 OOM），POET‑X 通过系统性的算子重写、内存压缩和内核融合，在单卡 H100 上实现 13B 预训练，且保持原 POET 的稳定性和泛化优势（图 5、表 6）。

  2. 输入中心式 + 检查点的 PEFT 级内存效率：POET‑X 仅更新左、右两个小参数块，其可训练参数量远小于全模型权重，但之前的实现因中间激活爆炸而达不到 PEFT 级别的内存节省。POET‑X 通过输入中心式计算和检查点，成功将其内存占用压低至 LoRA 同等水平（见表 9），却能够进行全谱预训练而非仅低秩微调。

  3. 独立的 CNP 优化方法：针对 Cayley‑Neumann 参数化的 Triton 融合内核以及仅仅存储上三角参数的技巧，具有独立于 POET 的推广价值，可直接用于任何需要批量生成正交矩阵的高维场景。

  4. 与量化训练的无缝结合：POET‑XQ 比相同内存量级下的 GaLore‑8bit 和 APOLLO‑8bit 获得更低的验证困惑度（14.78 vs. 17.74/20.49，表 7），证明稀疏正交变换与量化之间具有良好的相容性。

实验结果分析

单层剖析：在 Llama‑8B 的典型隐藏维度和序列长度下，POET‑X 将单层前向+反向延迟从原 POET 的 10.59ms 降至 1.38ms（POET‑X_fast）和 1.89ms（POET‑X_mem），接近高度优化的 PyTorch 线性层（图 3）。内存构成（图 4）也显示，原 POET 因存储完整转换矩阵 WRPW_{RP} 导致激活内存极高，而 POET‑X 两个版本均只保留极小额外激活，梯度与优化器状态的内存也因参数高效而显著低于 AdamW。

大规模预训练：在 C4 数据集上遵照 Chinchilla 法则训练 Llama‑3B（60B tokens），POET‑X_{b=512} 以 12.05 的验证困惑度优于 AdamW（12.69）和同内存量级的 GaLore（14.88）、APOLLO（12.97）等，仅次于 Muon（11.45），但内存占用仅为 68.52G，远低于 Muon 的 70.94G 和 AdamW 的 81.03G（表 6）。壁钟时间效率（图 5）也显示 POET‑X 收敛速度比 AdamW 更快，得益于其 DDP 友好的内存特性，避免了 FSDP 引入的通信开销。

内存与吞吐可扩展性：对 3B/8B/13B 模型在不同序列长度下的测试（表 9）表明，POET‑X_mem 在所有配置下的内存均小于 LoRA，例如 13B+2048 序列时仅占 47.21G（b=256）和 59.02G（b=512），而 LoRA_{r=320} 需要 71.55G。吞吐量方面（表 10、14），POET‑X_fast 在 64 卡配置下吞吐与 LoRA 相当甚至更优，且其单卡到多卡的吞吐缩放比（Ratio）接近理想线性，远好于因 FSDP 通信而严重退化的 AdamW（图 6）。

量化训练：POET‑XQ 在保持更低内存的同时，吞吐量高于 Q‑GaLore 和 Q‑APOLLO（表 8），并且验证困惑度显著优于两者（表 7），展示了稀疏正交参数化在低精度训练场景下的潜力。

实践建议

如果希望将 POET‑X 集成到自己的 LLM 预训练流程中，以下几个维度值得参考：

  1. 块大小选择：b=256b=256 提供最低的可训练参数量和内存消耗，但在本论文实验中 b=512b=512 通常获得更好的困惑度（表 6、7）。建议根据可用显存进行取舍：若单卡内存紧俏，优先选 b=256b=256 的 POET‑X_mem；如果内存宽裕且追求性能，可尝试 b=512b=512。

  2. fast 与 mem 版本切换：POET‑X_fast 不做检查点，速度快但额外保存一个中间激活张量；POET‑X_mem 通过重计算节省内存，适合大序列或大模型场景。在 13B 模型下使用长序列时，POET‑X_mem 的内存优势极为明显（表 9），因此推荐在接近显存上限时启用 mem 版本。

  3. 量化使用：若需要量化训练以达到更高压缩比，务必使用 POET‑X_mem 作为基础，因为 POET‑XQ 依赖于在反向时重新计算激活，无法与 fast 版本直接结合。同时，最好调整学习率缩放因子 γ\gamma（论文中设定为 0.5），以保证在低精度下正交矩阵更新的稳定性。

  4. 分布式训练：POET‑X 极其适合 DDP 模式，因为每块 GPU 上只需要放置一份完整的、无梯度的基础权重 W0W_0，以及量级极小的正交参数和优化器状态。这可以有效避免 FSDP 带来的分片通信负担，显著提升多节点扩展效率（图 6）。在搭建多节点训练环境时，建议直接采用 DDP 并适当调整梯度累积步数以匹配全局批次大小。

  5. 自定义内核部署：POET‑X 加速的核心在于置换、块对角乘法和 CNP 的定制 CUDA/Triton 内核。若将 POET‑X 迁移到其他框架或非 NVIDIA 硬件，需重写这些算子；一个折中方案是采用论文中表 1、2、3、5 所示的 PyTorch 原生实现作为原型，但需预期约 2‑20 倍的性能下降。

综合来看，POET‑X 在一个可实践的 GPU 条件下同时实现了 AdamW 无法企及的内存压缩和稳定的训练质量，对于资源受限的研究团队或需要最大限度利用单卡算力的场景，具有直接的应用价值。