code/models/tt_dit/experimental/tests/test_lora_variants.py · sunandbean/diffrhythm-1.2-p100a at main

diffrhythm-1.2-p100a / code /models /tt_dit /experimental /tests /test_lora_variants.py

Add files using upload-large-folder tool

33d23e1 verified 14 days ago

9.05 kB

# SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc.
# SPDX-License-Identifier: Apache-2.0

"""
Parity tests for the three LoRA-aware Linear subclasses, each exercised
under both ``lora_mode`` options:
  fuse     — bind merges ``scale * A.T @ B.T`` into the base weight in-place.
             Forward runs the un-modified base path.
  runtime  — bind uploads A, B to device; forward adds
             ``scale * (x @ A) @ B`` on top of the base output.
The unbind branch is checked in both modes — for ``fuse`` it subtracts
the delta back out of W; for ``runtime`` it frees the A/B tensors so the
forward path collapses to the base.
Mesh shape is 1x1 so sharding collapses to the replicated case for all
three variants. The tests still exercise the LoRA-specific upload and
forward code paths and catch class-level bugs in shape handling,
fused-activation fallback, and chunked-output handling.
Run with:
    pytest -xvs models/tt_dit/experimental/tests/test_lora_variants.py
"""
from __future__ import annotations

import pytest
import torch

import ttnn
from models.tt_dit.layers.linear import LoRAColParallelLinear, LoRALinear, LoRARowParallelLinear
from models.tt_dit.utils.check import assert_quality
from models.tt_dit.utils.tensor import bf16_tensor


def _make_pair(in_features: int, out_features: int, rank: int, dtype: torch.dtype):
    A = torch.randn(rank, in_features, dtype=dtype) * 0.02
    B = torch.randn(out_features, rank, dtype=dtype) * 0.02
    return A, B


def _ref_lora(x, W, b, A, B, scale):
    """y = x @ (W + scale * B @ A).T + b"""
    W_eff = W + scale * (B @ A)
    return torch.nn.functional.linear(x, W_eff, b)


# --------------------------------------------------------------------
# LoRALinear (replicated)
# --------------------------------------------------------------------
@pytest.mark.parametrize("mesh_device", [(1, 1)], indirect=True)
@pytest.mark.parametrize("lora_mode", ["fuse", "runtime"])
@pytest.mark.parametrize(
    ("seq_len", "in_features", "out_features", "rank"),
    [
        (128, 256, 256, 8),
        (256, 512, 1024, 16),
    ],
)
def test_lora_linear_variant(
    mesh_device: ttnn.MeshDevice,
    lora_mode: str,
    seq_len: int,
    in_features: int,
    out_features: int,
    rank: int,
) -> None:
    torch.manual_seed(0)
    dtype = torch.bfloat16
    scale = 0.6

    torch_lin = torch.nn.Linear(in_features, out_features, bias=True).to(dtype=dtype)
    torch_lin.eval()
    A_t, B_t = _make_pair(in_features, out_features, rank, dtype)
    x_t = torch.randn(1, 1, seq_len, in_features, dtype=dtype)

    with torch.no_grad():
        y_base = torch_lin(x_t)
        y_lora = _ref_lora(x_t, torch_lin.weight, torch_lin.bias, A_t, B_t, scale)

    tt = LoRALinear(in_features, out_features, bias=True, mesh_device=mesh_device, lora_mode=lora_mode)
    tt.load_torch_state_dict(torch_lin.state_dict())
    idx = tt.register_lora(A_t, B_t, scale=scale)
    x_tt = bf16_tensor(x_t, device=mesh_device)

    # 1) no adapter bound → base behavior
    y = tt(x_tt)
    for t in ttnn.get_device_tensors(y):
        assert_quality(y_base, ttnn.to_torch(t), pcc=0.999_5)

    # 2) bind active → LoRA effect applied
    tt.bind_active(idx, scale=scale)
    y = tt(x_tt)
    for t in ttnn.get_device_tensors(y):
        assert_quality(y_lora, ttnn.to_torch(t), pcc=0.999_0)

    # 3) unbind collapses back to base
    tt.unbind_active()
    y = tt(x_tt)
    for t in ttnn.get_device_tensors(y):
        assert_quality(y_base, ttnn.to_torch(t), pcc=0.999_5)

    tt.deallocate_lora()


