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