[논문 리뷰] Scaling DoRA: High-Rank Adaptation via Factored Norms and Fused Kernels

Alexandra Zelenin, Alexandra Zhuravlyova|arXiv (Cornell University)|2026. 03. 23.

Parallel Computing and Optimization Techniques인용 수 0

한 줄 요약

논문은 고랭크 DoRA를 위한 물질화된 Dense B@A를 피하기 위한 인자화(norm) 접근법과 전방/역전 파스를 가속화하고 VRAM 사용을 줄이는 융합 Triton 커널 구현을 제시하며, 여러 GPU에서 6개의 비전-언어 모델로 검증합니다.

ABSTRACT

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.

연구 동기 및 목표

• 고랭크 DoRA의 매개변수 효율적 미세조정에서의 메모리 및 속도 병목 현상을 동기 부여한다.
• Dense B@A 물질화 없이 ||W + sBA||_row^2의 메모리 효율적 인자화(norm) 공식을 개발한다.
• forward 및 backward 계산을 하나의 패스로 축소하는 융합 Triton 커널을 설계한다.
• 런타임 디스패치 전략 및 일반적인 PEFT 프레임워크와 분산 학습 체계와의 호환성을 제공한다.
• 다양한 GPU 아키텍처 및 모델 규모에 걸친 메모리, 속도, 충실도 및 수렴을 경험적으로 검증한다.

제안 방법

• 행렬 B@A를 물질화하지 않고도 행 단위 제곱합의 기저(base), 교차(cross), BA-노름 항으로 분해 가능한 인자화 분해를 도출한다(복잡도 O(d_out r + r^2)).
• 정밀도를 관리하기 위해 명시적 형변환과 함께 청킹된 fp32 누적으로 행별 노름을 구성한다(식 2-6).
• DoRA 파이프라인을 하나의 패스로 융합하기 위해 쿼트럴(전방/역방), 노름 조립 커널, 합성 커널의 네 가지 커널을 Triton으로 구현하여 수치적 안정성을 확보한다.
• 하드웨어와 형태에 따라 융합된 역전(학습), 융합된 전방(추론) 또는 즉시(fallback) 경로를 선택하는 3계층 런타임 디스패치를 제공한다.
• torch.compile, 그래디언트 체크포인팅, Deepspeed ZeRO, FSDP와의 호환성을 보장하고 정밀도를 유지하기 위한 명시적 규모-나눗셈 처리를 포함한다.
• 6개의 GPU 및 6개의 8–32B 비전-언어 모델에 대한 상세 마이크로벤치마크 및 모델 수준 벤치마크를 수행한다.

Figure 1 : The stable compose form achieves $3.0$ × lower peak error near $g\approx 1$ (bf16, $d_{\text{out}}=8192$ , $d_{\text{in}}=2048$ ). The naive form $g\odot(s\cdot\text{lora}+\text{base})-\text{base}$ exhibits catastrophic cancellation; the stable form and fused kernel both remain near the b

실험 결과

연구 질문

• RQ1 DoRA의 행 기준 노름을 Dense B@A 물질화 없이 계산할 수 있어 고랭크 적합화를 확장 가능하게 만드는가?
• RQ2 융합 커널이 DoRA의 전방 및 역방 통과에서 메모리 트래픽을 줄이고 다양한 GPU 및 모델 크기에서 처리량을 향상시키는가?
• RQ3 인자화된-노름 접근이 표준 DoRA 구현과 비교해 수치적 안정성 및 학습 수렴에 미치는 영향은 무엇인가?
• RQ4 제안된 시스템이 Fidelity를 보존하면서 기존 PEFT 프레임워크 및 분산 학습 도구 체인과 어떻게 통합되는가?
• RQ5 다양한 랭크(r=384–768) 및 모델 규모(8–32B)에서의 메모리 및 속도 트레이드오프는 무엇인가?

주요 결과

• 인자화 노름은 랭크 의존적 지속 메모리를 O(d_out d_in)에서 O(d_out r + r^2)로 축소하고 B@A 물질화의 필요성을 제거한다.
• 융합된 Triton 커널은 네 가지 DoRA 연산을 하나의 패스로 축소하여 전방 속도는 1.5–2.7×, 역방 속도는 1.06–1.23× 가속시키고 최대 7 GB의 피크 VRAM을 낮춘다.
• 6개의 8–32B VLM을 대상으로 3개의 GPU에서 역전 그래디언트 계산은 HF PEFT의 DoRA 기준선보다 1.46–1.87×, 본 저자의 즉시 기준선보다 1.18–1.24× 빠르며 추론 속도는 1.5–2.0× 향상된다.
• 융합된 피드가 있는 최종 로짓과 즉시 기반의 로짓 간 코사인 유사도는 모든 모델/GPU 쌍에서 0.9999를 초과하여 충실도가 높음을 나타낸다.
• 융합 커널하의 수렴은 2000스텝 동안 1회 평균 손실 변화가 7.1e-4 이내로 즉시 학습과 일치한다(여러 시드에서).
• 메모리 프로파일링 실험에서 역전이 전방 피크 VRAM을 감소시키고 동일 메모리 예산 내에서 더 큰 구성의 학습을 가능하게 한다.

Figure 2 : Three-tier dispatch: fused backward for training (Tier 1), fused forward for inference (Tier 2), eager fallback for CPU, no-Triton, or sub-crossover shapes (Tier 3).