code/models/tt_dit/tests/unit/test_linear.py · sunandbean/diffrhythm-1.2-p100a at main

diffrhythm-1.2-p100a / code /models /tt_dit /tests /unit /test_linear.py

Add files using upload-large-folder tool

33d23e1 verified 11 days ago

9.47 kB

# SPDX-FileCopyrightText: © 2025 Tenstorrent USA, Inc.

# SPDX-License-Identifier: Apache-2.0


import pytest
import torch

import ttnn

from ...layers.linear import ColParallelLinear, Linear, RowParallelLinear
from ...parallel.manager import CCLManager
from ...utils.check import assert_quality
from ...utils.tensor import bf16_tensor
from ...utils.test import mesh_device_config_to_string


@pytest.mark.parametrize(
    "mesh_device",
    [(1, 1), (1, 2), (2, 1)],
    ids=mesh_device_config_to_string,
    indirect=True,
)
@pytest.mark.parametrize(
    ("B, M, K, N"),
    [
        (1, 1, 256, 2432),  # SD3.5 timestep_embedder
        (1, 1, 2432, 2432),  # SD3.5 timestep_embedder
        (1, 333, 2048, 2432),  # SD3.5 text_embedder
        (1, 333, 2432, 2432),  # SD3.5 text_embedder
        (1, 4096, 2432, 2432 * 3),  # SD3.5 spatial qkv
        (1, 333, 2432, 2432 * 3),  # SD3.5 text qkv
        (1, 4096, 2432, 2432),  # SD3.5 spatial out_proj
        (1, 333, 2432, 2432),  # SD3.5 text out_proj
        (1, 4096, 2432, 9728),  # SD3.5 spatial FF
        (1, 4096, 9728, 2432),  # SD3.5 spatial FF
        (1, 333, 2432, 9728),  # SD3.5 text FF
        (1, 333, 9728, 2432),  # SD3.5 text FF
        (1, 1, 2432, 14592),  # SD3.5 context
        (1, 1, 2432, 4864),  # SD3.5 final context
        (1, 4096, 2432, 64),  # SD3.5 proj_out
        # Flux2 M=32 shapes — exercise the 1D matmul branch (M ≤ 64)
        (1, 1, 256, 6144),  # timestep/guidance embed linear_1
        (1, 1, 6144, 6144),  # timestep/guidance embed linear_2
        (1, 1, 6144, 4608),  # double_stream_mod img / txt
        (1, 1, 6144, 2304),  # single_stream_mod
        (1, 1, 6144, 1536),  # spatial_time_embed_out
    ],
)
@pytest.mark.parametrize(
    ("bias"),
    [
        True,
        # False,
    ],
)
def test_linear(
    mesh_device: ttnn.MeshDevice,
    B: int,
    M: int,
    K: int,
    N: int,
    bias: bool,
) -> None:
    torch_dtype = torch.bfloat16
    torch_model = torch.nn.Linear(K, N, bias=bias).to(dtype=torch_dtype)
    torch_model.eval()

    tt_model = Linear(K, N, bias=bias, mesh_device=mesh_device)
    tt_model.load_torch_state_dict(torch_model.state_dict())

    torch_input_tensor = torch.randn((1, B, M, K), dtype=torch_dtype)

    tt_input_tensor = bf16_tensor(torch_input_tensor, device=mesh_device)

    with torch.no_grad():
        torch_output = torch_model(torch_input_tensor)

    tt_output = tt_model(tt_input_tensor)

    for t in ttnn.get_device_tensors(tt_output):
        t = ttnn.to_torch(t)
        assert_quality(torch_output, t, pcc=0.999_500)


