Escalado de DoRA: Adaptación de Alto Rango mediante Normas Factorizadas y Núcleos Fusionados

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

March 23, 2026

Autores: Alexandra Zelenin, Alexandra Zhuravlyova

cs.AI

Resumen

La Adaptación de Bajo Rango con Descomposición de Peso (DoRA) extiende LoRA al desacoplar la magnitud del peso de su dirección, pero su paso hacia adelante requiere la norma por filas de W + sBA, un cálculo que todos los principales frameworks que estudiamos implementan materializando el producto denso [d_out, d_in] BA. Con d_in = 8192 y un rango r = 384, la norma de un solo módulo requiere aproximadamente 512 MB de memoria de trabajo transitoria en bf16, lo que hace que DoRA de alto rango sea costoso y a menudo inviable en configurasiones comunes de una sola GPU cuando se involucran cientos de módulos adaptados y checkpointing. Presentamos dos contribuciones de sistemas. Una norma factorizada descompone la norma al cuadrado en términos base, cruzados y de Gram calculables mediante intermediarios de O(d_out r + r^2), eliminando el producto denso. Kernels de Triton fusionados colapsan la composición de DoRA de cuatro kernels en una sola pasada, reduciendo el tráfico de memoria aproximadamente 4 veces y utilizando una forma numéricamente estable que evita la cancelación catastrófica en el régimen de reescalado cercano a la unidad, donde las escalas de magnitud se concentran en la práctica. En seis modelos de visión y lenguaje (VLM) de 8 a 32B probados en tres GPUs NVIDIA (RTX 6000 PRO, H200, B200) con r = 384 en bf16, la implementación fusionada es 1.5-2.0 veces más rápida que la implementación DoRA de Hugging Face PEFT para inferencia y 1.5-1.9 veces más rápida para el cálculo del gradiente (excluyendo el paso del optimizador), con un pico de VRAM hasta 7 GB menor. Los microbenchmarks en seis GPUs que abarcan cuatro generaciones de arquitecturas (L40S, A100, RTX 6000 PRO, H200, B200, B300) confirman una aceleración de 1.5-2.7 veces en el kernel de composición. La similitud del coseno en los logits finales supera 0.9999 en todos los pares modelo/GPU, y las curvas de entrenamiento multi-semilla coinciden dentro de un delta de pérdida medio por paso de 7.1 x 10^-4 a lo largo de 2000 pasos.

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.