# --------------------------------------------------------------------
# LoRAColParallelLinear — single output and chunked output
# --------------------------------------------------------------------
# Chunked configs run fuse-only: the runtime delta path doesn't split
# its delta into the same chunk layout as the base output.
_COL_PARALLEL_PARAMS = [
    # (label, lora_mode, seq, in, out, rank, chunks, activation_fn)
    ("plain_fuse", "fuse", 128, 256, 512, 8, None, None),
    ("plain_runtime", "runtime", 128, 256, 512, 8, None, None),
    ("gelu_fuse", "fuse", 128, 256, 1024, 8, None, "gelu"),
    ("gelu_runtime", "runtime", 128, 256, 1024, 8, None, "gelu"),
    ("chunked_qkv_fuse", "fuse", 128, 256, 3 * 256, 8, 3, None),
]


@pytest.mark.parametrize("mesh_device", [(1, 1)], indirect=True)
@pytest.mark.parametrize(
    ("lora_mode", "seq_len", "in_features", "out_features", "rank", "chunks", "activation_fn"),
    [params[1:] for params in _COL_PARALLEL_PARAMS],
    ids=[params[0] for params in _COL_PARALLEL_PARAMS],
)
def test_lora_col_parallel_variant(
    mesh_device: ttnn.MeshDevice,
    lora_mode: str,
    seq_len: int,
    in_features: int,
    out_features: int,
    rank: int,
    chunks: int | None,
    activation_fn: str | None,
) -> None:
    torch.manual_seed(0)
    dtype = torch.bfloat16
    scale = 0.5

    torch_lin = torch.nn.Linear(in_features, out_features, bias=True).to(dtype=dtype)
    torch_lin.eval()
    A_t, B_t = _make_pair(in_features, out_features, rank, dtype)
    x_t = torch.randn(1, 1, seq_len, in_features, dtype=dtype)

    def torch_act(y):
        if activation_fn == "gelu":
            return torch.nn.functional.gelu(y)
        return y

    with torch.no_grad():
        y_base = torch_act(torch_lin(x_t))
        y_lora = torch_act(_ref_lora(x_t, torch_lin.weight, torch_lin.bias, A_t, B_t, scale))

    tt = LoRAColParallelLinear(
        in_features,
        out_features,
        bias=True,
        mesh_device=mesh_device,
        mesh_axis=0,
        activation_fn=activation_fn,
        chunks=chunks,
        lora_mode=lora_mode,
    )
    tt.load_torch_state_dict(torch_lin.state_dict())
    idx = tt.register_lora(A_t, B_t, scale=scale)
    x_tt = bf16_tensor(x_t, device=mesh_device)

    def _join(out):
        """Concat chunked outputs along last dim; pass through single tensors."""
        if isinstance(out, list):
            torch_chunks = [ttnn.to_torch(ttnn.get_device_tensors(o)[0]) for o in out]
            return torch.cat(torch_chunks, dim=-1)
        return ttnn.to_torch(ttnn.get_device_tensors(out)[0])

    # 1) no adapter bound → base behavior
    assert_quality(y_base, _join(tt(x_tt)), pcc=0.999_5)

    # 2) bind active → LoRA effect applied
    tt.bind_active(idx, scale=scale)
    assert_quality(y_lora, _join(tt(x_tt)), pcc=0.999_0)

    # 3) unbind collapses to base
    tt.unbind_active()
    assert_quality(y_base, _join(tt(x_tt)), pcc=0.999_0)

    tt.deallocate_lora()


