File size: 2,608 Bytes
c92e699 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 | # SPDX-FileCopyrightText: © 2023 Tenstorrent USA, Inc.
# SPDX-License-Identifier: Apache-2.0
import torch
from loguru import logger
# import ttnn
from models.tt_transformers.tt.model_config import ModelArgs
@torch.no_grad()
def test_torch_inference(ensure_gc):
iterations = 20
model_args = ModelArgs(mesh_device=None, cache_hf=True)
state_dict = model_args.load_state_dict()
tokenizer = model_args.tokenizer
prompts = ["1 2 3 4 "] * model_args.max_batch_size
encoded_prompts = [model_args.encode_prompt(prompt, instruct=False) for prompt in prompts]
reference_model = model_args.reference_transformer()
reference_model.load_state_dict(state_dict)
# Embedding on host
embd = model_args.reference_embedding()
state_dict_prefix = model_args.get_state_dict_prefix("", None)
embd.load_state_dict({"emb.weight": state_dict[f"{state_dict_prefix}tok_embeddings.weight"]})
generation_start_pos = 0
generation_length = iterations
seqlen = 1 # Generating one token per user at a time
# Select the first token from the prompts for initial decoding
encoded_prompts_tensor = torch.tensor(encoded_prompts) # [:,0]
pt_decode_input = embd(encoded_prompts_tensor[:, 0]).view(model_args.max_batch_size, seqlen, -1)
logger.info(pt_decode_input.shape)
all_outputs_ref = []
for i in range(generation_length):
logger.info(f"[Decode] Generating token {i}")
start_pos = generation_start_pos + i
ref_output = reference_model(pt_decode_input, start_pos)
# While in "prefill" mode, use the prompt tokens as the output
if i in range(len(encoded_prompts[0])):
all_outputs_ref.append(encoded_prompts[0][i]) # Update list of ref outputs
pt_decode_input = embd(encoded_prompts_tensor[:, i]).view(model_args.max_batch_size, seqlen, -1)
else:
# pt_out_tok = torch.argmax(torch.nn.functional.log_softmax(ref_output, dim=-1), dim=-1)
pt_out_tok = torch.argmax(ref_output, dim=-1)
# pt_out_tok_logscores = top_k_top_p_filtering(ref_output.squeeze(1), top_k=0, top_p=0.9)
# probs = torch.nn.functional.softmax(pt_out_tok_logscores, dim=-1)
# pt_out_tok = torch.multinomial(probs, num_samples=1)#.squeeze(1)
pt_decode_input = embd(pt_out_tok)
all_outputs_ref.append(pt_out_tok.squeeze(1).tolist()[0]) # Update generated token to list of ref outputs
# TODO print all 32 users
logger.info("[User 0] Ref generation: '" + "".join(tokenizer.decode(all_outputs_ref)) + "'")
|