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

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

Add files using upload-large-folder tool

33d23e1 verified 12 days ago

10.3 kB

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

# SPDX-License-Identifier: Apache-2.0


import pytest
import torch

# This is the Conv3d used by the mochi models in huggingface's diffusers module.
from diffusers.models.autoencoders.autoencoder_kl_cogvideox import CogVideoXCausalConv3d as RefContextParallelConv3d
from loguru import logger

import ttnn

from ...layers.conv3d import ContextParallelConv3d as TtContextParallelConv3d
from ...parallel.config import MochiVAEParallelConfig, ParallelFactor
from ...parallel.manager import CCLManager
from ...utils.check import assert_quality
from ...utils.tensor import print_tensor_mem_info

# Common test configurations
PCC_REQUIRED = 0.999

conv3d_args = {
    "context_parallel": True,
    "causal": True,
    "padding_mode": "replicate",
    "bias": True,
    "groups": 1,
}


def create_random_models(
    mesh_device,
    **model_args,
):
    """Initialize both reference and TT models."""
    # CogVideoXCausalConv3d takes a narrower kwarg set than genmo's class: no
    # `causal`/`context_parallel`/`bias`/`groups` knobs (causality is always on;
    # bias=True and groups=1 are hardcoded). Rename `padding_mode` -> `pad_mode`.
    ref_args = {
        "in_channels": model_args["in_channels"],
        "out_channels": model_args["out_channels"],
        "kernel_size": model_args["kernel_size"],
        "stride": model_args["stride"],
        "pad_mode": model_args["padding_mode"],
    }
    reference_model = RefContextParallelConv3d(**ref_args)

    # Create TT model
    tt_model = TtContextParallelConv3d(mesh_device=mesh_device, **model_args)
    # Diffusers nests the conv at `self.conv`; flatten to TT's `weight`/`bias`.
    ref_state = {k.removeprefix("conv."): v for k, v in reference_model.state_dict().items()}
    tt_model.load_torch_state_dict(ref_state)

    return reference_model, tt_model


def validate_outputs(tt_output, ref_output, test_name):
    """Validate and compare model outputs."""
    pcc, mse, mae = compute_metrics(ref_output, tt_output)
    logger.info(f"Output - PCC: {pcc}, MSE: {mse}, MAE: {mae}")

    passing = pcc >= PCC_REQUIRED

    if passing:
        logger.info(f"{test_name} Passed!")
    else:
        logger.warning(f"{test_name} Failed!")
        logger.error(f"Output failed with PCC: {pcc}, MSE: {mse}, MAE: {mae}")

    assert passing, f"{test_name} output does not meet PCC requirement {PCC_REQUIRED}"