# --------------------------------------------------------------------
# LoRARowParallelLinear — N_tp=1 (1x1 mesh) sanity
# --------------------------------------------------------------------
@pytest.mark.parametrize("mesh_device", [(1, 1)], indirect=True)
@pytest.mark.parametrize("lora_mode", ["fuse", "runtime"])
@pytest.mark.parametrize(
    ("seq_len", "in_features", "out_features", "rank"),
    [
        (128, 512, 256, 8),
        (256, 1024, 512, 16),
    ],
)
def test_lora_row_parallel_variant(
    mesh_device: ttnn.MeshDevice,
    lora_mode: str,
    seq_len: int,
    in_features: int,
    out_features: int,
    rank: int,
) -> None:
    from models.tt_dit.parallel.manager import CCLManager

    torch.manual_seed(0)
    dtype = torch.bfloat16
    scale = 0.4

    torch_lin = torch.nn.Linear(in_features, out_features, bias=True).to(dtype=dtype)
    torch_lin.eval()
    A_t, B_t = _make_pair(in_features, out_features, rank, dtype)
    x_t = torch.randn(1, 1, seq_len, in_features, dtype=dtype)

    with torch.no_grad():
        y_base = torch_lin(x_t)
        y_lora = _ref_lora(x_t, torch_lin.weight, torch_lin.bias, A_t, B_t, scale)

    ccl_manager = CCLManager(mesh_device, topology=ttnn.Topology.Linear)
    tt = LoRARowParallelLinear(
        in_features,
        out_features,
        bias=True,
        mesh_device=mesh_device,
        mesh_axis=0,
        ccl_manager=ccl_manager,
        lora_mode=lora_mode,
    )
    tt.load_torch_state_dict(torch_lin.state_dict())
    idx = tt.register_lora(A_t, B_t, scale=scale)
    x_tt = bf16_tensor(x_t, device=mesh_device)

    # baseline
    y = tt(x_tt)
    assert_quality(y_base, ttnn.to_torch(ttnn.get_device_tensors(y)[0]), pcc=0.999_5)

    # bind active
    tt.bind_active(idx, scale=scale)
    y = tt(x_tt)
    assert_quality(y_lora, ttnn.to_torch(ttnn.get_device_tensors(y)[0]), pcc=0.999_0)

    # unbind
    tt.unbind_active()
    y = tt(x_tt)
    assert_quality(y_base, ttnn.to_torch(ttnn.get_device_tensors(y)[0]), pcc=0.999_5)

    tt.deallocate_lora()


