PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Alexandra Zelenin, Alexandra Zhuravlyova

arXiv:2603.22276v1cs.LGstat.ML

TL;DR

High-rank DoRA can require dense norm intermediates that create substantial memory pressure. The paper replaces this computation with a factored norm and combines it with fused Triton kernels; across VLMs and GPUs, the fused implementation is faster than HF PEFT while preserving output and training fidelity.

  • Problem

    DoRA’s row-wise norm computation materializes the dense [d_out, d_in] product BA, creating costly memory pressure at high rank and across many adapted modules.

  • Method

    The paper computes the norm from base, cross, and Gram terms using O(d_out r + r^2) intermediates, then fuses DoRA composition into single-pass Triton kernels.

  • Results

    1.5–2.0× faster inference and 1.5–1.9× faster gradient computation than HF PEFT’s DoRA implementation, with up to 7 GB lower peak VRAM.

  • Takeaways & Limitations

    The fused implementation preserves fidelity while reducing memory traffic and improving high-rank DoRA efficiency across the evaluated VLM and GPU configurations.

  • Takeaways & Limitations

    Fusion offers no advantage below ∼2048 × 6144 activations, Triton kernels are unavailable on non-CUDA platforms, and FSDP2/DTensor is unsupported.

Abstract

from arXiv · show

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.

1 Introduction

DoRA extends LoRA by separating weight magnitude from direction, but its row-wise norm computation can materialize costly dense products. This paper proposes factored norms and fused Triton kernels to reduce working memory and computation while preserving the same DoRA calculation.

  • DoRA extends LoRA by decomposing adapted weights into magnitude and direction.
  • The row-wise norm of W + sBA is a bottleneck because major frameworks materialize the dense BA product.
  • ∼512 MB of transient memory is allocated for one module at d_in = 8192, with checkpointing allocating these temporaries twice per step.
  • The proposed factored norm computes three terms through O(d_out r + r^2) intermediates without materializing BA.At d = 8192 and r = 512 in fp32, the theoretical persistent-memory reduction is 15×.
  • Fused Triton kernels reduce DoRA composition from four CUDA kernel launches to one pass while using a numerically stable near-unity rescaling form.A three-tier dispatch selects fused backward, fused forward, or eager fallback paths.
  • The implementation is evaluated across six GPUs and six 8–32B VLMs using PEFT, Dense (B@A), Eager, and Fused configurations.

2 Factored Norm Computation

The factored norm computes DoRA’s row-wise norm from base, cross, and Gram terms using low-rank intermediates, avoiding dense BA materialization. This reduces rank-dependent memory substantially, although chunked base-norm accumulation and fp32 arithmetic shape measured costs.

  • Algebraic decomposition: The squared row-wise norm expands into base, cross, and BA-norm terms that can be computed without materializing BA.The cross term uses row-wise inner products, while the BA-norm term factors through the Gram matrix.
  • Algebraic decomposition: O(d_out r + r^2) intermediates replace dense [d_out, d_in] temporaries in the factored computation.Chunking accumulates the required intermediates while keeping working memory within a configurable budget.
  • Complexity: 241 MB of the measured memory delta is attributed to the chunked base-norm buffer, which is rank-independent at the tested settings.The theoretical reduction still captures the asymptotic benefit of removing rank-dependent tensors as rank grows.
  • Complexity: 15× is the theoretical reduction in rank-dependent persistent memory at d=8192 and r=512 in fp32.The corresponding table reports 3.2× measured reduction because allocator deltas include the rank-independent base-norm transient.
  • Precision and implementation: The factored norm accumulates in fp32 regardless of weight dtype, so its isolated bf16 norm memory can exceed that of half-precision PEFT.Model-level VRAM savings remain because the fused compose kernel eliminates forward-pass intermediates.
  • Compute tradeoff: 4.8× slower is the factored norm in isolation than the dense reference on H200 fp32, because chunked matmuls replace one contiguous norm call.End-to-end speed can nevertheless improve because dense BA materialization dominates time and memory; on RTX 6000 PRO, the factored norm matches or outperforms the reference at r ≤384.

