规模化 DoRA：通过因式分解范数和融合内核实现高秩适应 Scaling DoRA: High-Rank Adaptation via Factored Norms and Fused Kernels
通过因式分解范数和融合 Triton 内核，解决 DoRA 高秩配置下的内存和效率瓶颈
前置知识
DoRA (Weight-Decomposed Low-Rank Adaptation)
DoRA 是 LoRA 的扩展方法，通过将适配权重分解为幅度和方向两部分来提升性能。具体公式为 $W' = m \odot \frac{W + sBA}{\|W + sBA\|_{row}}$，其中 $W$ 是冻结的基础权重，$B$ 和 $A$ 是低秩因子，$s$ 是缩放系数，$m$ 是可学习的幅度向量。相比于 LoRA，DoRA 通过幅度方向的解耦可以更精细地控制权重更新，在高秩配置下能更好地逼近全参数微调。
本文的核心问题就是 DoRA 的行向范数计算 $\|W + sBA\|_{row}$ 造成的内存和计算瓶颈。只有理解 DoRA 的数学定义，才能明白为什么需要计算行向范数，以及为什么现有实现会通过构造单位矩阵来密集化 BA 乘积。
LoRA (Low-Rank Adaptation)
LoRA 是参数高效微调的代表性方法，通过冻结预训练权重 $W$ 并注入低秩更新 $\Delta W = BA$ 来微调模型，其中 $B$ 和 $A$ 是低秩因子，rank 远小于维度。微调时只训练 $A$ 和 $B$ 参数，大幅减少可训练参数量。推理时可以将 $BA$ 合并到 $W$ 中，不增加推理开销。LoRA 已成为大模型微调的事实标准。
DoRA 是在 LoRA 基础上的扩展，理解 LoRA 的低秩分解思想有助于理解 DoRA 的改进之处。本文在处理 DoRA 的范数计算时，核心挑战正是避免材料化完整的 BA 低秩乘积，这与 LoRA 本身的设计思想相呼应。
Triton 语言
Triton 是一种类 Python 的语言和编译器，用于编写高性能 GPU 内核。它比 CUDA 更易用，同时能接近手写 CUDA 的性能。Triton 提供了 block-level 的抽象，开发者可以专注于算子的逻辑实现，而不需要手动管理线程块、共享内存等底层细节。Triton 的 autotuner 可以自动搜索最优的内核配置参数。
本文的核心贡献之一就是使用 Triton 编写了融合内核，将 DoRA 组合操作从四个 CUDA 内核调用合并为一次传递。理解 Triton 的作用有助于明白为什么融合内核能显著减少内存流量并提升性能。
梯度检查点 (Gradient Checkpointing)
梯度检查点是一种用计算换内存的技术，用于降低大模型训练时的内存占用。在前向传播时只保存部分中间激活值，其他激活值在反向传播时重新计算。这可以大幅减少峰值内存使用，但代价是额外的计算开销。在 Transformer 等深层网络中，梯度检查点通常按层为单位进行，每 $k$ 层保存一次检查点。
本文指出 DoRA 的临时内存分配在梯度检查点时会加倍（每个模块重新计算两次），这使得内存压力更加严重。理解梯度检查点的工作机制有助于明白为什么在高秩 DoRA 配置下，即使单模块的 512MB 临时内存也会累积成 OOM 问题。
研究动机
DoRA (Weight-Decomposed Low-Rank Adaptation) 虽然在参数高效微调中表现出色，但其在高秩配置下的前向传播需要计算行向范数 $\|W + sBA\|_{row}$。现有主流框架（包括 Hugging Face PEFT、torchtune、Unsloth、SWIFT、LLaMA-Factory、Axolotl）都通过构造单位矩阵来实现，这会材料化密集的 BA 乘积。在具体数据场景下，当维度为 8192、rank 为 384、使用 bf16 精度时，单个模块就需要约 512MB 的瞬态工作内存。更严重的是，在梯度检查点场景下，这些临时内存会被分配两次。当模型规模达到 8-32B 参数、包含数百个适配模块时，这种累积的内存压力不仅导致速度显著下降，还经常引发 OOM (Out of Memory) 失败。作者测量发现，HF PEFT 的 DoRA 实现在 32B 模型上比优化后的实现慢 46-87%，且需要多使用 1.3-6.7GB 的峰值显存。
本文的目标是本文的目标是在不改变 DoRA 算法本身的前提下，通过系统层面的优化实现更高效的工作集和更低的内存流量。具体来说，作者希望在保持数学计算完全一致的前提下，显著减少 DoRA 计算过程中的瞬态内存占用和内存带宽消耗，使高秩 DoRA 配置能够在常见的单 GPU 设置上可行。同时，优化后的实现需要保持数值精度和训练收敛性，与原始实现等价。
与已有工作不同的是，本文的独特切入角度是从系统优化而非算法改进的角度来解决问题。大多数关于 LoRA/DoRA 的工作关注改进适配器架构、优化器或训练策略，而本文专注于如何更高效地执行相同的 DoRA 计算。作者指出了一个被广泛忽视的问题：现有框架都使用相同的低效实现模式，而这个问题的解决方案并不是更聪明的算法，而是更聪明的计算方式。通过代数分解和内核融合，作者在不改变数学定义的情况下实现了显著的性能提升。这种系统层面的优化视角在参数高效微调领域相对少见，为后续研究提供了新的思路。
核心方法
本文的方法核心是将 DoRA 的计算过程进行两个关键优化：因式分解范数计算和融合 Triton 内核。直觉上，DoRA 的瓶颈在于需要计算 $\|W + sBA\|_{row}$，而现有方法通过构造单位矩阵来材料化完整的 BA 乘积，这在维度较大时非常低效。作者的思路是通过代数恒等式将范数平方分解为三个可独立计算的部分（基础范数、交叉项、BA 范数），每个部分都可以通过低秩中间量计算，完全避免材料化密集矩阵。然后，作者使用 Triton 语言编写融合内核，将 DoRA 组合操作从四个独立的 CUDA 内核调用合并为一次传递，大幅减少内存流量。整个系统还包括一个三层运行时分发机制，根据训练/推理场景、GPU 架构和激活大小自动选择最优的计算路径。
本文的核心创新点有两个：首先是因式分解范数的代数恒等式 $\|W + sBA\|_{row}^2 = \|W\|_{row}^2 + 2s\langle W, BA \rangle_{row} + s^2 \|BA\|_{row}^2$，其中行向内积。这个恒等式将原本需要计算的范数分解为三个项，每项都可以通过低秩中间量计算：基础范数可以通过分块累积计算；交叉项可以重写为 $(B \odot U) \cdot \mathbf{1}$，其中 $U = WA^\top$ 通过分块累积；BA 范数可以通过 Gram 矩阵 $G = AA^\top$ 计算，即 $(BG \odot B) \cdot \mathbf{1}$。其次是融合 Triton 内核，将 DoRA 组合的四步元素操作合并为单次传递，使得每个元素只需 3 次读取（base、lora、g）和 1 次写入，相比 eager 路径的约 12 次内存访问减少了约 4 倍的内存流量。这些创新与已有方法的本质区别在于，它们不改变 DoRA 的数学定义，而是通过更聪明的计算方式实现相同的数学结果。
方法步骤详情
方法的第一步是因式分解范数计算。输入是冻结权重 W、低秩因子 A 和 B、缩放系数 s，以及分块预算参数。具体操作包括：确定分块大小；初始化基础平方向量、交叉项向量、Gram 矩阵，所有累积都在 fp32 精度下进行；对每个分块，计算权重和 A 的分块并累积到基础平方、Gram 矩阵和交叉项中；最终计算 BA 平方和权重范数。第二步是融合 Triton 内核。DoRA 组合操作在标准 PyTorch 中需要四次元素操作，每次都启动独立的 CUDA 内核。融合内核将这四步合并为一次 Triton 内核调用，每个程序处理一个 tile 的数据，直接计算最终输出。为了数值稳定性，使用稳定形式而非会灾难性抵消的形式。反向内核同时计算 lora 和 base 的梯度。第三步是三层运行时分发。第一层是融合反向路径，用于训练且需要梯度时，输出和保存的张量在一次传递中计算；第二层是融合前向路径，用于推理时，没有自动微分图节点；第三层是 eager 回退，用于 CPU、无 Triton、形状小于交叉点等场景。
技术新颖性
本文的技术新颖性体现在多个方面。首先是代数恒等式的系统性应用，虽然将平方范数展开为基础、交叉和 Gram 项在数值线性代数中是标准技术，但将其应用于 DoRA 特定的范数计算并配合精确的 dtype 纪律、分块和融合管线是新颖的。其次是完整的内核融合方案，包括前向、反向和范数组装三个 Triton 内核，以及针对近统一缩放 regime 的数值稳定性处理。作者发现在实际训练中 100% 的 g 值落在 bf16 抵消区，因此数值稳定性不是理论问题而是实际需要。第三是三层运行时分发系统，根据训练/推理场景、GPU 架构和激活大小动态选择最优路径，这是对简单融合策略的重要改进。第四是跨 GPU 架构的一致性验证，作者在六款 GPU 上进行了详细的微基准测试和模型级基准测试，从 GDDR6 到 HBM3e，确认性能提升确实来自减少内存流量而非架构特定优化。这种系统性的验证在系统优化论文中相对少见。最后，作者还对数值精度进行了多层次验证，包括算子级测试、最终 logit 余弦相似度和多种子训练收敛曲线，确保优化不改变训练动力学。
实验结果
核心发现包括三个方面：性能提升、内存节省和数值精度。在性能方面，融合实现在六款 8-32B 视觉-语言模型上的推理速度比 HF PEFT 的 DoRA 实现快 1.5-2.0 倍，梯度计算（不包括优化器步骤）快 1.5-1.9 倍。具体数据上，在 H200 GPU 上，Qwen3-VL-32B 的推理速度提升 1.74 倍，梯度计算速度提升 1.66 倍；在 B200 GPU 上，Qwen2.5-VL-32B 的推理速度提升 2.03 倍，梯度计算速度提升 1.74 倍。在 RTX 6000 PRO 上，所有六个模型都能运行推理，但在梯度计算时 32B 模型会 OOM。微基准测试确认了组合内核加速 1.5-2.7 倍，带宽利用率达到峰值的约 50-55%。在内存节省方面，融合实现在模型级别使用比 eager 基线少 0.1-1.0 GB 的峰值显存，比 PEFT 少 1.3-6.7 GB。在高秩配置下，速度提升 vs. PEFT 随 rank 增长而增长，因为 PEFT 的材料化成本随 rank 增长。在数值精度方面，融合实现和 eager 实现的最终 logit 余弦相似度在所有模型/GPU 对上超过 0.9999。多种子收敛实验显示平均每步损失差异为极小值，训练曲线在视觉上无法区分。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| 推理 | 速度提升倍数 vs. HF PEFT DoRA | 1.5-2.0× | 1.0× (HF PEFT) | 最高 2.03× (Qwen2.5-VL-32B on B200) |
| 梯度计算 | 速度提升倍数 vs. HF PEFT DoRA | 1.5-1.9× | 1.0× (HF PEFT) | 最高 1.87× (Qwen2.5-VL-32B on H200) |
| 梯度计算 | 速度提升倍数 vs. eager 基线 | 1.18-1.24× | 1.0× (eager) | 最高 1.24× (Qwen3-VL-8B on B200) |
| 模型推理 | 峰值 VRAM (GB) - Qwen3-VL-8B | 29.5 GB | 30.7 GB (PEFT) | 降低 1.2 GB |
| 模型梯度计算 | 峰值 VRAM (GB) - Mistral-Sm-24B | 70.6-70.7 GB | 77.3-77.4 GB (PEFT) | 降低 6.7 GB |
| 组合内核 | 带宽利用率 (占峰值百分比) | 50-55% | ~20% (eager) | 约 2.5-2.75× |
| 训练收敛 | 平均每步损失差异 (多种子) | 极小 | 0 (理论) | 差异极小，训练等价 |
| 训练收敛 | 最终 logit 余弦相似度 | > 0.9999 | 1.0 (理论) | 数值精度损失可忽略 |
| 范数计算内存 | 测量减少倍数 (8192×8192, r=512) | 3.2× | 1.0× (PEFT) | 理论可达 15.1× |
| 范数计算内存 | 测量减少倍数 (8192×28672 MoE, r=384) | 11.0× | 1.0× (PEFT) | 对 MoE 架构特别有效 |
局限与改进
作者明确承认的局限性包括：FSDP2/DTensor 不支持，因为因式分解范数假设可以访问完整的基础权重 W，而 FSDP2 将权重分片到不同设备；收敛验证仅覆盖两个模型家族、两个优化器和一个 SFT 数据集，对 RL 管道的泛化性仍需确认；模型级基准测试只覆盖三款 GPU，其他只有微基准覆盖；分发交叉点选择基于经验启发式，可能需要针对未来硬件重新调优。我观察到作者没有提到但对实际应用可能有影响的局限性包括：当 rank 很小且维度不大时，因式分解范数的计算开销可能接近或超过其内存节省带来的收益；在 bf16 精度下，因式分解范数由于 fp32 累积临时量，在孤立范数微操作上的内存使用实际上比 PEFT 多，虽然这不影响模型级别的 VRAM 节省；融合内核的性能提升在大激活尺寸时最为明显，在小激活尺寸时启动延迟占主导；代码需要 Triton 和特定版本的 PyTorch/CUDA，在老硬件或非 CUDA 平台上会自动回退到 eager 路径；虽然作者提供了广泛的兼容性测试，但对更复杂的分布式训练设置的兼容性未充分验证。
独立分析的弱点
独立分析的弱点首先是小激活场景下的性能倒退。当激活尺寸小于约特定阈值时，融合内核的启动延迟占主导，可能比 eager 路径慢。这在某些应用场景中可能成为问题。改进方向可以是实现动态自适应的交叉点选择，根据实际测量到的启动延迟和计算时间自动调整分发阈值。其次是 fp32 累积的计算开销。因式分解范数在 bf16 权重下仍然使用 fp32 累积，这增加了计算开销。改进方向可以是引入混合精度策略。第三是向后兼容性问题。作者发现 PEFT 的嵌入路径只计算部分公式而忽略了另一部分，这意味着用 PEFT 训练的检查点需要重新微调或提供遗留组合回退。第四是有限的 FSDP2 支持。随着 FSDP2 在大模型训练中的普及，不支持它是一个显著限制。改进方向是实现分布式累积策略。第五是微基准覆盖不完整，缺少 A100 这个广泛使用的数据中心 GPU 的端到端数据。
未来方向
作者提出的未来工作方向包括：缓存基础范数，因为 W 是冻结的，可以预计算到缓冲区，这可以完全消除基础范数计算的临时内存；扩展 FSDP2/DTensor 支持，需要实现分布式累积策略；探索与 Meta KernelAgent 生成的双阶段部分约简策略的集成。基于成果可延伸的未来研究方向包括：将因式分解范数的思想应用于其他需要范数计算的适配器变体；探索更激进的内核融合，如将 DoRA 组合与下游操作融合在一起；研究自适应 rank 策略，根据层的维度和重要性动态调整 rank；扩展到其他模态和架构，验证方法的泛化性；开发自动调优工具，根据具体模型、硬件和工作负载自动选择最优的范数路径、融合策略和交叉点参数；研究在边缘设备上的优化；探索与量化、剪枝等模型压缩技术的结合。作者还提到嵌入公式修正需要进一步工作。
复现评估
复现评估方面，作者提供了完整的开源实现和实验设置，复现难度中等。所有源代码、基准测试脚本、原始 JSON 结果、Triton autotune 缓存和图表生成脚本都在 GitHub 上可用。收敛验证使用公开数据集，确保完全可复现。软件栈包括特定版本的 PyTorch、Triton、Transformers、CUDA，这些都是相对较新的版本，但主要云服务商和 HPC 环境应该能够满足。硬件方面，模型级基准需要 96GB 以上的显存或 141GB 以上的显存，这对研究者来说门槛较高，但微基准可以在更小的 GPU 上运行。每个微基准报告 200 次 CUDA-event 计时试验的中位数，模型级基准使用 20 次重复，统计方法严格。总的来说，虽然需要较高配置的硬件，但作者提供了完整的代码和数据，严格的统计方法，以及详细的环境说明，复现应该是可行的。
论文图表
