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

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

Add files using upload-large-folder tool

33d23e1 verified 12 days ago

8.39 kB

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

"""
Integration tests for the runtime-LoRA adapter loader.
Builds a tiny WanTransformer3DModel (1 block) with ``lora_enabled=True``,
writes a synthetic LoRA safetensors file, runs ``load_adapter_into``, and
checks that:
  - every LoRA-aware Linear in the block got an adapter registered
  - ``bind_active`` marks the right modules active across all paths
  - ``unbind_active`` clears the active state
The forward math is covered by ``test_lora_variants.py``; this file
exercises the loader's key-mapping logic and the pipeline-mixin bind
flow against the default (``fuse``) mode.
Run with:
    pytest -xvs models/tt_dit/experimental/tests/test_adapter_loader.py
"""
from __future__ import annotations

import tempfile
from pathlib import Path

import pytest
import torch

import ttnn
from models.tt_dit.experimental.lora.adapter_loader import load_adapter_into
from models.tt_dit.layers.lora import LoRAMixin
from models.tt_dit.models.transformers.wan2_2.transformer_wan import WanTransformer3DModel
from models.tt_dit.parallel.config import DiTParallelConfig, ParallelFactor
from models.tt_dit.parallel.manager import CCLManager


def _save_synthetic_adapter(
    path: Path,
    *,
    num_blocks: int = 1,
    dim: int = 128,
    ffn_dim: int = 256,
    rank: int = 8,
) -> None:
    """Write a fake LoRA safetensors that covers all 6 LoRA-targeted Linears
    per block, using the lightx2v naming convention."""
    from safetensors.torch import save_file

    state: dict[str, torch.Tensor] = {}
    dtype = torch.bfloat16

    def add_pair(base: str, in_dim: int, out_dim: int):
        state[f"{base}.lora_A.weight"] = torch.randn(rank, in_dim, dtype=dtype) * 0.02
        state[f"{base}.lora_B.weight"] = torch.randn(out_dim, rank, dtype=dtype) * 0.02

    for i in range(num_blocks):
        # self-attn Q, K, V, O
        add_pair(f"blocks.{i}.self_attn.q", dim, dim)
        add_pair(f"blocks.{i}.self_attn.k", dim, dim)
        add_pair(f"blocks.{i}.self_attn.v", dim, dim)
        add_pair(f"blocks.{i}.self_attn.o", dim, dim)
        # cross-attn Q, K, V, O
        add_pair(f"blocks.{i}.cross_attn.q", dim, dim)
        add_pair(f"blocks.{i}.cross_attn.k", dim, dim)
        add_pair(f"blocks.{i}.cross_attn.v", dim, dim)
        add_pair(f"blocks.{i}.cross_attn.o", dim, dim)
        # ffn ff1, ff2 (lightx2v naming: ffn.0 = ff1, ffn.2 = ff2)
        add_pair(f"blocks.{i}.ffn.0", dim, ffn_dim)
        add_pair(f"blocks.{i}.ffn.2", ffn_dim, dim)

    save_file(state, str(path))


def _count_registered(transformer) -> int:
    """Count LoRA modules with at least one entry in their bank."""
    n = 0
    for mod in _iter_lora_modules(transformer):
        if any(a is not None for a in mod.lora_bank):
            n += 1
    return n


def _iter_lora_modules(module):
    if isinstance(module, LoRAMixin):
        yield module
    for _, child in module.named_children():
        yield from _iter_lora_modules(child)