# --------------------------------------------------------------------
# Construction-time validation
# --------------------------------------------------------------------
@pytest.mark.parametrize("mesh_device", [(1, 1)], indirect=True)
def test_invalid_lora_mode_raises(mesh_device: ttnn.MeshDevice, expect_error) -> None:
    with expect_error(ValueError, "lora_mode"):
        LoRALinear(64, 64, bias=False, mesh_device=mesh_device, lora_mode="bogus")

	# SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc.
	# SPDX-License-Identifier: Apache-2.0

	"""
	Parity tests for the three LoRA-aware Linear subclasses, each exercised
	under both ``lora_mode`` options:

	fuse — bind merges ``scale * A.T @ B.T`` into the base weight in-place.
	Forward runs the un-modified base path.
	runtime — bind uploads A, B to device; forward adds
	``scale * (x @ A) @ B`` on top of the base output.

	The unbind branch is checked in both modes — for ``fuse`` it subtracts
	the delta back out of W; for ``runtime`` it frees the A/B tensors so the
	forward path collapses to the base.

	Mesh shape is 1x1 so sharding collapses to the replicated case for all
	three variants. The tests still exercise the LoRA-specific upload and
	forward code paths and catch class-level bugs in shape handling,
	fused-activation fallback, and chunked-output handling.

	Run with:
	pytest -xvs models/tt_dit/experimental/tests/test_lora_variants.py
	"""
	from __future__ import annotations

	import pytest
	import torch

	import ttnn
	from models.tt_dit.layers.linear import LoRAColParallelLinear, LoRALinear, LoRARowParallelLinear
	from models.tt_dit.utils.check import assert_quality
	from models.tt_dit.utils.tensor import bf16_tensor


	def _make_pair(in_features: int, out_features: int, rank: int, dtype: torch.dtype):
	A = torch.randn(rank, in_features, dtype=dtype) * 0.02
	B = torch.randn(out_features, rank, dtype=dtype) * 0.02
	return A, B


	def _ref_lora(x, W, b, A, B, scale):
	"""y = x @ (W + scale * B @ A).T + b"""
	W_eff = W + scale * (B @ A)
	return torch.nn.functional.linear(x, W_eff, b)


	# --------------------------------------------------------------------
	# LoRALinear (replicated)
	# --------------------------------------------------------------------
	@pytest.mark.parametrize("mesh_device", [(1, 1)], indirect=True)
	@pytest.mark.parametrize("lora_mode", ["fuse", "runtime"])
	@pytest.mark.parametrize(
	("seq_len", "in_features", "out_features", "rank"),
	[
	(128, 256, 256, 8),
	(256, 512, 1024, 16),
	],
	)
	def test_lora_linear_variant(
	mesh_device: ttnn.MeshDevice,
	lora_mode: str,
	seq_len: int,
	in_features: int,
	out_features: int,
	rank: int,
	) -> None:
	torch.manual_seed(0)
	dtype = torch.bfloat16
	scale = 0.6

	torch_lin = torch.nn.Linear(in_features, out_features, bias=True).to(dtype=dtype)
	torch_lin.eval()
	A_t, B_t = _make_pair(in_features, out_features, rank, dtype)
	x_t = torch.randn(1, 1, seq_len, in_features, dtype=dtype)

	with torch.no_grad():
	y_base = torch_lin(x_t)
	y_lora = _ref_lora(x_t, torch_lin.weight, torch_lin.bias, A_t, B_t, scale)

	tt = LoRALinear(in_features, out_features, bias=True, mesh_device=mesh_device, lora_mode=lora_mode)
	tt.load_torch_state_dict(torch_lin.state_dict())
	idx = tt.register_lora(A_t, B_t, scale=scale)
	x_tt = bf16_tensor(x_t, device=mesh_device)

	# 1) no adapter bound → base behavior
	y = tt(x_tt)
	for t in ttnn.get_device_tensors(y):
	assert_quality(y_base, ttnn.to_torch(t), pcc=0.999_5)

	# 2) bind active → LoRA effect applied
	tt.bind_active(idx, scale=scale)
	y = tt(x_tt)
	for t in ttnn.get_device_tensors(y):
	assert_quality(y_lora, ttnn.to_torch(t), pcc=0.999_0)

	# 3) unbind collapses back to base
	tt.unbind_active()
	y = tt(x_tt)
	for t in ttnn.get_device_tensors(y):
	assert_quality(y_base, ttnn.to_torch(t), pcc=0.999_5)

	tt.deallocate_lora()


	# --------------------------------------------------------------------
	# LoRAColParallelLinear — single output and chunked output
	# --------------------------------------------------------------------
	# Chunked configs run fuse-only: the runtime delta path doesn't split
	# its delta into the same chunk layout as the base output.
	_COL_PARALLEL_PARAMS = [
	# (label, lora_mode, seq, in, out, rank, chunks, activation_fn)
	("plain_fuse", "fuse", 128, 256, 512, 8, None, None),
	("plain_runtime", "runtime", 128, 256, 512, 8, None, None),
	("gelu_fuse", "fuse", 128, 256, 1024, 8, None, "gelu"),
	("gelu_runtime", "runtime", 128, 256, 1024, 8, None, "gelu"),
	("chunked_qkv_fuse", "fuse", 128, 256, 3 * 256, 8, 3, None),
	]


	@pytest.mark.parametrize("mesh_device", [(1, 1)], indirect=True)
	@pytest.mark.parametrize(
	("lora_mode", "seq_len", "in_features", "out_features", "rank", "chunks", "activation_fn"),
	[params[1:] for params in _COL_PARALLEL_PARAMS],
	ids=[params[0] for params in _COL_PARALLEL_PARAMS],
	)
	def test_lora_col_parallel_variant(
	mesh_device: ttnn.MeshDevice,
	lora_mode: str,
	seq_len: int,
	in_features: int,
	out_features: int,
	rank: int,
	chunks: int \| None,
	activation_fn: str \| None,
	) -> None:
	torch.manual_seed(0)
	dtype = torch.bfloat16
	scale = 0.5

	torch_lin = torch.nn.Linear(in_features, out_features, bias=True).to(dtype=dtype)
	torch_lin.eval()
	A_t, B_t = _make_pair(in_features, out_features, rank, dtype)
	x_t = torch.randn(1, 1, seq_len, in_features, dtype=dtype)

	def torch_act(y):
	if activation_fn == "gelu":
	return torch.nn.functional.gelu(y)
	return y

	with torch.no_grad():
	y_base = torch_act(torch_lin(x_t))
	y_lora = torch_act(_ref_lora(x_t, torch_lin.weight, torch_lin.bias, A_t, B_t, scale))

	tt = LoRAColParallelLinear(
	in_features,
	out_features,
	bias=True,
	mesh_device=mesh_device,
	mesh_axis=0,
	activation_fn=activation_fn,
	chunks=chunks,
	lora_mode=lora_mode,
	)
	tt.load_torch_state_dict(torch_lin.state_dict())
	idx = tt.register_lora(A_t, B_t, scale=scale)
	x_tt = bf16_tensor(x_t, device=mesh_device)

	def _join(out):
	"""Concat chunked outputs along last dim; pass through single tensors."""
	if isinstance(out, list):
	torch_chunks = [ttnn.to_torch(ttnn.get_device_tensors(o)[0]) for o in out]
	return torch.cat(torch_chunks, dim=-1)
	return ttnn.to_torch(ttnn.get_device_tensors(out)[0])

	# 1) no adapter bound → base behavior
	assert_quality(y_base, _join(tt(x_tt)), pcc=0.999_5)

	# 2) bind active → LoRA effect applied
	tt.bind_active(idx, scale=scale)
	assert_quality(y_lora, _join(tt(x_tt)), pcc=0.999_0)

	# 3) unbind collapses to base
	tt.unbind_active()
	assert_quality(y_base, _join(tt(x_tt)), pcc=0.999_0)

	tt.deallocate_lora()


	# --------------------------------------------------------------------
	# LoRARowParallelLinear — N_tp=1 (1x1 mesh) sanity
	# --------------------------------------------------------------------
	@pytest.mark.parametrize("mesh_device", [(1, 1)], indirect=True)
	@pytest.mark.parametrize("lora_mode", ["fuse", "runtime"])
	@pytest.mark.parametrize(
	("seq_len", "in_features", "out_features", "rank"),
	[
	(128, 512, 256, 8),
	(256, 1024, 512, 16),
	],
	)
	def test_lora_row_parallel_variant(
	mesh_device: ttnn.MeshDevice,
	lora_mode: str,
	seq_len: int,
	in_features: int,
	out_features: int,
	rank: int,
	) -> None:
	from models.tt_dit.parallel.manager import CCLManager

	torch.manual_seed(0)
	dtype = torch.bfloat16
	scale = 0.4

	torch_lin = torch.nn.Linear(in_features, out_features, bias=True).to(dtype=dtype)
	torch_lin.eval()
	A_t, B_t = _make_pair(in_features, out_features, rank, dtype)
	x_t = torch.randn(1, 1, seq_len, in_features, dtype=dtype)

	with torch.no_grad():
	y_base = torch_lin(x_t)
	y_lora = _ref_lora(x_t, torch_lin.weight, torch_lin.bias, A_t, B_t, scale)

	ccl_manager = CCLManager(mesh_device, topology=ttnn.Topology.Linear)
	tt = LoRARowParallelLinear(
	in_features,
	out_features,
	bias=True,
	mesh_device=mesh_device,
	mesh_axis=0,
	ccl_manager=ccl_manager,
	lora_mode=lora_mode,
	)
	tt.load_torch_state_dict(torch_lin.state_dict())
	idx = tt.register_lora(A_t, B_t, scale=scale)
	x_tt = bf16_tensor(x_t, device=mesh_device)

	# baseline
	y = tt(x_tt)
	assert_quality(y_base, ttnn.to_torch(ttnn.get_device_tensors(y)[0]), pcc=0.999_5)

	# bind active
	tt.bind_active(idx, scale=scale)
	y = tt(x_tt)
	assert_quality(y_lora, ttnn.to_torch(ttnn.get_device_tensors(y)[0]), pcc=0.999_0)

	# unbind
	tt.unbind_active()
	y = tt(x_tt)
	assert_quality(y_base, ttnn.to_torch(ttnn.get_device_tensors(y)[0]), pcc=0.999_5)

	tt.deallocate_lora()


	# --------------------------------------------------------------------
	# Construction-time validation
	# --------------------------------------------------------------------
	@pytest.mark.parametrize("mesh_device", [(1, 1)], indirect=True)
	def test_invalid_lora_mode_raises(mesh_device: ttnn.MeshDevice, expect_error) -> None:
	with expect_error(ValueError, "lora_mode"):
	LoRALinear(64, 64, bias=False, mesh_device=mesh_device, lora_mode="bogus")