Schalen van DoRA: Hoog-rang Adaptatie via Gefactoriseerde Normen en Gefuseerde Kernels

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

March 23, 2026

Auteurs: Alexandra Zelenin, Alexandra Zhuravlyova

cs.AI

Samenvatting

Gewichtsgedecomponeerde Low-Rank-aanpassing (DoRA) breidt LoRA uit door de gewichtsgrootte los te koppelen van de richting, maar de voorwaartse doorgang vereist de rijgewijze norm van W + sBA, een berekening die elk groot framework dat we onderzochten implementeert door het dichte [d_out, d_in] product BA te materialiseren. Bij d_in = 8192 en rang r = 384 vereist de norm van een enkele module ongeveer 512 MB tijdelijk werkgeheugen in bf16, wat DoRA met hoge rang kostbaar en vaak onhaalbaar maakt op gangbare single-GPU-opstellingen zodra honderden aangepaste modules en checkpointing betrokken zijn. Wij presenteren twee systeembijdragen. Een gefactoriseerde norm decomposeert de kwadratische norm in basis-, kruis- en Gram-termen die berekenbaar zijn via O(d_out r + r^2) tussenproducten, waardoor het dichte product wordt geëlimineerd. Gefuseerde Triton-kernels vouwen de vier-kernel DoRA-compositie samen tot een enkele doorgang, waardoor het geheugenverkeer met ongeveer 4x wordt verminderd en gebruik wordt gemaakt van een numeriek stabiele vorm die catastrofale cancelatie vermijdt in het near-unity-herschalingsregime waar grootte-schalingen zich in de praktijk concentreren. Over zes 8-32B vision-language modellen (VLM's) op drie NVIDIA GPU's (RTX 6000 PRO, H200, B200) bij r = 384 in bf16, is de gefuseerde implementatie 1.5-2.0x sneller dan Hugging Face PEFT's DoRA-implementatie voor inferentie en 1.5-1.9x sneller voor gradiëntberekening (optimizer-stap uitgesloten), met tot 7 GB lagere piek-VRAM. Microbenchmarks op zes GPU's over vier architectuurgeneraties (L40S, A100, RTX 6000 PRO, H200, B200, B300) bevestigen een 1.5-2.7x versnelling van de compose-kernel. De cosinusgelijkenis van de laatste logits overschrijdt 0.9999 over alle model/GPU-paren, en trainingscurves met meerdere seeds komen overeen binnen een gemiddeld verliesverschil per stap van 7.1 x 10^-4 over 2000 stappen.

English

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 we 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. We present two systems contributions. A 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. Fused Triton kernels collapse the four-kernel DoRA composition into a single pass, reducing memory traffic by about 4x and using a numerically stable form that avoids catastrophic cancellation in the near-unity rescaling regime where magnitude scales concentrate in practice. Across six 8-32B vision-language models (VLMs) on three NVIDIA GPUs (RTX 6000 PRO, H200, B200) at r = 384 in bf16, the fused implementation is 1.5-2.0x faster than Hugging Face PEFT's DoRA implementation for inference and 1.5-1.9x faster for gradient computation (optimizer step excluded), with up to 7 GB lower peak VRAM. Microbenchmarks on six GPUs spanning four architecture generations (L40S, A100, RTX 6000 PRO, H200, B200, B300) confirm 1.5-2.7x compose-kernel speedup. Final-logit cosine similarity exceeds 0.9999 across all model/GPU pairs, and multi-seed training curves match within 7.1 x 10^-4 mean per-step loss delta over 2000 steps.