LoRA/DoRA fast forward: no weight materialization (1.6-2.0x on MPS)#69
Open
Cortexelus wants to merge 1 commit into
Open
LoRA/DoRA fast forward: no weight materialization (1.6-2.0x on MPS)#69Cortexelus wants to merge 1 commit into
Cortexelus wants to merge 1 commit into
Conversation
Cortexelus
force-pushed
the
fast-lora-forward
branch
3 times, most recently
from
July 15, 2026 06:08
26654c8 to
cb6842d
Compare
…-2.0x MPS training) torch.nn.utils.parametrize materializes W' = f(W0) on every self.weight access, so a plain forward pays an [out,in] low-rank matmul + [out,in] norm (and several [out,in] autograd intermediates in backward) per layer per step. Reformulate to compute the same y without ever building an [out,in] tensor beyond reading W0: fold the low-rank update into the data GEMMs, and get the DoRA row/col norm from the closed-form norm^2 expansion (a cached rowsum(W0^2) constant plus two small rank-sized terms). Semantics mirror LoRAParametrization exactly — norm stays in the autograd graph, same 1e-12 eps, dropout drawn once and shared, strength==0 short-circuit, un-scaled bias — verified to gradient parity (<=4.4e-7 fp32). Installed unconditionally by add_lora as an instance-attribute forward, leaving state_dict layout and parametrization registration untouched. Anything the fast path can't reproduce exactly (stacked/multi-LoRA, bora/-xs, fan_in_fan_out, disable_lora, trainable base, custom Linear subclass) falls back to the naive materialized forward per module.
Cortexelus
force-pushed
the
fast-lora-forward
branch
from
July 15, 2026 06:14
cb6842d to
bcd2cb4
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
LoRA/DoRA training-forward reformulation — no full-weight materialization. Stacked on #68 (base branch
mps-training; uses itsdisable_autocasthelper — GitHub will retarget this PR to main when #68 merges). Speeds up LoRA training on every backend; measured on MPS: sm-music 0.668→1.053 it/s (1.58×), medium 0.254→0.501 it/s (1.97×), sm peak memory 3.67→2.58 GiB.What
torch.nn.utils.parametrizerebuilds the full[out,in]adapted weight on every forward. For lora / dora-rows / dora-cols that's algebraically unnecessary — dora-rows is exactly a row-scale of the (base + low-rank) output with the norm in closed form:enable_fast_lora_forward(model)installs an instance-attribute forward on each parametrized Linear — the parametrization registration and theparametrizations.weight.<i>.*checkpoint layout are untouched, so saves/loads/resume/set_lora_strength/remove_lora/merge_lorawork unchanged. Applied unconditionally byadd_lora— there is no toggle: the cached norm constant is a per-layer vector (no memory tradeoff; peak memory strictly drops), speed strictly improves, and the naive path remains in-tree for the fallback cases. Falls back to the standard path for: stacked parametrizations (ARC at lora_index=1), bora/-xs, fan_in_fan_out,disable_lora, trainable base, custom Linear subclasses; transparently no-ops afterremove_lora/merge_lora.Fidelity
The reference's in-graph norm (no DoRA-paper detach) is mirrored; strength multiplies inside the delta; weight-space dropout drawn once per forward and shared with the norm terms; dora-rows bias added un-scaled after the row scale; strength 0 is bit-exact vs the frozen base. 40 unit tests: fp32 output parity 1e-5, gradient parity ≤4.4e-7 rel (A/B/magnitude), autocast parity ~1e-3 (fp16 GEMM noise), stacking fallback bit-exact, state_dict key sets identical. Same-seed 30-step training runs deviate ≤1.2e-3 from the pre-change loss curves with identical RNG consumption; the fp32 controlled-parity forward matches the recorded loss exactly (3.98231220).
Caveat
Default-ON changes CUDA training numerics at the ~1e-3 autocast level (op reordering, same class as any fusion). Gradient-parity tests cover the math on CPU/MPS; one CUDA smoke would be welcome. Rollback path is reverting the PR (the change is confined to models/lora/model.py).