3 Fused Triton Kernels

Fused Triton kernels reduce DoRA composition overhead while preserving numerical stability near the practically common regime g ≈1. The kernels combine efficient forward and backward computation with device-aware tuning and fallback support.

  • Compose Kernel: A single fused pass reduces DoRA composition memory traffic by about 4×, while realized speedup reaches 2.0–2.7× over eager PyTorch.The eager path performs four element-wise operations and about 12 memory passes; the fused kernel uses three reads and one write.
  • Numerical Stability: Mean g ≈1.0 with std ≈0.0015 makes catastrophic cancellation in the naive composition form a practical numerical concern.DoRA magnitude initialization and training keep the rescaling factor tightly concentrated around unity.
  • Numerical Stability: The stable form uses fp32 intermediates and explicit (g−1) correction, while canonical evaluation order addresses bf16 non-associativity.The same ordered computation is enforced across PyTorch composition paths.
  • Numerical Stability: 3.0× lower peak error near g ≈1 is achieved by the stable compose form in bf16, with stable and fused paths near the bf16 quantization floor.The comparison uses d_out = 8192 and d_in = 2048 against an fp64 reference.
  • Kernel Tuning: Per-device autotuning is required because only ∼9% of optimal configurations agree pairwise across six GPUs.First-run autotuning takes 10–30 s per kernel, with cached configurations persisted by Triton.
  • Backward Kernel: The fused backward computes dlora and dbase in one pass, while dmag remains a separate PyTorch reduction to avoid atomic contention and nondeterministic ordering.Reducing ROWS_PER_PROGRAM also lowers register pressure when writing two output tensors.
  • Norm Assembly Kernel: A second Triton kernel computes wnorm from three factored terms, while magnitude division remains in PyTorch so both norm paths share one precision context.Store-reload barriers and an exact-rounding sqrt instruction reproduce PyTorch’s evaluation order for the norm kernel.
  • Runtime Selection: Three dispatch tiers provide fused backward training, fused forward inference, and eager fallback for CPU, unavailable Triton, or sub-crossover shapes.The dispatch organization is summarized in Figure 2 and Table 2.

4 Runtime Dispatch

Runtime dispatch selects fused training or inference kernels when appropriate and falls back to pure PyTorch for unsupported or small workloads. The implementation preserves precision contracts and supports several distributed and compilation environments.

  • Dispatch: _compose_with_dispatch selects the composition path at runtime using four environment variables, with defaults requiring no configuration.The dispatch mechanism is designed to account for kernel availability and working-set budgets.
  • Tier 1: Fused Backward: Tier 1 uses a dual-output fused backward kernel that computes output and inner together, eliminating the sequential forward-pass VRAM spike.When magnitude is frozen, the inner allocation is skipped; auto-mode uses a crossover based on d_out and activation size.
  • Tier 1: Fused Backward: ∼71% of adapted modules per layer use Tier 1 during training, while ∼29% fall back to Tier 3 because some KV projections are below crossover.In the evaluated VLMs, KV projections have d_out as low as 512.
  • Tier 2: Fused Forward: Tier 2 uses a forward-only Triton kernel without autograd graph nodes when requires_grad is false.This path targets inference workloads that do not need backward state.
  • Tier 3: Eager Fallback: Tier 3 uses pure PyTorch for CPU, unavailable Triton, and sub-crossover training, with out-of-place composition under autograd.The out-of-place behavior avoids aliasing when gradients are active.
  • Precision: All PyTorch compose paths are bitwise-identical, while Triton outputs stay within 10−4 max-abs error in fp32 and dtype-appropriate tolerances in bf16/fp16.Triton preserves the algebra but may differ in final bits because of FMA contraction and reduction trees.
  • Compatibility: The fused compose is graph-break-free under torch.compile when dropout is inactive and supports DeepSpeed ZeRO-2/3 and FSDP1, but not FSDP2/DTensor.The operation is registered as the custom op peft::fused_dora_compose.
  • Precision: Magnitude division uses g = m/ max(wnorm, ϵ) in PyTorch outside the no_grad norm context, keeping precision identical across execution tiers.This division is shared by all tiered norm paths.

