Mise à l'échelle de DoRA : Adaptation de haut rang via normes factorisées et noyaux fusionnés

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

March 23, 2026

Auteurs: Alexandra Zelenin, Alexandra Zhuravlyova

cs.AI

Résumé

L'adaptation à faible rang avec décomposition de poids (DoRA) étend LoRA en découplant la magnitude des poids de leur direction, mais son passage avant nécessite la norme ligne par ligne de W + sBA, un calcul que chaque framework majeur que nous avons étudié implémente en matérialisant le produit dense [d_out, d_in] BA. Avec d_in = 8192 et un rang r = 384, la norme d'un seul module nécessite environ 512 Mo de mémoire de travail transitoire en bf16, rendant DoRA à haut rang coûteux et souvent irréalisable sur les configurations courantes à un seul GPU lorsque des centaines de modules adaptés et la mise en point de contrôle sont impliqués. Nous présentons deux contributions systèmes. Une norme factorisée décompose la norme au carré en termes de base, croisés et de Gram calculables via des intermédiaires de complexité O(d_out r + r^2), éliminant ainsi le produit dense. Des noyaux Triton fusionnés réduisent la composition DoRA à quatre noyaux en une seule passe, réduisant le trafic mémoire d'environ 4x et utilisant une forme numériquement stable qui évite l'annulation catastrophique dans le régime de re-redimensionnement proche de l'unité où les échelles de magnitude se concentrent en pratique. Sur six modèles vision-langage (VLM) de 8 à 32B testés sur trois GPU NVIDIA (RTX 6000 PRO, H200, B200) à r = 384 en bf16, l'implémentation fusionnée est 1,5 à 2,0 fois plus rapide que l'implémentation DoRA de Hugging Face PEFT pour l'inférence et 1,5 à 1,9 fois plus rapide pour le calcul du gradient (étape de l'optimiseur exclue), avec un pic de mémoire VRAM inférieur jusqu'à 7 Go. Des micro-benchmarks sur six GPU couvrant quatre générations d'architecture (L40S, A100, RTX 6000 PRO, H200, B200, B300) confirment une accélération des noyaux de composition de 1,5 à 2,7x. La similarité cosinus des logits finaux dépasse 0,9999 pour toutes les paires modèle/GPU, et les courbes d'entraînement multi-graines correspondent avec un delta de perte moyen par pas inférieur à 7,1 x 10^-4 sur 2000 pas.

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.