Scaling DoRA: Adattamento ad Alto Rango tramite Norme Fattorizzate e Kernel Fusi

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

March 23, 2026

Autori: Alexandra Zelenin, Alexandra Zhuravlyova

cs.AI

Abstract

La Weight-Decomposed Low-Rank Adaptation (DoRA) estende LoRA disaccoppiando la magnitudine dei pesi dalla loro direzione, ma il suo passaggio in avanti richiede la norma riga per riga di W + sBA, un calcolo che ogni framework principale da noi esaminato implementa materializzando il prodotto denso [d_out, d_in] BA. Con d_in = 8192 e rango r = 384, la norma di un singolo modulo richiede circa 512 MB di memoria di lavoro transitoria in bf16, rendendo DoRA ad alto rango costoso e spesso infattibile sulle comuni configurazioni single-GPU quando sono coinvolte centinaia di moduli adattati e il checkpointing. Presentiamo due contributi sistemici. Una *norma fattorizzata* scompone la norma al quadrato in termini di base, incrociati e di Gram calcolabili tramite intermediari di complessità O(d_out r + r^2), eliminando il prodotto denso. Kernel Triton fusi condensano la composizione in quattro kernel di DoRA in un singolo passaggio, riducendo il traffico di memoria di circa 4x e utilizzando una forma numericamente stabile che evita la cancellazione catastrofica nel regime di riscalamento quasi unitario dove, in pratica, si concentrano le scale di magnitudine. Su sei modelli visione-linguaggio (VLM) da 8 a 32B, eseguiti su tre GPU NVIDIA (RTX 6000 PRO, H200, B200) con r = 384 in bf16, l'implementazione fusa è 1.5-2.0x più veloce dell'implementazione DoRA di Hugging Face PEFT per l'inferenza e 1.5-1.9x più veloce per il calcolo del gradiente (escluso lo step dell'ottimizzatore), con un picco di VRAM inferiore fino a 7 GB. Microbenchmark su sei GPU che coprono quattro generazioni architetturali (L40S, A100, RTX 6000 PRO, H200, B200, B300) confermano un speedup del kernel di composizione di 1.5-2.7x. La similarità coseno sui logit finali supera 0.9999 su tutte le coppie modello/GPU, e le curve di addestramento multi-seed corrispondono entro un delta di loss medio per step di 7.1 x 10^-4 su 2000 step.

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.