5 Experiments

Across model-level and microbenchmark evaluations, fused DoRA consistently improves speed and reduces peak VRAM, with gains depending on workload size, rank, and hardware. Numerical fidelity remains effectively unchanged, while end-to-end training gains are smaller after non-adapter overheads.

  • Model-Level Performance: 1.46–1.87× faster gradient computation than HF PEFT and 1.18–1.24× faster than eager, with 1.3–6.7 GB lower peak VRAM.These measurements cover six 8–32B VLMs on three GPUs and exclude optimizer updates.
  • Model-Level Performance: 1.5–2.0× faster inference than HF PEFT, with all six models—including 32B—running on the RTX 6000 PRO at 84–88 GB peak.The same 32B models OOM during gradient computation on that GPU.
  • High-Rank Scaling: 1.66× →1.74× speedup over PEFT as rank increases from 384 to 768 for the 32B model, while eager-relative speedup decreases from 1.18× →1.14×.PEFT materialization cost grows with rank, whereas larger LoRA matmuls dilute the compose-kernel contribution.
  • Compose Kernel Performance: 2.70×, 2.62×, 2.00×, 1.92×, 1.73×, and 1.47× forward compose speedups occur on B200, B300, H200, RTX 6000 PRO, A100, and L40S, respectively.The cross-architecture pattern is consistent with reduced memory traffic across GPUs ranging from GDDR6 to HBM3e.
  • Compose Kernel Performance: Above approximately 8192 × 8192, fused backward wins on all six GPUs, while below approximately 2048 × 6144 it can trail eager at 0.88–0.99×.Geometric-mean backward speedups range from 1.06× on L40S to 1.23× on B200.
  • Fidelity and End-to-End Training: Cosine similarity between fused and eager final logits exceeds 0.9999 across all six models and three GPUs, while a 2000-step run reduces total training time by 8.3%.The wall-clock improvement is smaller than adapter-computation speedups after optimizer, data-loading, and framework overheads.

6 Discussion

The discussion weighs deployment tradeoffs: fused paths reduce model-level memory and support key training and inference scenarios, while limitations remain around hardware, distributed training, embeddings, and ablation scope.

  • Tradeoffs and Limitations: Below ∼2048 × 6144 activations, launch latency dominates, and Triton kernels are unavailable on non-CUDA platforms.The dispatch therefore uses a conservative crossover heuristic and an eager fallback.
  • Tradeoffs and Limitations: 0.1–1.0 GB less peak VRAM than eager is achieved by fused backward through eliminating an activation-sized tensor and sequential forward-pass spikes.With frozen magnitude, the inner tensor is skipped entirely.
  • Tradeoffs and Limitations: 46–87% slower than the PEFT path, the factored norm presents a framework-level latency tradeoff despite the broader fused system’s benefits.Dense temporaries can also compete with inference memory budgets.
  • Tradeoffs and Limitations: FSDP2/DTensor is unsupported because the factored norm assumes access to the full base weight W.Supporting FSDP2 would require chunk-wise partial-sum accumulation followed by an all-reduce.
  • Tradeoffs and Limitations: Adapted embeddings are excluded from headline benchmarks, and PEFT-finetuned embedding checkpoints may require re-fine-tuning or a legacy composition fallback.The implementation applies the full DoRA formula to embeddings, correcting PEFT’s omission of the base term.
  • Tradeoffs and Limitations: Model-level speedups reflect factored norms and fused kernels jointly, while a fuller factorial ablation across additional model families remains needed.Microbenchmarks and eager-versus-fused comparisons provide only component-level and partial kernel-fusion evidence.

7 Related Work