# TODO: Merge this with the conv3d tests in test_vae_mochi.py
@torch.no_grad()
@pytest.mark.parametrize(
    "input_shape, out_channels, kernel_size, stride",
    [
        [(1, 768, 28, 60, 106), 768, (3, 3, 3), (1, 1, 1)],
        #  [(1, 512, 82, 120, 212), 512, (3, 3, 3), (1, 1, 1)],
    ],
    ids=["768"],  # , "512"],
)
@pytest.mark.parametrize(
    "mesh_device",
    [(1, 1)],
    indirect=True,
)
def test_context_parallel_conv3d_forward_noshard(
    mesh_device, input_shape, out_channels, kernel_size, stride, reset_seeds
):
    input_channels = input_shape[1]
    model_args = conv3d_args.copy()
    model_args.update(
        {
            "in_channels": input_channels,
            "out_channels": out_channels,
            "kernel_size": kernel_size,
            "stride": stride,
            "parallel_config": MochiVAEParallelConfig(
                time_parallel=ParallelFactor(factor=1, mesh_axis=1),
                h_parallel=ParallelFactor(factor=1, mesh_axis=0),
                w_parallel=ParallelFactor(factor=1, mesh_axis=1),
            ),
            "ccl_manager": CCLManager(
                mesh_device=mesh_device,
                topology=ttnn.Topology.Linear,
            ),
        }
    )

    reference_model, tt_model = create_random_models(mesh_device, **model_args)

    # Create input tensor (NCTHW format for PyTorch)
    torch_input = torch.randn(input_shape)

    # Convert to NTHWC format for TT
    tt_input_NTHWC = torch_input.permute(0, 2, 3, 4, 1)
    logger.info(f"Creating tensor: {tt_input_NTHWC.shape}")
    tt_input_NTHWC = ttnn.from_torch(
        tt_input_NTHWC,
        device=mesh_device,
        dtype=ttnn.bfloat16,
        layout=ttnn.TILE_LAYOUT,
        memory_config=ttnn.DRAM_MEMORY_CONFIG,
        mesh_mapper=None,
    )
    tt_input_NTHWC = ttnn.to_layout(tt_input_NTHWC, ttnn.ROW_MAJOR_LAYOUT)

    logger.info("Run TtContextParallelConv3d forward")
    tt_output = tt_model(tt_input_NTHWC)

    # Convert TT output to torch tensor (from NTHWC to NCHW format)
    tt_output_torch = ttnn.to_torch(ttnn.get_device_tensors(tt_output)[0])
    tt_output_torch = tt_output_torch.permute(0, 4, 1, 2, 3)  # NTHWC -> NCHW

    # Get reference output
    logger.info("Run reference model forward")
    with torch.no_grad():
        ref_output, _ = reference_model(torch_input)  # returns (output, new_conv_cache)

    # Compare shapes
    logger.info(f"TT output shape: {tt_output_torch.shape}, Ref output shape: {ref_output.shape}")
    assert tt_output_torch.shape == ref_output.shape, "Output shapes don't match"

    assert_quality(ref_output, tt_output_torch, pcc=PCC_REQUIRED)