@pytest.mark.parametrize("mesh_device", [(1, 1)], indirect=True)
def test_adapter_loader_registers_all_lora_linears(mesh_device: ttnn.MeshDevice) -> None:
    """Loader picks up Q/K/V/O for both attn1/attn2 plus ff1/ff2 per block."""
    num_heads = 4
    head_dim = 32
    dim = num_heads * head_dim
    ffn_dim = 2 * dim
    rank = 8
    num_blocks = 1

    parallel_config = DiTParallelConfig(
        cfg_parallel=ParallelFactor(factor=1, mesh_axis=0),
        tensor_parallel=ParallelFactor(factor=1, mesh_axis=1),
        sequence_parallel=ParallelFactor(factor=1, mesh_axis=0),
    )
    ccl_manager = CCLManager(mesh_device, topology=ttnn.Topology.Linear)

    transformer = WanTransformer3DModel(
        patch_size=(1, 2, 2),
        num_heads=num_heads,
        dim=dim,
        in_channels=16,
        out_channels=16,
        text_dim=128,
        freq_dim=64,
        ffn_dim=ffn_dim,
        num_layers=num_blocks,
        mesh_device=mesh_device,
        ccl_manager=ccl_manager,
        parallel_config=parallel_config,
        is_fsdp=False,
        lora_enabled=True,
    )

    # Every block should have 5 LoRA-aware Linears wired in: to_qkv (self),
    # to_q (cross), to_kv (cross), to_out for each attention (×2), ff1, ff2.
    # That's 1 + 1 + 1 + 2 + 1 + 1 = 7 modules per block.
    expected_per_block = 7
    n_modules = sum(1 for _ in _iter_lora_modules(transformer))
    assert (
        n_modules == num_blocks * expected_per_block
    ), f"expected {num_blocks * expected_per_block} LoRA modules; got {n_modules}"

    with tempfile.TemporaryDirectory() as tmp:
        adapter_path = Path(tmp) / "synthetic.safetensors"
        _save_synthetic_adapter(adapter_path, num_blocks=num_blocks, dim=dim, ffn_dim=ffn_dim, rank=rank)

        handle = load_adapter_into(transformer, str(adapter_path), scale=1.0, name="synthetic")

    # All 7 LoRA modules per block should have a registered adapter.
    registered = _count_registered(transformer)
    assert (
        registered == num_blocks * expected_per_block
    ), f"expected {num_blocks * expected_per_block} registered adapters; got {registered}"

    # Fused-QKV (to_qkv self-attn): rank should be 3*r. Fused KV (cross-attn
    # to_kv): rank should be 2*r. Singletons: rank = r.
    block0 = transformer.blocks[0]
    assert block0.attn1.to_qkv.lora_bank[0].rank == 3 * rank
    assert block0.attn2.to_kv.lora_bank[0].rank == 2 * rank
    assert block0.attn2.to_q.lora_bank[0].rank == rank
    assert block0.attn1.to_out.lora_bank[0].rank == rank
    assert block0.attn2.to_out.lora_bank[0].rank == rank
    assert block0.ffn.ff1.lora_bank[0].rank == rank
    assert block0.ffn.ff2.lora_bank[0].rank == rank

    # handle.rank tracks the max rank across targets (3*rank for the fused-QKV group).
    assert handle.rank == 3 * rank