The paper positions its contribution as a complementary systems optimization for DoRA, applying factored computation and kernel fusion to transient memory and composition costs rather than adapter statistics or parameter counts.

  • Parameter-efficient fine-tuning: DoRA separates magnitude from direction, while rsLoRA’s rank-stabilized scaling appears in all three terms of the factored norm.The work optimizes DoRA execution rather than introducing a new adapter architecture.
  • DoRA variants: EDoRA and DoRAN target statistical efficiency, whereas this work targets transient memory and is therefore complementary.EDoRA reduces static parameter count through SVD; DoRAN perturbs the normalization denominator.
  • Framework implementations: Major surveyed frameworks use torch.eye materialization, and none avoids materializing dense BA as of February 2026.Unsloth disables custom kernels when DoRA is active, while orchestration frameworks delegate to PEFT.
  • Kernel fusion: Factored norms and fused kernels apply established numerical-linear-algebra and GPU-fusion principles specifically to DoRA’s memory-bound composition.The contribution includes dtype discipline, chunking, and integration into the fused pipeline.
  • LLM-guided optimization: KernelAgent achieved 3.58× over eager for backward through a two-stage partial-reduction strategy, exceeding the paper’s 1.06–1.23× backward speedup.The release prioritizes drop-in compatibility and end-to-end wins across real models.

8 Conclusion

The paper concludes that factored norms and fused Triton kernels make high-rank DoRA more memory-efficient and faster without materially changing model outputs or training behavior, within documented scope limits.

  • Conclusion: O(d_out × d_in) working memory is reduced to O(d_out × r + r^2) through a factored norm and single-pass fused GPU composition.These are systems changes to DoRA execution rather than a new adapter architecture.
  • Conclusion: 1.5–2.0× faster inference and 1.5–1.9× faster gradient computation than HF PEFT are achieved on six 8–32B VLMs, with up to 7 GB lower peak VRAM.Microbenchmarks confirm 1.5–2.7× compose-kernel speedup across six GPUs and four architecture generations.
  • Conclusion: Final-logit cosine similarity exceeds 0.9999, and matched multi-seed training curves support fidelity across operator, model-output, and convergence checks.The convergence scope is limited to two model families, two optimizers, and one SFT dataset.
  • Conclusion: FSDP2 remains unsupported, model-level benchmarks cover only three of six GPUs, and the empirical dispatch crossover may require retuning on future hardware.Generalization to RL pipelines remains to be confirmed.

Data Availability

The release provides source code and benchmark artifacts, while documenting the composition contract and numerical handling needed for replication.

  • Data Availability: All source code, benchmark scripts, raw JSON results, Triton autotune caches, and figure-generation scripts are publicly available at the cited repository.The release is tagged v1.0 and includes reproducible convergence validation using a public dataset.
  • Data Availability: The module returns a delta ΔY, which the caller adds to Ybase using the full DoRA composition formula.The contract specifies ΔY = g ⊙(sXAᵀBᵀ) + (g −1) ⊙Ybase.
  • Data Availability: Norm quantities are recomputed each forward pass, detached, accumulated in FP32 with autocast disabled, and use dtype-specific epsilon values.Bias is subtracted before composition and re-added afterward.

B Implementation Details

The implementation combines factored norm assembly with fused Triton composition, while preserving numerical behavior through explicit precision, shape, and fallback rules.

  • Compatibility and fallbacks: Per-device autotuning is essential: exact configuration agreement across GPUs is only ∼9%, while RPP=1 is selected in 1149/1206 autotuned entries.Chunking, dropout handling, and scale-is-zero fast paths further specialize execution.
  • Kernel design: The forward kernel fuses base and LoRA composition into one pass, broadcasting the magnitude vector along all dimensions except the last.Inputs are base, LoRA output, magnitude g, and scalar s; tensors remain in the input dtype.
  • Kernel design: The backward kernel computes LoRA and base gradients in one pass, while dmag uses a separate reduction to avoid nondeterministic atomic ordering.Eager training uses Triton for both directions; compiled training uses Inductor for the PyTorch backward graph.
  • Numerical stability: The norm assembly computes max(base_sq + two_s · cross + s2 · ba_sq, 0) in fp32, with controlled evaluation ordering and correctly rounded square roots.Store-reload barriers prevent FMA fusion, and inline PTX supplies IEEE 754 sqrt.rn.f32 behavior.
  • Numerical stability: Magnitude division remains a PyTorch operation, ensuring identical precision across norm paths at the cost of one negligible element-wise kernel launch.The division is g = m / max(wnorm, ε).
  • Compatibility and fallbacks: Unsupported broadcast shapes, non-contiguous inputs, and incompatible output dimensions route execution to eager fallback paths.The fused path requires last-dimension magnitude broadcasting and dout divisible by BLOCK_SIZE 128.