# Original test iterated over the following arches/device combos. Unclear how to
# iterate in this exact fashion within the current test framework.
# {
#     "N150": (1, 1),
#     "N300": (1, 2),
#     "T3K": (1, 1),
#     "TG": (8, 4)
# }
@torch.no_grad()
@pytest.mark.parametrize(
    "input_shape, out_channels, kernel_size, stride",
    [
        [(1, 768, 28, 60, 106), 768, (3, 3, 3), (1, 1, 1)],
        [(1, 512, 82, 120, 212), 512, (3, 3, 3), (1, 1, 1)],
        [(1, 256, 163, 240, 424), 256, (3, 3, 3), (1, 1, 1)],
        [(1, 128, 163, 480, 848), 128, (3, 3, 3), (1, 1, 1)],
    ],
    ids=["768", "512", "256", "128"],
)
@pytest.mark.parametrize(
    "mesh_device",
    [(1, 1)],  # TODO: make this the full suite of meshes want to test.
    indirect=True,
)
@pytest.mark.skip(
    reason=(
        "This test is currently busted. "
        "Some  tt tensors hang on calls to from_torch,"
        "The sharding is also not quite right as written. "
        "Fix this test when we bring up the mochi model tests, as "
        "test_vae_mochi.py has a lot of overlap and we can consolidate them."
    )
)
def test_context_parallel_conv3d_forward(mesh_device, input_shape, out_channels, kernel_size, stride, reset_seeds):
    """Test complete forward pass of TtContextParallelConv3d."""
    input_channels = input_shape[1]

    model_args = conv3d_args.copy()
    model_args.update(
        {
            "in_channels": input_channels,
            "out_channels": out_channels,
            "kernel_size": kernel_size,
            "stride": stride,
            # TODO: see test_vae_mochi.py for how to call these with various mesh device shapes
            # "parallel_config": MochiVAEParallelConfig(
            #     time_parallel=ParallelFactor(factor=4, mesh_axis=1),
            #     h_parallel=ParallelFactor(factor=1, mesh_axis=0),
            #     w_parallel=ParallelFactor(factor=1, mesh_axis=1),
            # ),
            # "ccl_manager": CCLManager(
            #     mesh_device=mesh_device,
            #     topology=ttnn.Topology.Linear,
            # ),
        }
    )
    # Create the models
    reference_model, tt_model = create_random_models(mesh_device, **model_args)

    # Create input tensor (NCTHW format for PyTorch)
    torch_input = torch.randn(input_shape)

    # Convert to NTHWC format for TT
    tt_input_NTHWC = torch_input.permute(0, 2, 3, 4, 1)

    # TODO: see test_vae_mochi.py for how to set this up.
    dims = [None, 1]
    mesh_mapper = ttnn.create_mesh_mapper(
        mesh_device,
        ttnn.MeshMapperConfig(
            [
                ttnn.PlacementReplicate() if dims[0] is None else ttnn.PlacementShard(dims[0]),
                ttnn.PlacementReplicate() if dims[1] is None else ttnn.PlacementShard(dims[1]),
            ],
            ttnn.MeshShape(mesh_device.shape[0], mesh_device.shape[1]),
        ),
    )

    tt_input_NTHWC = ttnn.from_torch(
        tt_input_NTHWC,
        device=mesh_device,
        dtype=ttnn.bfloat16,
        layout=ttnn.TILE_LAYOUT,
        memory_config=ttnn.DRAM_MEMORY_CONFIG,
        mesh_mapper=mesh_mapper,
    )

    tt_input_NTHWC = ttnn.to_layout(tt_input_NTHWC, ttnn.ROW_MAJOR_LAYOUT)

    logger.info("Run TtContextParallelConv3d forward")
    tt_output = tt_model(tt_input_NTHWC)

    # Convert TT output to torch tensor (from NTHWC to NCHW format)
    tt_output_torch = ttnn.to_torch(ttnn.get_device_tensors(tt_output)[0])
    tt_output_torch = tt_output_torch.permute(0, 4, 1, 2, 3)  # NTHWC -> NCHW

    # Get reference output
    logger.info("Run reference model forward")
    with torch.no_grad():
        ref_output, _ = reference_model(torch_input)  # returns (output, new_conv_cache)

    # Compare shapes
    logger.info(f"TT output shape: {tt_output_torch.shape}, Ref output shape: {ref_output.shape}")
    assert tt_output_torch.shape == ref_output.shape, "Output shapes don't match"

    assert_quality(ref_output, tt_output_torch, pcc=PCC_REQUIRED)


