# 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!"