File size: 3,380 Bytes
3e0ec22 | 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 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 | # 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.tt.embedding import Embedding, ScaledEmbedding
from models.tt_transformers.tt.model_config import ModelArgs
@torch.no_grad()
@pytest.mark.parametrize("use_scaled_embedding", (False, True))
@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(
"batch_size",
(1,),
)
@pytest.mark.parametrize(
"max_seq_len",
(128,), # For decode-only unit test, there's no need to run with large sequence lengths
)
def test_embedding(max_seq_len, batch_size, mesh_device, reset_seeds, ensure_gc, use_scaled_embedding):
dtype = ttnn.bfloat16
model_args = ModelArgs(mesh_device, max_batch_size=batch_size, max_seq_len=max_seq_len, cache_hf=True)
model_args.n_layers = 1
state_dict = model_args.load_state_dict()
tokenizer = model_args.tokenizer
if use_scaled_embedding:
model_args.embed_scale = model_args.dim**0.5
logger.info(f"Using scaled embedding with scale {model_args.embed_scale}")
reference_emb = model_args.reference_embedding()
if model_args.is_llama_vision():
weight = torch.cat(
[
state_dict["text_model.tok_embeddings.weight"],
state_dict["text_model.learnable_embedding.weight"],
],
dim=0,
)
else:
weight = state_dict["tok_embeddings.weight"]
reference_emb.load_state_dict({"emb.weight": weight})
emb_kwargs = {
"mesh_device": mesh_device,
"args": model_args,
"weight_cache_path": model_args.weight_cache_path(dtype),
"state_dict": state_dict,
"dtype": dtype,
}
if use_scaled_embedding:
emb_kwargs["embed_scale"] = model_args.embed_scale
emb_cls = ScaledEmbedding
else:
emb_cls = Embedding
tt_emb = emb_cls(**emb_kwargs)
prompts = ["Joy"] * 32
pt_input = torch.tensor([model_args.encode_prompt(prompt, instruct=False) for prompt in prompts])
reference_output = reference_emb(pt_input)
logger.info(f"reference_output: {reference_output.shape}")
tt_input = ttnn.from_torch(
pt_input.squeeze(1),
device=mesh_device,
mesh_mapper=ttnn.ReplicateTensorToMesh(mesh_device),
dtype=ttnn.uint32,
layout=ttnn.ROW_MAJOR_LAYOUT,
)
tt_output = tt_emb(tt_input)
tt_output_torch = ttnn.to_torch(
tt_output,
mesh_composer=ttnn.ConcatMesh2dToTensor(mesh_device, dims=(0, -1), mesh_shape=model_args.cluster_shape),
)[:32].view(reference_output.shape)
logger.info(f"tt_output_torch: {tt_output_torch.shape}")
passing, pcc_message = comp_pcc(reference_output, tt_output_torch)
logger.info(comp_allclose(reference_output, tt_output_torch))
logger.info(f"PCC: {pcc_message}")
if passing:
logger.info("embedding Passed!")
else:
logger.warning("embedding Failed!")
assert passing, f"embedding output does not meet PCC requirement {0.99}."
|