Scaling DoRA: High-Rank Adaptation via Factored Norms and Fused Kernels

Paper: 2603.22276 Authors: Alexandra Zelenin, Alexandra Zhuravlyova Categories: cs.LG

Abstract

Weight-Decomposed Low-Rank Adaptation (DoRA) extends LoRA by decoupling weight magnitude from direction, but its forward pass requires computing a row-wise norm that every major framework implements by materializing the dense product BA — consuming ~512 MB of transient memory per module at typical dimensions. This paper presents two systems contributions: (1) a factored norm that decomposes the squared norm into base, cross, and Gram terms computable through O(d_out·r + r²) intermediates, eliminating the dense product; and (2) fused Triton kernels that collapse four CUDA kernel launches into a single pass with ~4x reduced memory traffic. Across six 8-32B vision-language models on three NVIDIA GPUs at rank 384, the implementation achieves 1.5-2.0x inference speedup and 1.5-1.9x faster gradient computation with up to 7 GB lower peak VRAM.

Key Contributions

• Factored norm computation: Algebraic decomposition of ∥W + sBA∥² into three terms (base, cross, BA norm), each computable without materializing the full dense matrix
• Fused Triton kernels: Single-pass kernel replacing four separate CUDA launches, with numerically stable form avoiding catastrophic cancellation in the near-unity rescaling regime
• Three-tier runtime dispatch: Automatically selects optimal path (fused backward for training, fused forward for inference, eager fallback for CPU), compatible with torch.compile, gradient checkpointing, DeepSpeed ZeRO, and FSDP
• Comprehensive validation: Tested on 6 GPUs spanning 4 architecture generations (L40S through B300) with final-logit cosine similarity exceeding 0.9999

The Problem

DoRA’s forward pass requires computing:

$W' = m \odot \frac{W + sBA}{\|W + sBA\|_{\text{row}}}$

Every major framework (HF PEFT, torchtune, Unsloth, SWIFT, LLaMA-Factory, Axolotl) computes the row-wise norm by constructing an identity matrix and materializing the full dense product BA. At d_in = 8192 and rank r = 384, a single module allocates 3-4 dense temporaries totaling ~512 MB in bf16. With gradient checkpointing, these are allocated twice per step. Across hundreds of adapted modules in an 8-32B model, this causes severe speed degradation and OOM failures at high rank.

Factored Norm

The key insight is that the row-wise squared norm decomposes as:

$\|W + sBA\|^2_{\text{row}} = \|W\|^2_{\text{row}} + 2s\langle W, BA\rangle_{\text{row}} + s^2\|BA\|^2_{\text{row}}$

Each term can be computed through low-rank intermediates:

• Base norm: ∥W∥² via chunked accumulation along d_in
• Cross term: Row-wise inner product via U = WA^T ∈ R^(d_out × r), then element-wise product with B
• BA norm: Via Gram matrix G = AA^T ∈ R^(r × r) (only 1 MB at r = 512)

Theoretical persistent-memory reduction: 15x at d = 8192, r = 512.

Fused Kernels

The DoRA composition (g-1) ⊙ base + g ⊙ s ⊙ lora normally requires four separate CUDA kernel launches. The fused Triton kernel performs this in a single pass, using a numerically stable form that avoids catastrophic cancellation when g ≈ 1 (where magnitude scales concentrate in practice).

Results

Metric	Improvement
Inference speedup	1.5-2.0x over HF PEFT
Gradient computation	1.5-1.9x faster
Peak VRAM reduction	Up to 7 GB
Compose-kernel speedup	1.5-2.7x (microbenchmarks)
Numerical fidelity	Cosine similarity > 0.9999
Training match	< 7.1×10⁻⁴ mean per-step loss delta over 2000 steps

Takeaways

• High-rank DoRA has been impractical on single-GPU setups due to dense matrix materialization — this work makes it feasible
• The factored norm is a purely algebraic insight: decompose the squared norm to avoid ever constructing the full BA product
• Fused kernels address memory traffic, not just compute — collapsing multiple passes into one
• The approach is a drop-in replacement requiring no changes to the adapter architecture, optimizer, or training recipe
• All six major PEFT frameworks surveyed use the same suboptimal identity-matrix approach, suggesting broad applicability of this fix