dora-factored-kernel / tests /test_reference_parity.py

github-actions[bot]

DoRA-factored kernels-library package — Stage A scaffold (#14)

bdcc32f 3 days ago

4.82 kB

	# 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)