@pytest.mark.parametrize("mesh_device", [(1, 1)], indirect=True)
def test_pipeline_bind_unbind_walks_all_modules(mesh_device: ttnn.MeshDevice) -> None:
    """The pipeline mixin's bind/unbind walks every LoRA module on each set_active_lora."""
    from models.tt_dit.experimental.pipelines.pipeline_wan_runtime_lora import _iter_lora_modules as pipeline_iter

    num_heads = 4
    head_dim = 32
    dim = num_heads * head_dim
    ffn_dim = 2 * dim
    rank = 8

    parallel_config = DiTParallelConfig(
        cfg_parallel=ParallelFactor(factor=1, mesh_axis=0),
        tensor_parallel=ParallelFactor(factor=1, mesh_axis=1),
        sequence_parallel=ParallelFactor(factor=1, mesh_axis=0),
    )
    ccl_manager = CCLManager(mesh_device, topology=ttnn.Topology.Linear)

    transformer = WanTransformer3DModel(
        patch_size=(1, 2, 2),
        num_heads=num_heads,
        dim=dim,
        in_channels=16,
        out_channels=16,
        text_dim=128,
        freq_dim=64,
        ffn_dim=ffn_dim,
        num_layers=1,
        mesh_device=mesh_device,
        ccl_manager=ccl_manager,
        parallel_config=parallel_config,
        is_fsdp=False,
        lora_enabled=True,
    )

    with tempfile.TemporaryDirectory() as tmp:
        adapter_path = Path(tmp) / "synthetic.safetensors"
        _save_synthetic_adapter(adapter_path, num_blocks=1, dim=dim, ffn_dim=ffn_dim, rank=rank)
        handle = load_adapter_into(transformer, str(adapter_path), name="synthetic")

    # Simulate the bind path the pipeline mixin runs.
    for dotted, bank_idx in handle.target_indices.items():
        cur = transformer
        for part in dotted.split("."):
            cur = cur[int(part)] if part.isdigit() else getattr(cur, part)
        cur.bind_active(bank_idx)

    # All modules covered by the handle should be active (delta merged).
    for dotted in handle.target_indices:
        cur = transformer
        for part in dotted.split("."):
            cur = cur[int(part)] if part.isdigit() else getattr(cur, part)
        assert cur.is_lora_active, f"{dotted} not active after bind_active"

    # Unbind everything → delta subtracted, is_lora_active False.
    for mod in pipeline_iter(transformer):
        if mod.is_lora_active:
            mod.unbind_active()
    for dotted in handle.target_indices:
        cur = transformer
        for part in dotted.split("."):
            cur = cur[int(part)] if part.isdigit() else getattr(cur, part)
        assert not cur.is_lora_active
        assert cur.active_idx is None

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

	"""
	Integration tests for the runtime-LoRA adapter loader.

	Builds a tiny WanTransformer3DModel (1 block) with ``lora_enabled=True``,
	writes a synthetic LoRA safetensors file, runs ``load_adapter_into``, and
	checks that:

	- every LoRA-aware Linear in the block got an adapter registered
	- ``bind_active`` marks the right modules active across all paths
	- ``unbind_active`` clears the active state

	The forward math is covered by ``test_lora_variants.py``; this file
	exercises the loader's key-mapping logic and the pipeline-mixin bind
	flow against the default (``fuse``) mode.

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

	import tempfile
	from pathlib import Path

	import pytest
	import torch

	import ttnn
	from models.tt_dit.experimental.lora.adapter_loader import load_adapter_into
	from models.tt_dit.layers.lora import LoRAMixin
	from models.tt_dit.models.transformers.wan2_2.transformer_wan import WanTransformer3DModel
	from models.tt_dit.parallel.config import DiTParallelConfig, ParallelFactor
	from models.tt_dit.parallel.manager import CCLManager


	def _save_synthetic_adapter(
	path: Path,
	*,
	num_blocks: int = 1,
	dim: int = 128,
	ffn_dim: int = 256,
	rank: int = 8,
	) -> None:
	"""Write a fake LoRA safetensors that covers all 6 LoRA-targeted Linears
	per block, using the lightx2v naming convention."""
	from safetensors.torch import save_file

	state: dict[str, torch.Tensor] = {}
	dtype = torch.bfloat16

	def add_pair(base: str, in_dim: int, out_dim: int):
	state[f"{base}.lora_A.weight"] = torch.randn(rank, in_dim, dtype=dtype) * 0.02
	state[f"{base}.lora_B.weight"] = torch.randn(out_dim, rank, dtype=dtype) * 0.02

	for i in range(num_blocks):
	# self-attn Q, K, V, O
	add_pair(f"blocks.{i}.self_attn.q", dim, dim)
	add_pair(f"blocks.{i}.self_attn.k", dim, dim)
	add_pair(f"blocks.{i}.self_attn.v", dim, dim)
	add_pair(f"blocks.{i}.self_attn.o", dim, dim)
	# cross-attn Q, K, V, O
	add_pair(f"blocks.{i}.cross_attn.q", dim, dim)
	add_pair(f"blocks.{i}.cross_attn.k", dim, dim)
	add_pair(f"blocks.{i}.cross_attn.v", dim, dim)
	add_pair(f"blocks.{i}.cross_attn.o", dim, dim)
	# ffn ff1, ff2 (lightx2v naming: ffn.0 = ff1, ffn.2 = ff2)
	add_pair(f"blocks.{i}.ffn.0", dim, ffn_dim)
	add_pair(f"blocks.{i}.ffn.2", ffn_dim, dim)

	save_file(state, str(path))


	def _count_registered(transformer) -> int:
	"""Count LoRA modules with at least one entry in their bank."""
	n = 0
	for mod in _iter_lora_modules(transformer):
	if any(a is not None for a in mod.lora_bank):
	n += 1
	return n


	def _iter_lora_modules(module):
	if isinstance(module, LoRAMixin):
	yield module
	for _, child in module.named_children():
	yield from _iter_lora_modules(child)


	@pytest.mark.parametrize("mesh_device", [(1, 1)], indirect=True)
	def test_adapter_loader_registers_all_lora_linears(mesh_device: ttnn.MeshDevice) -> None:
	"""Loader picks up Q/K/V/O for both attn1/attn2 plus ff1/ff2 per block."""
	num_heads = 4
	head_dim = 32
	dim = num_heads * head_dim
	ffn_dim = 2 * dim
	rank = 8
	num_blocks = 1

	parallel_config = DiTParallelConfig(
	cfg_parallel=ParallelFactor(factor=1, mesh_axis=0),
	tensor_parallel=ParallelFactor(factor=1, mesh_axis=1),
	sequence_parallel=ParallelFactor(factor=1, mesh_axis=0),
	)
	ccl_manager = CCLManager(mesh_device, topology=ttnn.Topology.Linear)

	transformer = WanTransformer3DModel(
	patch_size=(1, 2, 2),
	num_heads=num_heads,
	dim=dim,
	in_channels=16,
	out_channels=16,
	text_dim=128,
	freq_dim=64,
	ffn_dim=ffn_dim,
	num_layers=num_blocks,
	mesh_device=mesh_device,
	ccl_manager=ccl_manager,
	parallel_config=parallel_config,
	is_fsdp=False,
	lora_enabled=True,
	)

	# Every block should have 5 LoRA-aware Linears wired in: to_qkv (self),
	# to_q (cross), to_kv (cross), to_out for each attention (×2), ff1, ff2.
	# That's 1 + 1 + 1 + 2 + 1 + 1 = 7 modules per block.
	expected_per_block = 7
	n_modules = sum(1 for _ in _iter_lora_modules(transformer))
	assert (
	n_modules == num_blocks * expected_per_block
	), f"expected {num_blocks * expected_per_block} LoRA modules; got {n_modules}"

	with tempfile.TemporaryDirectory() as tmp:
	adapter_path = Path(tmp) / "synthetic.safetensors"
	_save_synthetic_adapter(adapter_path, num_blocks=num_blocks, dim=dim, ffn_dim=ffn_dim, rank=rank)

	handle = load_adapter_into(transformer, str(adapter_path), scale=1.0, name="synthetic")

	# All 7 LoRA modules per block should have a registered adapter.
	registered = _count_registered(transformer)
	assert (
	registered == num_blocks * expected_per_block
	), f"expected {num_blocks * expected_per_block} registered adapters; got {registered}"

	# Fused-QKV (to_qkv self-attn): rank should be 3*r. Fused KV (cross-attn
	# to_kv): rank should be 2*r. Singletons: rank = r.
	block0 = transformer.blocks[0]
	assert block0.attn1.to_qkv.lora_bank[0].rank == 3 * rank
	assert block0.attn2.to_kv.lora_bank[0].rank == 2 * rank
	assert block0.attn2.to_q.lora_bank[0].rank == rank
	assert block0.attn1.to_out.lora_bank[0].rank == rank
	assert block0.attn2.to_out.lora_bank[0].rank == rank
	assert block0.ffn.ff1.lora_bank[0].rank == rank
	assert block0.ffn.ff2.lora_bank[0].rank == rank

	# handle.rank tracks the max rank across targets (3*rank for the fused-QKV group).
	assert handle.rank == 3 * rank


	@pytest.mark.parametrize("mesh_device", [(1, 1)], indirect=True)
	def test_pipeline_bind_unbind_walks_all_modules(mesh_device: ttnn.MeshDevice) -> None:
	"""The pipeline mixin's bind/unbind walks every LoRA module on each set_active_lora."""
	from models.tt_dit.experimental.pipelines.pipeline_wan_runtime_lora import _iter_lora_modules as pipeline_iter

	num_heads = 4
	head_dim = 32
	dim = num_heads * head_dim
	ffn_dim = 2 * dim
	rank = 8

	parallel_config = DiTParallelConfig(
	cfg_parallel=ParallelFactor(factor=1, mesh_axis=0),
	tensor_parallel=ParallelFactor(factor=1, mesh_axis=1),
	sequence_parallel=ParallelFactor(factor=1, mesh_axis=0),
	)
	ccl_manager = CCLManager(mesh_device, topology=ttnn.Topology.Linear)

	transformer = WanTransformer3DModel(
	patch_size=(1, 2, 2),
	num_heads=num_heads,
	dim=dim,
	in_channels=16,
	out_channels=16,
	text_dim=128,
	freq_dim=64,
	ffn_dim=ffn_dim,
	num_layers=1,
	mesh_device=mesh_device,
	ccl_manager=ccl_manager,
	parallel_config=parallel_config,
	is_fsdp=False,
	lora_enabled=True,
	)

	with tempfile.TemporaryDirectory() as tmp:
	adapter_path = Path(tmp) / "synthetic.safetensors"
	_save_synthetic_adapter(adapter_path, num_blocks=1, dim=dim, ffn_dim=ffn_dim, rank=rank)
	handle = load_adapter_into(transformer, str(adapter_path), name="synthetic")

	# Simulate the bind path the pipeline mixin runs.
	for dotted, bank_idx in handle.target_indices.items():
	cur = transformer
	for part in dotted.split("."):
	cur = cur[int(part)] if part.isdigit() else getattr(cur, part)
	cur.bind_active(bank_idx)

	# All modules covered by the handle should be active (delta merged).
	for dotted in handle.target_indices:
	cur = transformer
	for part in dotted.split("."):
	cur = cur[int(part)] if part.isdigit() else getattr(cur, part)
	assert cur.is_lora_active, f"{dotted} not active after bind_active"

	# Unbind everything → delta subtracted, is_lora_active False.
	for mod in pipeline_iter(transformer):
	if mod.is_lora_active:
	mod.unbind_active()
	for dotted in handle.target_indices:
	cur = transformer
	for part in dotted.split("."):
	cur = cur[int(part)] if part.isdigit() else getattr(cur, part)
	assert not cur.is_lora_active
	assert cur.active_idx is None