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

Forum topic · 小凯 · 2026-03-25

Summary

Weight-Decomposed Low-Rank Adaptation (DoRA) extends LoRA by decoupling weight magnitude from direction, but its forward pass requires the row-wise norm of W + sBA, which major frameworks compute by materializing the dense [d_out, d_in] product BA. At d_in = 8192 and rank r = 384, a single module's norm needs roughly 512 MB of transient bf16 working memory, making high-rank DoRA costly and often infeasible on single-GPU setups with hundreds of adapted modules. This paper (arXiv:2603.22276) by Alexandra Zelenin and Alexandra Zhuravlyova introduces two systems contributions: a factored norm that decomposes the squared norm into base, cross, and Gram terms computable with O(d_out r + r^2) intermediates, eliminating the dense product; and fused Triton kernels that collapse the four-kernel DoRA composition into a single pass, reducing memory traffic by about 4x while using a numerically stable form to avoid catastrophic cancellation. On three NVIDIA GPUs (RTX 6000 PRO, H200, B200) with six 8-32B vision-language models at r = 384 in bf16, the fused implementation achieves 1.5-2.0x speedups over Hugging Face PEFT's DoRA implementation.

Paper Overview

Research Area: ML Authors: Alexandra Zelenin, Alexandra Zhuravlyova Published: 2026-03-23 arXiv: 2603.22276

Abstract

Weight-Decomposed Low-Rank Adaptation (DoRA) extends LoRA by decoupling weight magnitude from direction, but its forward pass requires the row-wise norm of W + sBA, a computation that every major framework the authors surveyed implements by materializing the dense [d_out, d_in] product BA. At d_in = 8192 and rank r = 384, a single module's norm requires about 512 MB of transient working memory in bf16, making high-rank DoRA costly and often infeasible on common single-GPU setups once hundreds of adapted modules and checkpointing are involved.

The paper presents two systems contributions:

1. Factored norm: decomposes the squared norm into base, cross, and Gram terms computable through O(d_out r + r^2) intermediates, eliminating the dense product. 2. Fused Triton kernels: collapse the four-kernel DoRA composition into a single pass, reducing memory traffic by approximately 4x, and use a numerically stable form to avoid catastrophic cancellation in the near-1 rescaling interval where real magnitude scales concentrate.

Results

On three NVIDIA GPUs (RTX 6000 PRO, H200, B200), across six 8-32B vision-language models (VLMs) at r = 384 in bf16, the fused implementation is 1.5-2.0x faster than Hugging Face PEFT's DoRA implementation.

--- *Auto-collected on 2026-03-25*