@torch.no_grad()
@pytest.mark.parametrize(
    "input_shape, out_channels, kernel_size, stride",
    [
        [(1, 512, 82, 120, 212), 512, (3, 3, 3), (1, 1, 1)],
    ],
    ids=["512"],
)
@pytest.mark.parametrize(
    "mesh_device",
    [(1, 1)],
    indirect=True,
)
@pytest.mark.skip(reason="This allocation hangs on blackhole. TODO Debug and fix in the future.")
def test_sharding(mesh_device, input_shape, out_channels, kernel_size, stride):
    # Create input tensor (NCTHW format for PyTorch)
    torch_input = torch.randn(input_shape)

    logger.info(f"initial shape {torch_input.shape}")
    # Convert to NTHWC format for TT
    tt_input_NTHWC = torch_input.permute(0, 2, 3, 4, 1).contiguous()

    logger.info(f"initial reshape {tt_input_NTHWC.shape}")
    logger.info(f"contiguous? : {tt_input_NTHWC.is_contiguous()}")

    tt_input_NTHWC = ttnn.from_torch(
        tt_input_NTHWC,
        device=mesh_device,
        dtype=ttnn.bfloat16,
        layout=ttnn.ROW_MAJOR_LAYOUT,
        memory_config=ttnn.DRAM_MEMORY_CONFIG,
        mesh_mapper=None,
    )
    logger.info("Create DONE!")
    print_tensor_mem_info(tt_input_NTHWC)

    logger.info("reorder")
    tt_input_NTHWC_rm = ttnn.to_layout(tt_input_NTHWC, ttnn.ROW_MAJOR_LAYOUT)
    print_tensor_mem_info(tt_input_NTHWC_rm)

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

	# SPDX-License-Identifier: Apache-2.0


	import pytest
	import torch

	# This is the Conv3d used by the mochi models in huggingface's diffusers module.
	from diffusers.models.autoencoders.autoencoder_kl_cogvideox import CogVideoXCausalConv3d as RefContextParallelConv3d
	from loguru import logger

	import ttnn

	from ...layers.conv3d import ContextParallelConv3d as TtContextParallelConv3d
	from ...parallel.config import MochiVAEParallelConfig, ParallelFactor
	from ...parallel.manager import CCLManager
	from ...utils.check import assert_quality
	from ...utils.tensor import print_tensor_mem_info

	# Common test configurations
	PCC_REQUIRED = 0.999

	conv3d_args = {
	"context_parallel": True,
	"causal": True,
	"padding_mode": "replicate",
	"bias": True,
	"groups": 1,
	}


	def create_random_models(
	mesh_device,
	**model_args,
	):
	"""Initialize both reference and TT models."""
	# CogVideoXCausalConv3d takes a narrower kwarg set than genmo's class: no
	# `causal`/`context_parallel`/`bias`/`groups` knobs (causality is always on;
	# bias=True and groups=1 are hardcoded). Rename `padding_mode` -> `pad_mode`.
	ref_args = {
	"in_channels": model_args["in_channels"],
	"out_channels": model_args["out_channels"],
	"kernel_size": model_args["kernel_size"],
	"stride": model_args["stride"],
	"pad_mode": model_args["padding_mode"],
	}
	reference_model = RefContextParallelConv3d(**ref_args)

	# Create TT model
	tt_model = TtContextParallelConv3d(mesh_device=mesh_device, **model_args)
	# Diffusers nests the conv at `self.conv`; flatten to TT's `weight`/`bias`.
	ref_state = {k.removeprefix("conv."): v for k, v in reference_model.state_dict().items()}
	tt_model.load_torch_state_dict(ref_state)

	return reference_model, tt_model


	def validate_outputs(tt_output, ref_output, test_name):
	"""Validate and compare model outputs."""
	pcc, mse, mae = compute_metrics(ref_output, tt_output)
	logger.info(f"Output - PCC: {pcc}, MSE: {mse}, MAE: {mae}")

	passing = pcc >= PCC_REQUIRED

	if passing:
	logger.info(f"{test_name} Passed!")
	else:
	logger.warning(f"{test_name} Failed!")
	logger.error(f"Output failed with PCC: {pcc}, MSE: {mse}, MAE: {mae}")

	assert passing, f"{test_name} output does not meet PCC requirement {PCC_REQUIRED}"


	# TODO: Merge this with the conv3d tests in test_vae_mochi.py
	@torch.no_grad()
	@pytest.mark.parametrize(
	"input_shape, out_channels, kernel_size, stride",
	[
	[(1, 768, 28, 60, 106), 768, (3, 3, 3), (1, 1, 1)],
	# [(1, 512, 82, 120, 212), 512, (3, 3, 3), (1, 1, 1)],
	],
	ids=["768"], # , "512"],
	)
	@pytest.mark.parametrize(
	"mesh_device",
	[(1, 1)],
	indirect=True,
	)
	def test_context_parallel_conv3d_forward_noshard(
	mesh_device, input_shape, out_channels, kernel_size, stride, reset_seeds
	):
	input_channels = input_shape[1]
	model_args = conv3d_args.copy()
	model_args.update(
	{
	"in_channels": input_channels,
	"out_channels": out_channels,
	"kernel_size": kernel_size,
	"stride": stride,
	"parallel_config": MochiVAEParallelConfig(
	time_parallel=ParallelFactor(factor=1, mesh_axis=1),
	h_parallel=ParallelFactor(factor=1, mesh_axis=0),
	w_parallel=ParallelFactor(factor=1, mesh_axis=1),
	),
	"ccl_manager": CCLManager(
	mesh_device=mesh_device,
	topology=ttnn.Topology.Linear,
	),
	}
	)

	reference_model, tt_model = create_random_models(mesh_device, **model_args)

	# Create input tensor (NCTHW format for PyTorch)
	torch_input = torch.randn(input_shape)

	# Convert to NTHWC format for TT
	tt_input_NTHWC = torch_input.permute(0, 2, 3, 4, 1)
	logger.info(f"Creating tensor: {tt_input_NTHWC.shape}")
	tt_input_NTHWC = ttnn.from_torch(
	tt_input_NTHWC,
	device=mesh_device,
	dtype=ttnn.bfloat16,
	layout=ttnn.TILE_LAYOUT,
	memory_config=ttnn.DRAM_MEMORY_CONFIG,
	mesh_mapper=None,
	)
	tt_input_NTHWC = ttnn.to_layout(tt_input_NTHWC, ttnn.ROW_MAJOR_LAYOUT)

	logger.info("Run TtContextParallelConv3d forward")
	tt_output = tt_model(tt_input_NTHWC)

	# Convert TT output to torch tensor (from NTHWC to NCHW format)
	tt_output_torch = ttnn.to_torch(ttnn.get_device_tensors(tt_output)[0])
	tt_output_torch = tt_output_torch.permute(0, 4, 1, 2, 3) # NTHWC -> NCHW

	# Get reference output
	logger.info("Run reference model forward")
	with torch.no_grad():
	ref_output, _ = reference_model(torch_input) # returns (output, new_conv_cache)

	# Compare shapes
	logger.info(f"TT output shape: {tt_output_torch.shape}, Ref output shape: {ref_output.shape}")
	assert tt_output_torch.shape == ref_output.shape, "Output shapes don't match"

	assert_quality(ref_output, tt_output_torch, pcc=PCC_REQUIRED)


	# Original test iterated over the following arches/device combos. Unclear how to
	# iterate in this exact fashion within the current test framework.
	# {
	# "N150": (1, 1),
	# "N300": (1, 2),
	# "T3K": (1, 1),
	# "TG": (8, 4)
	# }
	@torch.no_grad()
	@pytest.mark.parametrize(
	"input_shape, out_channels, kernel_size, stride",
	[
	[(1, 768, 28, 60, 106), 768, (3, 3, 3), (1, 1, 1)],
	[(1, 512, 82, 120, 212), 512, (3, 3, 3), (1, 1, 1)],
	[(1, 256, 163, 240, 424), 256, (3, 3, 3), (1, 1, 1)],
	[(1, 128, 163, 480, 848), 128, (3, 3, 3), (1, 1, 1)],
	],
	ids=["768", "512", "256", "128"],
	)
	@pytest.mark.parametrize(
	"mesh_device",
	[(1, 1)], # TODO: make this the full suite of meshes want to test.
	indirect=True,
	)
	@pytest.mark.skip(
	reason=(
	"This test is currently busted. "
	"Some tt tensors hang on calls to from_torch,"
	"The sharding is also not quite right as written. "
	"Fix this test when we bring up the mochi model tests, as "
	"test_vae_mochi.py has a lot of overlap and we can consolidate them."
	)
	)
	def test_context_parallel_conv3d_forward(mesh_device, input_shape, out_channels, kernel_size, stride, reset_seeds):
	"""Test complete forward pass of TtContextParallelConv3d."""
	input_channels = input_shape[1]

	model_args = conv3d_args.copy()
	model_args.update(
	{
	"in_channels": input_channels,
	"out_channels": out_channels,
	"kernel_size": kernel_size,
	"stride": stride,
	# TODO: see test_vae_mochi.py for how to call these with various mesh device shapes
	# "parallel_config": MochiVAEParallelConfig(
	# time_parallel=ParallelFactor(factor=4, mesh_axis=1),
	# h_parallel=ParallelFactor(factor=1, mesh_axis=0),
	# w_parallel=ParallelFactor(factor=1, mesh_axis=1),
	# ),
	# "ccl_manager": CCLManager(
	# mesh_device=mesh_device,
	# topology=ttnn.Topology.Linear,
	# ),
	}
	)
	# Create the models
	reference_model, tt_model = create_random_models(mesh_device, **model_args)

	# Create input tensor (NCTHW format for PyTorch)
	torch_input = torch.randn(input_shape)

	# Convert to NTHWC format for TT
	tt_input_NTHWC = torch_input.permute(0, 2, 3, 4, 1)

	# TODO: see test_vae_mochi.py for how to set this up.
	dims = [None, 1]
	mesh_mapper = ttnn.create_mesh_mapper(
	mesh_device,
	ttnn.MeshMapperConfig(
	[
	ttnn.PlacementReplicate() if dims[0] is None else ttnn.PlacementShard(dims[0]),
	ttnn.PlacementReplicate() if dims[1] is None else ttnn.PlacementShard(dims[1]),
	],
	ttnn.MeshShape(mesh_device.shape[0], mesh_device.shape[1]),
	),
	)

	tt_input_NTHWC = ttnn.from_torch(
	tt_input_NTHWC,
	device=mesh_device,
	dtype=ttnn.bfloat16,
	layout=ttnn.TILE_LAYOUT,
	memory_config=ttnn.DRAM_MEMORY_CONFIG,
	mesh_mapper=mesh_mapper,
	)

	tt_input_NTHWC = ttnn.to_layout(tt_input_NTHWC, ttnn.ROW_MAJOR_LAYOUT)

	logger.info("Run TtContextParallelConv3d forward")
	tt_output = tt_model(tt_input_NTHWC)

	# Convert TT output to torch tensor (from NTHWC to NCHW format)
	tt_output_torch = ttnn.to_torch(ttnn.get_device_tensors(tt_output)[0])
	tt_output_torch = tt_output_torch.permute(0, 4, 1, 2, 3) # NTHWC -> NCHW

	# Get reference output
	logger.info("Run reference model forward")
	with torch.no_grad():
	ref_output, _ = reference_model(torch_input) # returns (output, new_conv_cache)

	# Compare shapes
	logger.info(f"TT output shape: {tt_output_torch.shape}, Ref output shape: {ref_output.shape}")
	assert tt_output_torch.shape == ref_output.shape, "Output shapes don't match"

	assert_quality(ref_output, tt_output_torch, pcc=PCC_REQUIRED)


	@torch.no_grad()
	@pytest.mark.parametrize(
	"input_shape, out_channels, kernel_size, stride",
	[
	[(1, 512, 82, 120, 212), 512, (3, 3, 3), (1, 1, 1)],
	],
	ids=["512"],
	)
	@pytest.mark.parametrize(
	"mesh_device",
	[(1, 1)],
	indirect=True,
	)
	@pytest.mark.skip(reason="This allocation hangs on blackhole. TODO Debug and fix in the future.")
	def test_sharding(mesh_device, input_shape, out_channels, kernel_size, stride):
	# Create input tensor (NCTHW format for PyTorch)
	torch_input = torch.randn(input_shape)

	logger.info(f"initial shape {torch_input.shape}")
	# Convert to NTHWC format for TT
	tt_input_NTHWC = torch_input.permute(0, 2, 3, 4, 1).contiguous()

	logger.info(f"initial reshape {tt_input_NTHWC.shape}")
	logger.info(f"contiguous? : {tt_input_NTHWC.is_contiguous()}")

	tt_input_NTHWC = ttnn.from_torch(
	tt_input_NTHWC,
	device=mesh_device,
	dtype=ttnn.bfloat16,
	layout=ttnn.ROW_MAJOR_LAYOUT,
	memory_config=ttnn.DRAM_MEMORY_CONFIG,
	mesh_mapper=None,
	)
	logger.info("Create DONE!")
	print_tensor_mem_info(tt_input_NTHWC)

	logger.info("reorder")
	tt_input_NTHWC_rm = ttnn.to_layout(tt_input_NTHWC, ttnn.ROW_MAJOR_LAYOUT)
	print_tensor_mem_info(tt_input_NTHWC_rm)