@pytest.mark.parametrize(
    "mesh_device, device_params",
    [
        (
            shape,
            {
                "fabric_config": None if shape == (1, 1) else ttnn.FabricConfig.FABRIC_1D,
                "require_exact_physical_num_devices": False if shape == (1, 1) else True,
            },
        )
        for shape in [(1, 1), (1, 2), (2, 1), (2, 2), (2, 4), (4, 2)]
    ],
    ids=mesh_device_config_to_string,
    indirect=True,
)
@pytest.mark.parametrize(
    ("tp_mesh_axis"),
    [
        0,
        1,
    ],
)
@pytest.mark.parametrize("is_fsdp", [True, False], ids=["yes_fsdp", "no_fsdp"])
@pytest.mark.parametrize(
    ("B, M, K, N"),
    [
        (1, 1, 256, 2432),  # SD3.5 timestep_embedder
        (1, 1, 2432, 2432),  # SD3.5 timestep_embedder
        (1, 333, 2048, 2432),  # SD3.5 text_embedder
        (1, 333, 2432, 2432),  # SD3.5 text_embedder
        (1, 4096, 2432, 2432 * 3),  # SD3.5 spatial qkv
        (1, 333, 2432, 2432 * 3),  # SD3.5 text qkv
        (1, 4096, 2432, 2432),  # SD3.5 spatial out_proj
        (1, 333, 2432, 2432),  # SD3.5 text out_proj
        (1, 4096, 2432, 9728),  # SD3.5 spatial FF
        (1, 4096, 9728, 2432),  # SD3.5 spatial FF
        (1, 333, 2432, 9728),  # SD3.5 text FF
        (1, 333, 9728, 2432),  # SD3.5 text FF
        (1, 1, 2432, 14592),  # SD3.5 context
        (1, 1, 2432, 4864),  # SD3.5 final context
        (1, 1, 3072, 12288),  # Mochi timestep norm1 linear
        (1, 1, 3072, 6144),  # Mochi timestep norm1_context linear
        (1, 44520, 3072, 3072),  # Mochi spatial Q, K, V and out_proj
        (1, 118, 1536, 3072),  # Mochi text Q, K, V
        (1, 118, 3072, 1536),  # Mochi text out_proj
    ],
    ids=[
        "sd35_timestep_embedder",
        "sd35_timestep_embedder_2",
        "sd35_text_embedder",
        "sd35_text_embedder_2",
        "sd35_spatial_qkv",
        "sd35_text_qkv",
        "sd35_spatial_out_proj",
        "sd35_text_out_proj",
        "sd35_spatial_FF",
        "sd35_spatial_FF_2",
        "sd35_text_FF",
        "sd35_text_FF_2",
        "sd35_context",
        "sd35_final_context",
        "mochi_timestep_norm1_linear",
        "mochi_timestep_norm1_context_linear",
        "mochi_spatial_qkv_and_out_proj",
        "mochi_text_qkv",
        "mochi_text_out_proj",
    ],
)
@pytest.mark.parametrize(
    ("bias"),
    [
        True,
        # False,
    ],
)
def test_col_parallel_linear(
    mesh_device: ttnn.MeshDevice,
    B: int,
    M: int,
    K: int,
    N: int,
    bias: bool,
    tp_mesh_axis: int,
    is_fsdp: bool,
) -> None:
    torch_dtype = torch.bfloat16
    torch_model = torch.nn.Linear(K, N, bias=bias).to(dtype=torch_dtype)
    torch_model.eval()

    if is_fsdp:
        fsdp_mesh_axis = 1 - tp_mesh_axis
        ccl_manager = CCLManager(mesh_device, topology=ttnn.Topology.Linear)
        tt_model = ColParallelLinear(
            K,
            N,
            bias=bias,
            mesh_device=mesh_device,
            mesh_axis=tp_mesh_axis,
            fsdp_mesh_axis=fsdp_mesh_axis,
            ccl_manager=ccl_manager,
        )
    else:
        tt_model = ColParallelLinear(K, N, bias=bias, mesh_device=mesh_device, mesh_axis=tp_mesh_axis)
    tt_model.load_torch_state_dict(torch_model.state_dict())

    torch_input_tensor = torch.randn((1, B, M, K), dtype=torch_dtype)

    tt_input_tensor = bf16_tensor(torch_input_tensor, device=mesh_device)

    with torch.no_grad():
        torch_output = torch_model(torch_input_tensor)

    tt_output = tt_model(tt_input_tensor)

    shard_dims = [None, None]
    shard_dims[tp_mesh_axis] = -1
    shard_dims[1 - tp_mesh_axis] = 0
    tt_output = ttnn.to_torch(
        tt_output,
        mesh_composer=ttnn.ConcatMesh2dToTensor(mesh_device, dims=shard_dims, mesh_shape=tuple(mesh_device.shape)),
    )
    for i in range(tt_output.shape[0]):
        assert_quality(torch_output.squeeze(), tt_output[i].squeeze(), pcc=0.999_500)


