File size: 4,373 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 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 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 | # 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
from models.tt_transformers.tt.mlp import MLP
from models.tt_transformers.tt.model_config import ModelArgs
from models.tt_transformers.tt.prefetcher import Prefetcher
@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(
"seq_len",
(64 * 1024, 32 * 1024, 512, 32),
)
@pytest.mark.parametrize(
"batch_size",
(1,),
)
@pytest.mark.parametrize("device_params", [{"fabric_config": True}], indirect=True)
def test_mlp_inference(seq_len, batch_size, mesh_device, reset_seeds, ensure_gc, use_prefetcher):
dtype = ttnn.bfloat8_b
mode = Mode.DECODE if seq_len <= 32 else Mode.PREFILL
# Setup prefetcher (FF1, FF2, FF3 weights are prefetched)
num_tensors = 3 if mode == Mode.DECODE else 0
prefetcher = Prefetcher(mesh_device, num_tensors=num_tensors, num_layers=1) if use_prefetcher else None
if use_prefetcher:
prefetcher.init(mode)
model_args = ModelArgs(
mesh_device,
max_batch_size=batch_size,
max_seq_len=128,
cache_hf=True,
prefetcher=prefetcher,
)
model_args.n_layers = 1
state_dict = model_args.load_state_dict()
# Ref model needs partial state dict, but our models use full state dict keys as cached weight names
first_layer_prefix = model_args.get_state_dict_prefix("MLP", 0)
partial_state_dict = {
k[len(first_layer_prefix) + 1 :]: v for k, v in state_dict.items() if (k.startswith(first_layer_prefix))
}
reference_model = model_args.reference_mlp()
reference_model.load_state_dict(partial_state_dict)
if model_args.is_90b:
# float32 ~3x faster than bfloat16.
# bfloat16 fails on CI (32k and 64k seq_len) with "This test seems to have hung... Timing out test case"
reference_model.to(torch.float32)
tt_ccl = TT_CCL(mesh_device)
tt_model = MLP(
mesh_device=mesh_device,
tt_ccl=tt_ccl,
args=model_args,
state_dict=state_dict,
weight_cache_path=model_args.weight_cache_path(dtype),
layer_num=0,
dtype=dtype,
model_config=model_args.get_model_config(),
prefetcher=prefetcher,
)
# Run prefetcher if it is used
if prefetcher is not None and mode == Mode.DECODE:
prefetcher.prefetch()
prefetcher.run()
torch_input = torch.randn(
1, 1, seq_len, model_args.dim, dtype=get_ref_model_dype(reference_model, model_args.model_name)
)
reference_output = reference_model(torch_input)
tt_input = ttnn.from_torch(
torch_input,
device=mesh_device,
mesh_mapper=ttnn.ShardTensor2dMesh(
mesh_device,
dims=(None, 3) if model_args.is_galaxy else (None, None),
mesh_shape=model_args.cluster_shape,
), # When both dims are None, the mapper used is `ReplicateTensorToMesh`
dtype=ttnn.bfloat8_b,
memory_config=model_args.get_mlp_input_mem_config(mode, prefetcher),
layout=ttnn.TILE_LAYOUT,
)
logger.info("Run MLP")
tt_output = tt_model(tt_input, mode)
tt_output_torch = ttnn.to_torch(
tt_output,
mesh_composer=ttnn.ConcatMesh2dToTensor(mesh_device, dims=(1, 3), mesh_shape=model_args.cluster_shape),
)
tt_output_torch = tt_output_torch[:, :1, :, :]
pcc_required = 0.99
passing, pcc_message = comp_pcc(reference_output, tt_output_torch, pcc_required)
logger.info(comp_allclose(reference_output, tt_output_torch))
logger.info(f"PCC: {pcc_message}")
if passing:
logger.info("MLP Passed!")
else:
logger.warning("MLP Failed!")
assert passing, f"MLP output does not meet PCC requirement {pcc_required}: {pcc_message}."
|