Escalonamento de DoRA: Adaptação de Alto Rank via Normas Fatoradas e Kernels Fundidos

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

March 23, 2026

Autores: Alexandra Zelenin, Alexandra Zhuravlyova

cs.AI

Resumo

A Adaptação de Baixa Posto com Decomposição de Peso (DoRA) estende o LoRA ao desacoplar a magnitude do peso da sua direção, mas o seu *forward pass* requer a norma linha a linha de W + sBA, um cálculo que todos os principais *frameworks* que pesquisamos implementam materializando o produto denso [d_out, d_in] BA. Com d_in = 8192 e posto r = 384, a norma de um único módulo requer cerca de 512 MB de memória de trabalho transitória em bf16, tornando a DoRA de alto posto dispendiosa e muitas vezes inviável em configurações comuns de GPU única quando centenas de módulos adaptados e *checkpointing* estão envolvidos. Apresentamos duas contribuições de sistema. Uma **norma fatorada** decompõe a norma ao quadrado em termos de base, cruzado e de Gram, calculáveis através de intermediários de O(d_out r + r²), eliminando o produto denso. **Kernels Triton fundidos** colapsam a composição de quatro *kernels* da DoRA em uma única passada, reduzindo o tráfego de memória em cerca de 4x e usando uma forma numericamente estável que evita o cancelamento catastrófico no regime de reescalonamento próximo à unidade, onde as escalas de magnitude se concentram na prática. Em seis modelos visão-linguagem (VLMs) de 8-32B, executados em três GPUs NVIDIA (RTX 6000 PRO, H200, B200) com r = 384 em bf16, a implementação fundida é 1.5-2.0x mais rápida que a implementação DoRA do Hugging Face PEFT para inferência e 1.5-1.9x mais rápida para o cálculo do gradiente (etapa do otimizador excluída), com um pico de VRAM até 7 GB menor. *Microbenchmarks* em seis GPUs abrangendo quatro gerações de arquitetura (L40S, A100, RTX 6000 PRO, H200, B200, B300) confirmam uma aceleração de 1.5-2.7x no *kernel* de composição. A similaridade do cosseno nos *logits* finais excede 0.9999 em todos os pares modelo/GPU, e as curvas de treinamento com múltiplas sementes coincidem dentro de um delta de perda médio por etapa de 7.1 x 10^-4 ao longo de 2000 etapas.

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.