03 {多模态} DoRA 的内存墙在高秩时突破:系统层面的修复,而非数学层面的修复
DoRA(权重分解低秩适应)将权重矩阵分解为幅度和方向分量,这使其在微调质量上优于标准的 LoRA(低秩适应)。问题在于:DoRA 所需的逐行范数计算迫使每个框架都必须实例化完整的稠密乘积 BA——一个形状为 [d_out × d_in] 的矩阵。当 d_in = 8,192 且秩 r = 384 时,单个范数计算会在 bf16 格式下,为每个模块消耗约 512 MB 的瞬态工作内存。如果将此内存消耗乘以数百个经过适应的模块,并考虑梯度检查点,那么在开始实际训练逻辑之前,单 GPU 微调运行就已经变得不可行了。
此修复是代数层面的,而非架构层面的。W + sBA 的平方逐行范数可以分解为三项:基项、交叉项和格拉姆项,每项都可以通过 O(d_out × r + r²) 的中间值计算,而非 O(d_out × d_in)。稠密的实例化完全消失了。融合的 Triton 核函数(消除冗余内存读写的底层 GPU 计算核函数)将 DoRA 的四核函数前向传播合并为单一操作,同时减少了内存移动和核函数启动开销。这两项贡献是独立的:因子分解范数消除了内存峰值;融合核函数减少了延迟。它们共同使得高秩 DoRA(即 DoRA 相对于 LoRA 的质量优势最显著的范畴)在标准的单 GPU 设置上变得可行。
局限性在于其范围。这是一篇系统论文。它没有在高秩下进行下游微调质量比较,以展示当 r 扩展到 384 时,DoRA 相对于 LoRA 的质量优势是否依然存在。这方面的论证仍依赖于之前的 DoRA 结果。Triton 核函数也是硬件特定的;非 NVIDIA 硬件的团队将需要进行移植或近似实现。对于那些已经在中等秩下使用 DoRA,但在提高秩时遇到内存墙的实践者来说,这是一个直接的解决方案,而不是一个从 LoRA 切换的理由,如果 DoRA 的质量提升本身还不足以激励他们的话。
主要收获:
- 因子分解范数消除了 DoRA 前向传播中稠密的 [d_out × d_in] 矩阵实例化,将 O(d_out × d_in) 的峰值内存替换为 O(d_out × r + r²) 的中间值,这在高秩单 GPU 上标志着可行与不可行之间的区别。
- 高秩 DoRA 中的内存墙是一个系统层面的“人造物”,而非权重分解适应的根本约束。数学上它始终是可分解的;只是没有任何框架以这种方式实现它。
- 使用 DoRA 微调大型模型,并因内存限制而只能在低秩下工作的团队,应直接测试因子分解范数的实现;这是在单 GPU 设置上实现高秩 DoRA 的途径,无需转向多 GPU 并行。
Source: Scaling DoRA: High-Rank Adaptation via Factored Norms and Fused Kernels
Source: Scaling DoRA: High-Rank Adaptation via Factored Norms and Fused Kernels