@pytest.mark.parametrize(
    "mesh_device, device_params",
    [
        (
            shape,
            {
                "fabric_config": None if shape == (1, 1) else ttnn.FabricConfig.FABRIC_1D,
                "require_exact_physical_num_devices": False if shape == (1, 1) else True,
            },
        )
        for shape in [(1, 1), (1, 2), (2, 1), (2, 2), (2, 4), (4, 2)]
    ],
    ids=mesh_device_config_to_string,
    indirect=True,
)
@pytest.mark.parametrize(
    ("tp_mesh_axis"),
    [
        0,
        1,
    ],
)
@pytest.mark.parametrize("is_fsdp", [True, False], ids=["yes_fsdp", "no_fsdp"])
@pytest.mark.parametrize(
    ("B, M, K, N"),
    [
        (1, 1, 256, 2432),  # SD3.5 timestep_embedder
        (1, 1, 2432, 2432),  # SD3.5 timestep_embedder
        (1, 333, 2048, 2432),  # SD3.5 text_embedder
        (1, 333, 2432, 2432),  # SD3.5 text_embedder
        (1, 4096, 2432, 2432 * 3),  # SD3.5 spatial qkv
        (1, 333, 2432, 2432 * 3),  # SD3.5 text qkv
        (1, 4096, 2432, 2432),  # SD3.5 spatial out_proj
        (1, 333, 2432, 2432),  # SD3.5 text out_proj
        (1, 4096, 2432, 9728),  # SD3.5 spatial FF
        (1, 4096, 9728, 2432),  # SD3.5 spatial FF
        (1, 333, 2432, 9728),  # SD3.5 text FF
        (1, 333, 9728, 2432),  # SD3.5 text FF
        (1, 1, 2432, 14592),  # SD3.5 context
        (1, 1, 2432, 4864),  # SD3.5 final context
    ],
)
@pytest.mark.parametrize(
    ("bias"),
    [
        True,
        # False,
    ],
)
def test_row_parallel_linear(
    mesh_device: ttnn.MeshDevice,
    B: int,
    M: int,
    K: int,
    N: int,
    bias: bool,
    tp_mesh_axis: int,
    is_fsdp: bool,
) -> None:
    torch_dtype = torch.bfloat16
    torch_model = torch.nn.Linear(K, N, bias=bias).to(dtype=torch_dtype)
    torch_model.eval()

    fsdp_mesh_axis = 1 - tp_mesh_axis if is_fsdp else None
    ccl_manager = CCLManager(mesh_device, topology=ttnn.Topology.Linear)
    tt_model = RowParallelLinear(
        K,
        N,
        bias=bias,
        mesh_device=mesh_device,
        mesh_axis=tp_mesh_axis,
        fsdp_mesh_axis=fsdp_mesh_axis,
        ccl_manager=ccl_manager,
    )
    tt_model.load_torch_state_dict(torch_model.state_dict())

    torch_input_tensor = torch.randn((1, B, M, K), dtype=torch_dtype)

    tt_input_tensor = bf16_tensor(torch_input_tensor, device=mesh_device, mesh_axis=tp_mesh_axis, shard_dim=-1)

    with torch.no_grad():
        torch_output = torch_model(torch_input_tensor)

    tt_output = tt_model(tt_input_tensor)

    shard_dims = [None, None]
    shard_dims[tp_mesh_axis] = -1
    shard_dims[1 - tp_mesh_axis] = 0
    tt_output = ttnn.to_torch(
        tt_output,
        mesh_composer=ttnn.ConcatMesh2dToTensor(mesh_device, dims=shard_dims, mesh_shape=tuple(mesh_device.shape)),
    )
    for i in range(tt_output.shape[0]):
        assert_quality(torch_output.squeeze(), tt_output[i].squeeze(), pcc=0.999_500)

	# SPDX-FileCopyrightText: © 2025 Tenstorrent USA, Inc.

	# SPDX-License-Identifier: Apache-2.0


	import pytest
	import torch

	import ttnn

	from ...layers.linear import ColParallelLinear, Linear, RowParallelLinear
	from ...parallel.manager import CCLManager
	from ...utils.check import assert_quality
	from ...utils.tensor import bf16_tensor
	from ...utils.test import mesh_device_config_to_string


	@pytest.mark.parametrize(
	"mesh_device",
	[(1, 1), (1, 2), (2, 1)],
	ids=mesh_device_config_to_string,
	indirect=True,
	)
	@pytest.mark.parametrize(
	("B, M, K, N"),
	[
	(1, 1, 256, 2432), # SD3.5 timestep_embedder
	(1, 1, 2432, 2432), # SD3.5 timestep_embedder
	(1, 333, 2048, 2432), # SD3.5 text_embedder
	(1, 333, 2432, 2432), # SD3.5 text_embedder
	(1, 4096, 2432, 2432 * 3), # SD3.5 spatial qkv
	(1, 333, 2432, 2432 * 3), # SD3.5 text qkv
	(1, 4096, 2432, 2432), # SD3.5 spatial out_proj
	(1, 333, 2432, 2432), # SD3.5 text out_proj
	(1, 4096, 2432, 9728), # SD3.5 spatial FF
	(1, 4096, 9728, 2432), # SD3.5 spatial FF
	(1, 333, 2432, 9728), # SD3.5 text FF
	(1, 333, 9728, 2432), # SD3.5 text FF
	(1, 1, 2432, 14592), # SD3.5 context
	(1, 1, 2432, 4864), # SD3.5 final context
	(1, 4096, 2432, 64), # SD3.5 proj_out
	# Flux2 M=32 shapes — exercise the 1D matmul branch (M ≤ 64)
	(1, 1, 256, 6144), # timestep/guidance embed linear_1
	(1, 1, 6144, 6144), # timestep/guidance embed linear_2
	(1, 1, 6144, 4608), # double_stream_mod img / txt
	(1, 1, 6144, 2304), # single_stream_mod
	(1, 1, 6144, 1536), # spatial_time_embed_out
	],
	)
	@pytest.mark.parametrize(
	("bias"),
	[
	True,
	# False,
	],
	)
	def test_linear(
	mesh_device: ttnn.MeshDevice,
	B: int,
	M: int,
	K: int,
	N: int,
	bias: bool,
	) -> None:
	torch_dtype = torch.bfloat16
	torch_model = torch.nn.Linear(K, N, bias=bias).to(dtype=torch_dtype)
	torch_model.eval()

	tt_model = Linear(K, N, bias=bias, mesh_device=mesh_device)
	tt_model.load_torch_state_dict(torch_model.state_dict())

	torch_input_tensor = torch.randn((1, B, M, K), dtype=torch_dtype)

	tt_input_tensor = bf16_tensor(torch_input_tensor, device=mesh_device)

	with torch.no_grad():
	torch_output = torch_model(torch_input_tensor)

	tt_output = tt_model(tt_input_tensor)

	for t in ttnn.get_device_tensors(tt_output):
	t = ttnn.to_torch(t)
	assert_quality(torch_output, t, pcc=0.999_500)


	@pytest.mark.parametrize(
	"mesh_device, device_params",
	[
	(
	shape,
	{
	"fabric_config": None if shape == (1, 1) else ttnn.FabricConfig.FABRIC_1D,
	"require_exact_physical_num_devices": False if shape == (1, 1) else True,
	},
	)
	for shape in [(1, 1), (1, 2), (2, 1), (2, 2), (2, 4), (4, 2)]
	],
	ids=mesh_device_config_to_string,
	indirect=True,
	)
	@pytest.mark.parametrize(
	("tp_mesh_axis"),
	[
	0,
	1,
	],
	)
	@pytest.mark.parametrize("is_fsdp", [True, False], ids=["yes_fsdp", "no_fsdp"])
	@pytest.mark.parametrize(
	("B, M, K, N"),
	[
	(1, 1, 256, 2432), # SD3.5 timestep_embedder
	(1, 1, 2432, 2432), # SD3.5 timestep_embedder
	(1, 333, 2048, 2432), # SD3.5 text_embedder
	(1, 333, 2432, 2432), # SD3.5 text_embedder
	(1, 4096, 2432, 2432 * 3), # SD3.5 spatial qkv
	(1, 333, 2432, 2432 * 3), # SD3.5 text qkv
	(1, 4096, 2432, 2432), # SD3.5 spatial out_proj
	(1, 333, 2432, 2432), # SD3.5 text out_proj
	(1, 4096, 2432, 9728), # SD3.5 spatial FF
	(1, 4096, 9728, 2432), # SD3.5 spatial FF
	(1, 333, 2432, 9728), # SD3.5 text FF
	(1, 333, 9728, 2432), # SD3.5 text FF
	(1, 1, 2432, 14592), # SD3.5 context
	(1, 1, 2432, 4864), # SD3.5 final context
	(1, 1, 3072, 12288), # Mochi timestep norm1 linear
	(1, 1, 3072, 6144), # Mochi timestep norm1_context linear
	(1, 44520, 3072, 3072), # Mochi spatial Q, K, V and out_proj
	(1, 118, 1536, 3072), # Mochi text Q, K, V
	(1, 118, 3072, 1536), # Mochi text out_proj
	],
	ids=[
	"sd35_timestep_embedder",
	"sd35_timestep_embedder_2",
	"sd35_text_embedder",
	"sd35_text_embedder_2",
	"sd35_spatial_qkv",
	"sd35_text_qkv",
	"sd35_spatial_out_proj",
	"sd35_text_out_proj",
	"sd35_spatial_FF",
	"sd35_spatial_FF_2",
	"sd35_text_FF",
	"sd35_text_FF_2",
	"sd35_context",
	"sd35_final_context",
	"mochi_timestep_norm1_linear",
	"mochi_timestep_norm1_context_linear",
	"mochi_spatial_qkv_and_out_proj",
	"mochi_text_qkv",
	"mochi_text_out_proj",
	],
	)
	@pytest.mark.parametrize(
	("bias"),
	[
	True,
	# False,
	],
	)
	def test_col_parallel_linear(
	mesh_device: ttnn.MeshDevice,
	B: int,
	M: int,
	K: int,
	N: int,
	bias: bool,
	tp_mesh_axis: int,
	is_fsdp: bool,
	) -> None:
	torch_dtype = torch.bfloat16
	torch_model = torch.nn.Linear(K, N, bias=bias).to(dtype=torch_dtype)
	torch_model.eval()

	if is_fsdp:
	fsdp_mesh_axis = 1 - tp_mesh_axis
	ccl_manager = CCLManager(mesh_device, topology=ttnn.Topology.Linear)
	tt_model = ColParallelLinear(
	K,
	N,
	bias=bias,
	mesh_device=mesh_device,
	mesh_axis=tp_mesh_axis,
	fsdp_mesh_axis=fsdp_mesh_axis,
	ccl_manager=ccl_manager,
	)
	else:
	tt_model = ColParallelLinear(K, N, bias=bias, mesh_device=mesh_device, mesh_axis=tp_mesh_axis)
	tt_model.load_torch_state_dict(torch_model.state_dict())

	torch_input_tensor = torch.randn((1, B, M, K), dtype=torch_dtype)

	tt_input_tensor = bf16_tensor(torch_input_tensor, device=mesh_device)

	with torch.no_grad():
	torch_output = torch_model(torch_input_tensor)

	tt_output = tt_model(tt_input_tensor)

	shard_dims = [None, None]
	shard_dims[tp_mesh_axis] = -1
	shard_dims[1 - tp_mesh_axis] = 0
	tt_output = ttnn.to_torch(
	tt_output,
	mesh_composer=ttnn.ConcatMesh2dToTensor(mesh_device, dims=shard_dims, mesh_shape=tuple(mesh_device.shape)),
	)
	for i in range(tt_output.shape[0]):
	assert_quality(torch_output.squeeze(), tt_output[i].squeeze(), pcc=0.999_500)


	@pytest.mark.parametrize(
	"mesh_device, device_params",
	[
	(
	shape,
	{
	"fabric_config": None if shape == (1, 1) else ttnn.FabricConfig.FABRIC_1D,
	"require_exact_physical_num_devices": False if shape == (1, 1) else True,
	},
	)
	for shape in [(1, 1), (1, 2), (2, 1), (2, 2), (2, 4), (4, 2)]
	],
	ids=mesh_device_config_to_string,
	indirect=True,
	)
	@pytest.mark.parametrize(
	("tp_mesh_axis"),
	[
	0,
	1,
	],
	)
	@pytest.mark.parametrize("is_fsdp", [True, False], ids=["yes_fsdp", "no_fsdp"])
	@pytest.mark.parametrize(
	("B, M, K, N"),
	[
	(1, 1, 256, 2432), # SD3.5 timestep_embedder
	(1, 1, 2432, 2432), # SD3.5 timestep_embedder
	(1, 333, 2048, 2432), # SD3.5 text_embedder
	(1, 333, 2432, 2432), # SD3.5 text_embedder
	(1, 4096, 2432, 2432 * 3), # SD3.5 spatial qkv
	(1, 333, 2432, 2432 * 3), # SD3.5 text qkv
	(1, 4096, 2432, 2432), # SD3.5 spatial out_proj
	(1, 333, 2432, 2432), # SD3.5 text out_proj
	(1, 4096, 2432, 9728), # SD3.5 spatial FF
	(1, 4096, 9728, 2432), # SD3.5 spatial FF
	(1, 333, 2432, 9728), # SD3.5 text FF
	(1, 333, 9728, 2432), # SD3.5 text FF
	(1, 1, 2432, 14592), # SD3.5 context
	(1, 1, 2432, 4864), # SD3.5 final context
	],
	)
	@pytest.mark.parametrize(
	("bias"),
	[
	True,
	# False,
	],
	)
	def test_row_parallel_linear(
	mesh_device: ttnn.MeshDevice,
	B: int,
	M: int,
	K: int,
	N: int,
	bias: bool,
	tp_mesh_axis: int,
	is_fsdp: bool,
	) -> None:
	torch_dtype = torch.bfloat16
	torch_model = torch.nn.Linear(K, N, bias=bias).to(dtype=torch_dtype)
	torch_model.eval()

	fsdp_mesh_axis = 1 - tp_mesh_axis if is_fsdp else None
	ccl_manager = CCLManager(mesh_device, topology=ttnn.Topology.Linear)
	tt_model = RowParallelLinear(
	K,
	N,
	bias=bias,
	mesh_device=mesh_device,
	mesh_axis=tp_mesh_axis,
	fsdp_mesh_axis=fsdp_mesh_axis,
	ccl_manager=ccl_manager,
	)
	tt_model.load_torch_state_dict(torch_model.state_dict())

	torch_input_tensor = torch.randn((1, B, M, K), dtype=torch_dtype)

	tt_input_tensor = bf16_tensor(torch_input_tensor, device=mesh_device, mesh_axis=tp_mesh_axis, shard_dim=-1)

	with torch.no_grad():
	torch_output = torch_model(torch_input_tensor)

	tt_output = tt_model(tt_input_tensor)

	shard_dims = [None, None]
	shard_dims[tp_mesh_axis] = -1
	shard_dims[1 - tp_mesh_axis] = 0
	tt_output = ttnn.to_torch(
	tt_output,
	mesh_composer=ttnn.ConcatMesh2dToTensor(mesh_device, dims=shard_dims, mesh_shape=tuple(mesh_device.shape)),
	)
	for i in range(tt_output.shape[0]):
	assert_quality(torch_output.squeeze(), tt_output[i].squeeze(), pcc=0.999_500)