code/models/tt_transformers/tests/test_decoder.py

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

# SPDX-License-Identifier: Apache-2.0
import os

import pytest
import torch
from loguru import logger

import ttnn
from models.common.utility_functions import comp_allclose, comp_pcc
from models.tt_transformers.tests.test_utils import get_ref_model_dype
from models.tt_transformers.tt.ccl import TT_CCL
from models.tt_transformers.tt.common import Mode, PagedAttentionConfig, precompute_freqs
from models.tt_transformers.tt.decoder import TransformerBlock
from models.tt_transformers.tt.model_config import ModelArgs
from models.tt_transformers.tt.prefetcher import Prefetcher
from models.tt_transformers.tt.rope import HfRotarySetup, RotarySetup


@torch.no_grad()
@pytest.mark.parametrize(
    "use_prefetcher",
    ([False]),
)
@pytest.mark.parametrize(
    "mesh_device",
    [
        {"N150": (1, 1), "N300": (1, 2), "T3K": (1, 8), "TG": (8, 4)}.get(
            os.environ.get("MESH_DEVICE"), len(ttnn.get_device_ids())
        )
    ],
    indirect=True,
)
@pytest.mark.parametrize(
    "paged_attention",
    (
        True,
        # False
    ),
    ids=(
        "paged_attention",
        # "default_attention"
    ),
)
@pytest.mark.parametrize(
    "page_params",
    [{"page_block_size": 32, "page_max_num_blocks": 1024}],
)
@pytest.mark.parametrize(
    "batch_size",
    (1, 32),
)
@pytest.mark.parametrize(
    "max_seq_len",
    (256,),  # For decode-only unit test, there's no need to run with large sequence lengths
)
@pytest.mark.parametrize(
    "generation_length",
    (10,),  # For decode-only unit test, there's no need to run with large sequence lengths
)
@pytest.mark.parametrize("device_params", [{"fabric_config": True}], indirect=True)
def test_decoder_inference(
    max_seq_len,
    batch_size,
    paged_attention,
    page_params,
    mesh_device,
    reset_seeds,
    ensure_gc,
    generation_length,
    use_prefetcher,
):
    dtype = ttnn.bfloat8_b

    mode = Mode.DECODE
    num_tensors = 5 if use_prefetcher else 0
    prefetcher = Prefetcher(mesh_device, num_tensors=num_tensors, num_layers=1) if use_prefetcher else None

    if use_prefetcher:
        prefetcher.init(mode=mode)

    model_args = ModelArgs(
        mesh_device,
        max_batch_size=batch_size,
        max_seq_len=max_seq_len,
        cache_hf=True,
        prefetcher=prefetcher,
        use_hf_rope=False,
    )
    model_args.n_layers = 1

    state_dict = model_args.load_state_dict()
    reference_model = model_args.reference_decoder(load_checkpoint=True)

    generation_start_pos = 0
    all_tests_pass = True

    # Setup RoPE transformation matrices
    DefaultRopeSetup = HfRotarySetup if model_args.use_hf_rope else RotarySetup
    rope_setup = DefaultRopeSetup(
        mesh_device,
        model_args.max_batch_size,
        model_args.head_dim,
        model_args.max_seq_len,
        model_args.rope_theta,
        model_args.rope_scaling,
        model_args.use_qk_fused,
        prefetcher=prefetcher,
    )

    if model_args.rope_theta_local is not None:
        rope_setup_local = RotarySetup(
            mesh_device,
            model_args.max_batch_size,
            model_args.head_dim,
            model_args.max_seq_len,
            model_args.rope_theta_local,
            None,
            # Match model.py: when fused QK is on, the local rope must also double
            # its cos/sin batch so one tensor serves both Q and K (Gemma-2 sliding layers).
            use_qk_fused=model_args.use_qk_fused,
        )
    else:
        rope_setup_local = None

    transformation_mats = rope_setup.get_both_trans_mats()

    # Prepare page table for paged attention
    page_table_tt = None
    paged_attention_config = None

    if paged_attention:
        paged_attention_config = PagedAttentionConfig(
            block_size=page_params["page_block_size"],
            max_num_blocks=page_params["page_max_num_blocks"],
        )
        # Implied shuffling of blocks
        permutation = torch.randperm(paged_attention_config.max_num_blocks)
        # Page table which maps virtual blocks to physical
        reverse_permutation = torch.argsort(permutation)
        page_table = reverse_permutation.reshape(
            model_args.max_batch_size, paged_attention_config.max_num_blocks // model_args.max_batch_size
        )
        page_table_tt = ttnn.from_torch(
            page_table,
            device=mesh_device,
            dtype=ttnn.int32,
            layout=ttnn.ROW_MAJOR_LAYOUT,
            mesh_mapper=ttnn.ShardTensor2dMesh(
                mesh_device,
                dims=(None, -2) if (model_args.is_galaxy and batch_size > 1) else (None, None),
                mesh_shape=model_args.cluster_shape,
            ),
        )

    # Initialize TT model
    tt_ccl = TT_CCL(mesh_device)
    tt_model = TransformerBlock(
        args=model_args,
        mesh_device=mesh_device,
        tt_ccl=tt_ccl,
        dtype=dtype,
        state_dict=state_dict,
        layer_num=0,
        weight_cache_path=model_args.weight_cache_path(dtype),
        transformation_mats=transformation_mats,
        paged_attention_config=paged_attention_config,
        prefetcher=prefetcher,
    )
    if use_prefetcher:
        tt_model.prefetcher.prefetch()

    seqlen = 1

    # Precompute freqs_cis for reference model
    cos, sin = precompute_freqs(
        model_args.head_dim,
        model_args.max_seq_len * 2,
        model_args.rope_theta,
        model_args.rope_scaling.factor if model_args.rope_scaling else None,
        model_args.rope_scaling.original_max_position_embeddings if model_args.rope_scaling else None,
        model_args.rope_scaling.rope_type.value if model_args.rope_scaling else "llama3",
    )
    freqs_cis = torch.complex(cos, sin)

    # Initial positions
    current_pos = torch.tensor([generation_start_pos for _ in range(batch_size)])
    current_pos_tensor = ttnn.from_torch(
        current_pos,
        device=mesh_device,
        dtype=ttnn.int32,
        mesh_mapper=ttnn.ShardTensor2dMesh(
            mesh_device,
            dims=(None, 0) if (model_args.is_galaxy and batch_size > 1) else (None, None),
            mesh_shape=model_args.cluster_shape,
        ),
    )
    for i in range(generation_length):
        logger.info(f"[Decoder] Generating token {i}")

        if prefetcher is not None:
            prefetcher.run()

        # input = torch.randn(1, 32, 4096)
        pt_decode_input = (
            torch.rand(
                batch_size, seqlen, model_args.dim, dtype=get_ref_model_dype(reference_model, model_args.model_name)
            )
            * 2
        ) - 1
        tt_decode_input = pt_decode_input.clone()

        decode_input = model_args.prepare_residual_tensor_decode(
            tt_decode_input,
            model_args.get_residual_mem_config(mode, prefetcher),
        )

        # Get cos/sin matrices for the current position of each user
        rot_mats = rope_setup.get_rot_mats(current_pos)
        rot_mats_local = None if rope_setup_local is None else rope_setup_local.get_rot_mats(current_pos)

        # Run TT model
        tt_out = tt_model(
            decode_input,
            current_pos_tensor,
            rot_mats_global=rot_mats,
            rot_mats_local=rot_mats_local,
            mode=mode,
            page_table=page_table_tt,
        )

        tt_out = ttnn.to_torch(
            tt_out,
            mesh_composer=ttnn.ConcatMesh2dToTensor(mesh_device, dims=(1, 3), mesh_shape=model_args.cluster_shape),
        )

        tt_output_torch = tt_out[:, 0:1, : model_args.max_batch_size, : model_args.dim].view(-1, 1, model_args.dim)

        # In this test all users have the same position
        freqs_cis_i = freqs_cis[current_pos[0], :].unsqueeze(0)

        # Reference model
        ref_output = reference_model(pt_decode_input, current_pos[0], freqs_cis_i, mask=None)
        if ref_output.dim() == 2:
            ref_output = ref_output.unsqueeze(1)

        # For some model variants the HF decoder returns output only for the first batch item.
        # Since all users share the same position in this test, compare the first ref_output.shape[0]
        # items from TT output to ref_output.
        batch_cmp = ref_output.shape[0]
        tt_output_cmp = tt_output_torch[:batch_cmp]
        passing, pcc_message = comp_pcc(ref_output, tt_output_cmp)

        logger.info(comp_allclose(ref_output, tt_output_cmp))
        logger.info(f"PCC: {pcc_message}")

        if passing:
            logger.info("Decoder Block Passed!")
        else:
            logger.warning("Decoder Block Failed!")
            all_tests_pass = False

        # Increment position
        current_pos = torch.tensor([generation_start_pos + i + 1 for _ in range(batch_size)])
        current_pos_tensor = ttnn.from_torch(
            current_pos,
            device=mesh_device,
            dtype=ttnn.int32,
            mesh_mapper=ttnn.ShardTensor2dMesh(
                mesh_device,
                dims=(None, 0) if (model_args.is_galaxy and batch_size > 1) else (None, None),
                mesh_shape=model_args.cluster_shape,
            ),
        )

    if all_tests_pass:
        logger.info(f"All {generation_length} decode iterations Passed!")
    else:
        logger.warning("One or more iterations of decode Failed!")
        assert all_tests_pass, f"PCC value is lower than {0.99} for some of the outputs. Check Warnings!"