Масштабирование DoRA: Высокорнковая адаптация через факторизованные нормы и слияние ядер

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

March 23, 2026

Авторы: Alexandra Zelenin, Alexandra Zhuravlyova

cs.AI

Аннотация

Метод Weight-Decomposed Low-Rank Adaptation (DoRA) расширяет LoRA за счет разделения величины и направления весов, однако его прямой проход требует вычисления поэлементной нормы по строкам для матрицы W + sBA. Все изученные нами основные фреймворки реализуют это вычисление через материализацию плотного произведения [d_out, d_in] матриц BA. При d_in = 8192 и ранге r = 384 вычисление нормы для одного модуля требует около 512 МБ временной рабочей памяти в формате bf16, что делает DoRA с высоким рангом дорогостоящим и часто неосуществимым в типичных конфигурациях с одной GPU при работе с сотнями адаптируемых модулей и использованием контрольных точек. Мы представляем два системных усовершенствования. *Факторизованная норма* разлагает квадрат нормы на базовую, перекрестную и Грама компоненты, вычислимые через промежуточные значения сложности O(d_out r + r^2), что исключает необходимость плотного произведения. *Слияние ядер Triton* объединяет четырехъядерную композицию DoRA в один проход, сокращая объем пересылаемых данных в памяти примерно в 4 раза и используя численно устойчивую форму, которая избегает катастрофической потери точности в режиме масштабирования близком к единице, где на практике концентрируются масштабы величин. На шести Vision-Language моделях (VLM) объемом 8-32B, протестированных на трех GPU NVIDIA (RTX 6000 PRO, H200, B200) при r = 384 в bf16, объединенная реализация оказывается в 1.5-2.0 раза быстрее реализации DoRA от Hugging Face PEFT при выводе и в 1.5-1.9 раза быстрее при вычислении градиента (без учета шага оптимизатора), с пиковым потреблением VRAM ниже до 7 ГБ. Микробенчмарки на шести GPU, охватывающих четыре поколения архитектур (L40S, A100, RTX 6000 PRO, H200, B200, B300), подтверждают ускорение ядра композиции в 1.5-2.7 раза. Косинусное сходство финальных логитов превышает 0.9999 для всех пар модель/GPU, а кривые обучения при нескольких сидах совпадают со средним отклонением потерь на шаг в пределах 7.1 x 10^-4 за 2000 шагов.

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.