D Reproducibility

The paper provides code, benchmark artifacts, pinned environments, explicit memory definitions, and compatibility coverage to support reproducibility across model-level, convergence, and operator tests.

  • Artifacts: All source code, benchmark scripts, raw JSON results, autotune caches, and figure-generation scripts are released at the project repository.The patched PEFT module is included as a git submodule.
  • Coverage: The compatibility matrix separately records model benchmarks, convergence runs, and CI-only operator coverage, with the full test suite containing 1041 tests.Linear layers are the primary benchmark target, while convolution and embedding coverage is CI-only.
  • Environment: All benchmarks run under a single pinned Docker-compatible software stack, including PyTorch 2.10.0+cu130, Triton 3.6.0, and Transformers 5.2.0.The environment includes CUDA 13.1, driver 580.126.09, Python 3.12.12, and Linux 6.8.0.
  • Measurement: Memory reporting distinguishes allocator peak, working-set delta, and reserved VRAM according to microbenchmark, model-level, and colocated-workload use cases.These metrics separate operation footprint, transient DoRA overhead, and physically withheld GPU memory.
  • Regeneration: Figures and benchmark outputs are regenerable from included artifacts, scripts, and commands covering extended microbenchmark shapes and six-model runs.The extended suite uses 200 repeats, while model-level benchmarking uses rank 384 with gradient accumulation.
  • Evaluation inputs: The benchmarks use six specified vision-language model identifiers, with weights downloaded in March 2026 and exact hashes stored in the JSON artifacts.The model suite spans Qwen, Gemma, and Mistral checkpoints from 8B to 32B.
  • Evaluation inputs: Convergence validation uses a filtered, repacked MMFineReason dataset with token length at most 4096 and a published preprocessing script.Training uses SWIFT with the pinned PyTorch, Transformers, Triton, DeepSpeed, and FlashAttention versions.
  • Coverage: Table 13 extends the model-level peak-VRAM comparison to all six models.Its values use the peak_vram_mb measurement reported for the same setup as Table 8.

F Single-Layer E2E Decomposition

Single-layer end-to-end measurements isolate per-layer overhead but do not predict model-level speedup, because compose gains accumulate across many modules while backward overhead is amortized.

  • Interpretation: Compose gains compound across ∼500 DoRA modules in real models, whereas per-layer backward overhead is amortized.Consequently, single-layer E2E measurements can understate model-level benefit.
  • Microbenchmarks: The fp32 geometric-mean microbenchmark speedups range from 1.53× to 2.35× across six GPUs, with norm memory speedup fixed at 3.2×.The reported rows cover L40S, A100, RTX 6000 PRO, H200, B200, and B300.
  • Framework comparison: The surveyed five major fine-tuning frameworks all materialize the dense product for DoRA norms and provide no memory-efficient alternative.The comparison reflects manually inspected source implementations as of February 2026.
  • Cross-GPU results: All GPUs show consistent single-layer E2E improvement across eager and fused implementations and evaluated ranks.The figure uses bf16 with d = 4096, batch size 4, and sequence length 2048.
  • Scaling behavior: The single-layer E2E benefit peaks at hidden dimension h = 3072–4096, corresponding to common LLM sizes.This comparison fixes r = 384 across six GPUs.
Loading 2603.22276v1…