File size: 1,513 Bytes
b025706 | 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 | # SPDX-FileCopyrightText: © 2023 Tenstorrent USA, Inc.
# SPDX-License-Identifier: Apache-2.0
import ttnn
from models.common.lightweightmodule import LightweightModule
class TtLlamaClassEmbedding(LightweightModule):
"""Class Embedding layer."""
def __init__(
self,
mesh_device,
state_dict,
state_dict_prefix,
weight_cache_path,
dtype,
configuration,
):
super().__init__()
self.mesh_device = mesh_device
# Add batch and ntoks dimensions
class_embedding = state_dict[f"{state_dict_prefix}class_embedding"]
class_embedding = class_embedding.reshape(1, 1, 1, *class_embedding.shape)
self.class_embedding = ttnn.as_tensor(
class_embedding,
dtype=ttnn.bfloat16,
layout=ttnn.ROW_MAJOR_LAYOUT, # [INFO] TILE_LAYOUT shouldn't be used here because of the `dim=2` concat that this tensor is used in
device=self.mesh_device,
memory_config=ttnn.DRAM_MEMORY_CONFIG,
mesh_mapper=ttnn.ReplicateTensorToMesh(self.mesh_device),
)
def forward(self, x):
bsz = x.shape[1]
# Broadcast class embedding to match input batch size
class_embedding = ttnn.concat([self.class_embedding] * bsz, dim=1) # Broadcast batch size
x = ttnn.concat([class_embedding, x], dim=2) # Output ROW_MAJOR_LAYOUT
x = ttnn.to_layout(x, layout=ttnn.TILE_LAYOUT) # Convert back to TILE_LAYOUT
return x
|