File size: 4,824 Bytes
bdcc32f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 | # Copyright 2024-present the HuggingFace Inc. team.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Structural parity tests for the DoRA-factored reference kernel.
Asserts that the pure-PyTorch reference (``dora_factored_forward``) is numerically equivalent to a
naive dense DoRA effective-weight baseline written inline below. Mirrors the discipline of PEFT's
merged ``tests/test_dora_factored_norm.py`` (huggingface/peft#3382): deterministic seed, small
shapes, dense-vs-factored comparison. Runs on CPU with no GPU / no PEFT / no Triton.
"""
import pytest
import torch
from dora_factored import dora_factored_forward
def _dense_dora_weight(base_weight, lora_a, lora_b, scaling, dora_scale):
"""Naive dense DoRA effective weight: materialize ``W + s·BA`` then column-normalize.
This is the ~10-line baseline the factored reference must match. It mirrors
``DoraLinearLayer.get_weight_norm`` (the dense path) composed with the DoRA magnitude rescale.
"""
weight = base_weight + scaling * (lora_b @ lora_a)
col_norm = torch.linalg.norm(weight, dim=1)
scale = (dora_scale / col_norm).unsqueeze(-1)
return scale * weight
@pytest.mark.parametrize("scaling", [1.0, 0.5, 2.0])
@pytest.mark.parametrize(
"dtype, atol, rtol",
[
# fp32: same bar as the merged PR #3382 suite (tests/test_dora_factored_norm.py).
(torch.float32, 1e-5, 1e-5),
# bf16: the factored accumulation order diverges from torch.linalg.norm's internal
# accumulation more than in fp32; allow structural slack while still catching any
# formula bug (which diverges by orders of magnitude).
(torch.bfloat16, 5e-2, 5e-2),
],
)
def test_reference_matches_dense_dora(scaling, dtype, atol, rtol):
torch.manual_seed(0)
d_out, d_in, r = 32, 48, 4
base_weight = torch.randn(d_out, d_in, dtype=dtype)
lora_a = torch.randn(r, d_in, dtype=dtype)
lora_b = torch.randn(d_out, r, dtype=dtype)
# DoRA magnitude vector: one (positive) value per output column, like DoraLinearLayer.weight.
dora_scale = torch.randn(d_out, dtype=dtype).abs() + 0.5
dense = _dense_dora_weight(base_weight, lora_a, lora_b, scaling, dora_scale)
factored = dora_factored_forward(base_weight, lora_a, lora_b, scaling, dora_scale)
assert factored.shape == dense.shape == (d_out, d_in)
torch.testing.assert_close(factored, dense, atol=atol, rtol=rtol)
def test_factored_norm_matches_dense_norm():
# The factored decomposition's core claim, ported verbatim from PEFT's factored_weight_norm:
# the column-wise norm computed from the factors equals torch.linalg.norm of the dense weight.
# Mirrors tests/test_dora_factored_norm.py::test_factored_norm_matches_dense (huggingface/peft#3382)
# — the strongest provenance check that the port did not drift from the merged algorithm.
from dora_factored.reference import _factored_weight_norm
torch.manual_seed(1)
d_out, d_in, r = 16, 24, 3
base_weight = torch.randn(d_out, d_in)
lora_a = torch.randn(r, d_in)
lora_b = torch.randn(d_out, r)
scaling = 0.75
dense_norm = torch.linalg.norm(base_weight + scaling * (lora_b @ lora_a), dim=1)
factored_norm = _factored_weight_norm(base_weight, lora_a, lora_b, scaling)
assert factored_norm.shape == (d_out,)
torch.testing.assert_close(factored_norm, dense_norm, atol=1e-5, rtol=1e-5)
def test_reference_shape_and_finite():
torch.manual_seed(2)
base_weight = torch.randn(16, 24)
lora_a = torch.randn(3, 24)
lora_b = torch.randn(16, 3)
out = dora_factored_forward(base_weight, lora_a, lora_b, scaling=0.75, dora_scale=torch.ones(16))
assert out.shape == (16, 24)
assert torch.isfinite(out).all()
def test_scalar_dora_scale_broadcasts():
# dora_scale may be supplied as a scalar; it must broadcast across output columns identically
# to the per-column form filled with the same value.
torch.manual_seed(3)
base_weight = torch.randn(8, 12)
lora_a = torch.randn(2, 12)
lora_b = torch.randn(8, 2)
per_column = dora_factored_forward(base_weight, lora_a, lora_b, 1.0, torch.full((8,), 1.5))
scalar = dora_factored_forward(base_weight, lora_a, lora_b, 1.0, torch.tensor(1.5))
torch.testing.assert_close(scalar, per_column, atol=1e-5, rtol=1e-5)
|