File size: 17,287 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 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 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 347 348 349 350 351 352 353 354 355 356 357 358 359 360 361 362 363 364 365 366 367 368 369 370 371 372 373 374 375 376 377 378 379 380 381 382 383 384 385 386 387 388 389 390 391 392 393 394 395 396 397 398 399 400 401 402 403 404 405 406 407 408 409 410 411 412 413 | # SPDX-FileCopyrightText: © 2024 Tenstorrent USA, Inc.
# SPDX-License-Identifier: Apache-2.0
import ttnn
from models.common.lightweightmodule import LightweightModule
from models.common.rmsnorm import RMSNorm
from models.tt_transformers.tt.attention import Attention as DefaultAttention
from models.tt_transformers.tt.common import Mode
from models.tt_transformers.tt.distributed_norm import DistributedNorm
from models.tt_transformers.tt.mixtral_mlp import TtMixtralMLP
from models.tt_transformers.tt.mixtral_moe import TtMoeLayer
from models.tt_transformers.tt.mlp import MLP
from models.tt_transformers.tt.model_config import TensorGroup
class TransformerBlock(LightweightModule):
def __init__(
self,
args,
mesh_device,
tt_ccl,
dtype,
state_dict,
layer_num,
weight_cache_path,
transformation_mats,
paged_attention_config=None,
use_paged_kv_cache=False,
attention_class=None,
prefetcher=None,
):
super().__init__()
self.mesh_device = mesh_device
self.tt_ccl = tt_ccl
self.prefetcher = prefetcher
self.num_devices = args.num_devices
self.args = args
self.hidden_size = args.dim
self.n_heads = args.n_heads
self.head_dim = self.hidden_size // self.n_heads
self.max_seq_len = args.max_seq_len
self.dim = args.dim
self.max_batch_size = args.max_batch_size
self.n_kv_heads = args.n_kv_heads
self.current = 0
self.model_config = args.get_model_config()
self.is_mixture_of_experts = False
self.layer_num = layer_num
ActualAttentionClass = attention_class if attention_class is not None else DefaultAttention
self.attention = ActualAttentionClass(
mesh_device=mesh_device,
tt_ccl=self.tt_ccl,
args=args,
state_dict=state_dict,
weight_cache_path=weight_cache_path,
layer_num=layer_num,
dtype=dtype,
transformation_mats=transformation_mats,
configuration=args,
paged_attention_config=paged_attention_config,
use_paged_kv_cache=use_paged_kv_cache,
prefetcher=prefetcher,
)
if getattr(self.args, "is_mixture_of_experts", False):
self.feed_forward = TtMoeLayer(
mesh_device=mesh_device,
state_dict=state_dict,
experts=TtMixtralMLP(
mesh_device=mesh_device,
state_dict=state_dict,
args=args,
layer_num=layer_num,
dtypes={
"w1": dtype,
"w2": dtype,
"w3": dtype,
},
),
args=args,
layer_num=layer_num,
dtype=dtype,
tt_ccl=self.tt_ccl,
)
else:
self.feed_forward = MLP(
mesh_device=mesh_device,
tt_ccl=self.tt_ccl,
args=args,
state_dict=state_dict,
weight_cache_path=weight_cache_path,
layer_num=layer_num,
dtype=dtype,
model_config=self.model_config,
prefetcher=prefetcher,
)
# TODO: remove after https://github.com/tenstorrent/tt-metal/issues/35650 is fixed
extra_rmsnorm_kwargs = {}
# Llama 8B on a Galaxy DP4 row submesh runs out of L1 with fp32 RMSNorm
# accumulation, matching the existing Qwen workaround below.
use_galaxy_row_submesh_rmsnorm_l1_workaround = (
args.base_model_name == "Llama-3.1-8B"
and args.num_devices == 8
and args.mesh_device is not None
and tuple(args.mesh_device.shape) == (1, 8)
and ttnn.cluster.get_cluster_type() == ttnn.cluster.ClusterType.GALAXY
)
if (
args.base_model_name
in (
"Qwen2.5-7B",
"Qwen2.5-VL-7B",
)
or use_galaxy_row_submesh_rmsnorm_l1_workaround
):
extra_rmsnorm_kwargs["fp32_dest_acc_en"] = False
# Post-norm decoders (EXAONE-4.x) have no input_layernorm: attention reads
# the raw residual stream (h = x + post_attention_layernorm(attn(x))). The
# attention_norm slot degenerates to its gather role (norm=None), keeping
# the fractured->replicated all-gather the norm normally provides.
self.use_post_norm = getattr(args, "use_post_norm", False)
self.attention_norm = DistributedNorm(
None
if self.use_post_norm
else RMSNorm(
device=mesh_device,
dim=args.dim,
eps=args.norm_eps,
state_dict=state_dict,
state_dict_prefix=args.get_state_dict_prefix("", layer_num),
weight_cache_path=None if args.dummy_weights else weight_cache_path,
weight_dtype=ttnn.bfloat16,
weight_key="attention_norm",
is_distributed=self.args.is_distributed_norm,
add_unit_offset=self.args.rms_norm_add_unit_offset,
ccl_topology=self.args.ccl_topology(),
tt_ccl=self.tt_ccl,
**extra_rmsnorm_kwargs,
),
args,
tt_ccl=self.tt_ccl,
prefetcher=self.prefetcher,
TG=args.is_galaxy,
ag_config_key="ATTN_LN_AG_CONFIG",
)
self.ff_norm = DistributedNorm(
RMSNorm(
device=mesh_device,
dim=args.dim,
eps=args.norm_eps,
state_dict=state_dict,
state_dict_prefix=args.get_state_dict_prefix("", layer_num),
weight_cache_path=None if args.dummy_weights else weight_cache_path,
weight_dtype=ttnn.bfloat16,
weight_key="ffn_norm",
is_distributed=self.args.is_distributed_norm,
add_unit_offset=self.args.rms_norm_add_unit_offset,
ccl_topology=self.args.ccl_topology(),
tt_ccl=self.tt_ccl,
**extra_rmsnorm_kwargs,
),
args,
tt_ccl=self.tt_ccl,
prefetcher=self.prefetcher,
TG=args.is_galaxy,
ag_config_key="FFN_LN_AG_CONFIG",
)
if f"layers.{layer_num}.pre_feedforward_layernorm.weight" in state_dict:
self.pre_ff_norm = DistributedNorm( # pre_feedforward_layernorm
RMSNorm(
device=mesh_device,
dim=args.dim,
eps=args.norm_eps,
state_dict=state_dict,
add_unit_offset=self.args.rms_norm_add_unit_offset,
state_dict_prefix=args.get_state_dict_prefix("", layer_num),
weight_cache_path=None if args.dummy_weights else weight_cache_path,
weight_dtype=ttnn.bfloat16,
weight_key="pre_feedforward_layernorm",
is_distributed=self.args.is_distributed_norm,
ccl_topology=self.args.ccl_topology(),
tt_ccl=self.tt_ccl,
),
args,
tt_ccl=self.tt_ccl,
prefetcher=self.prefetcher,
TG=args.is_galaxy,
)
self.ff_norm.enable_all_gather = (
False # output of ff_norm should be sharded if model uses pre_ff_norm, so skip all_gather
)
elif self.use_post_norm:
# Post-norm (EXAONE-4.x): route forward() through the sandwich branch so
# ff_norm (HF post_attention_layernorm) is applied to the attention OUTPUT
# before the residual add, with a gather-only identity in the pre-MLP slot
# (MLP reads the raw residual stream: out = h + post_ff_norm(mlp(h))).
self.pre_ff_norm = DistributedNorm(
None,
args,
tt_ccl=self.tt_ccl,
prefetcher=self.prefetcher,
TG=args.is_galaxy,
)
self.ff_norm.enable_all_gather = (
False # output of ff_norm should be sharded if model uses pre_ff_norm, so skip all_gather
)
else:
# If pre_feedforward_layernorm is not in state_dict, we do not use it
self.pre_ff_norm = None
if f"layers.{layer_num}.post_feedforward_layernorm.weight" in state_dict:
self.post_ff_norm = DistributedNorm( # post_feedforward_layernorm
RMSNorm(
device=mesh_device,
dim=args.dim,
eps=args.norm_eps,
add_unit_offset=self.args.rms_norm_add_unit_offset,
state_dict=state_dict,
state_dict_prefix=args.get_state_dict_prefix("", layer_num),
weight_cache_path=None if args.dummy_weights else weight_cache_path,
weight_dtype=ttnn.bfloat16,
weight_key="post_feedforward_layernorm",
is_distributed=self.args.is_distributed_norm,
ccl_topology=self.args.ccl_topology(),
tt_ccl=self.tt_ccl,
),
args,
tt_ccl=self.tt_ccl,
prefetcher=self.prefetcher,
TG=args.is_galaxy,
enable_all_gather=False,
)
else:
# If post_feedforward_layernorm is not in state_dict, we do not use it
self.post_ff_norm = None
def update_weights(
self,
layer_hf_state_dict: dict[str, ttnn.Tensor],
*,
hf_rope: bool = False,
) -> None:
"""Strict layer-local weight update from an HF-keyed dict of on-device
4D ttnn tensors. Keys are the suffix after ``model.layers.{i}.`` (e.g.
``self_attn.q_proj.weight``); ``Transformer.update_weights`` strips the
layer prefix and routes each layer's slice here.
Strict: missing required keys raise ``KeyError``; any unconsumed key
(e.g. an extra bias, q/k-norm, or Gemma-style pre/post-FF norm not yet
wired into the leaf ``update()`` methods) raises ``ValueError``.
``hf_rope`` is forwarded to ``Attention.update``.
"""
unconsumed = set(layer_hf_state_dict.keys())
def consume(key: str) -> ttnn.Tensor:
if key not in layer_hf_state_dict:
raise KeyError(
f"TransformerBlock.update_weights (layer {self.layer_num}): " f"missing required HF key {key!r}"
)
unconsumed.discard(key)
return layer_hf_state_dict[key]
self.attention.update(
q_proj=consume("self_attn.q_proj.weight"),
k_proj=consume("self_attn.k_proj.weight"),
v_proj=consume("self_attn.v_proj.weight"),
o_proj=consume("self_attn.o_proj.weight"),
hf_rope=hf_rope,
)
self.feed_forward.update(
gate_proj=consume("mlp.gate_proj.weight"),
up_proj=consume("mlp.up_proj.weight"),
down_proj=consume("mlp.down_proj.weight"),
)
self.attention_norm.update(weight=consume("input_layernorm.weight"))
self.ff_norm.update(weight=consume("post_attention_layernorm.weight"))
if unconsumed:
sample = sorted(unconsumed)[:10]
raise ValueError(
f"TransformerBlock.update_weights (layer {self.layer_num}): "
f"{len(unconsumed)} HF key(s) not consumed within this layer. "
f"This usually means an extra weight (a bias, q/k norm, or "
f"Gemma-style pre/post-FF norm) that the leaf .update() "
f"doesn't yet support. Showing up to 10: {sample}"
)
def forward(
self,
x: ttnn.Tensor,
current_pos,
rot_mats_global=None,
rot_mats_local=None,
user_id=0,
mode="decode",
page_table=None,
chunk_page_table=None,
chunk_start_idx=None,
kv_cache=None,
batch_size=1,
) -> ttnn.Tensor:
TG = self.args.is_galaxy
residual = x
# x is fractured across devices and interleaved in DRAM (for prefill) and sharded in L1 (for decode)
skip_mem_cfg = self.args.get_residual_mem_config(mode, self.prefetcher)
assert (
x.memory_config() == skip_mem_cfg
), f"decoder input memcfg mismatch: {x.memory_config()} != {skip_mem_cfg}"
# Choose the correct rotation matrices based on the mode
rot_mats = (
rot_mats_local if (hasattr(self.attention, "is_sliding") and self.attention.is_sliding) else rot_mats_global
)
# Norms take fractured inputs and output replicated across devices
attn_norm_config = self.args.get_norm_config("attn", mode, self.prefetcher)
attn_in = self.attention_norm(x, mode, norm_config=attn_norm_config)
# Reshape to [B, 1, S_per_user, H] so attention infers batch_size from shape[0]
if batch_size > 1:
attn_in = ttnn.reshape(attn_in, [batch_size, 1, attn_in.shape[-2] // batch_size, -1])
attn_out = self.attention.forward(
attn_in,
current_pos,
rot_mats,
user_id,
mode,
page_table=page_table,
chunk_page_table=chunk_page_table,
chunk_start_idx=chunk_start_idx,
kv_cache=kv_cache,
)
# To match the batch-related reshape inside the attention module
# Use the batch_size parameter instead of inferring from shape[-3]
# because for [32, 1, S, H] tensors, shape[-3] is 1, not 32
# This reshape is only applicable in prefill mode with batched prefill
if mode == Mode.PREFILL and batch_size > 1:
residual = ttnn.reshape(residual, [1, 1, residual.shape[-2] * residual.shape[-3] * residual.shape[0], -1])
# TODO: create correct memory config in RopeSetup (issue is in ttnn.add op because of different shape in memory config for residual and rot_mats)
attn_out = ttnn.to_memory_config(attn_out, skip_mem_cfg)
if self.pre_ff_norm is None:
hidden_states = ttnn.add(
residual, attn_out, memory_config=skip_mem_cfg, dtype=ttnn.bfloat16 if TG else None
)
residual = hidden_states
if mode == "prefill":
x.deallocate(True)
else:
hidden_states = attn_out
ff_norm_config = self.args.get_norm_config("ff", mode, self.prefetcher)
hidden_states = self.ff_norm(hidden_states, mode, norm_config=ff_norm_config)
if self.pre_ff_norm is not None:
# Mesh partition ff_norm output to match residual sharding, skip if using distributed norm, because output is already sharded
if self.num_devices > 1 and not self.args.is_distributed_norm(mode):
hidden_states = ttnn.mesh_partition(
hidden_states,
memory_config=hidden_states.memory_config(),
dim=3,
cluster_axis=1,
)
hidden_states = ttnn.add(
residual, hidden_states, memory_config=skip_mem_cfg, dtype=ttnn.bfloat16 if TG else None
)
residual = hidden_states
pre_ff_norm_config = self.args.get_norm_config("ff", mode, self.prefetcher)
hidden_states = self.pre_ff_norm(hidden_states, mode, norm_config=pre_ff_norm_config)
ttnn.deallocate(attn_out)
if TG and mode == "decode":
hidden_states = ttnn.to_memory_config(hidden_states, memory_config=self.args.get_mlp_act_mem_config(mode))
# MLP takes replicated inputs and produces fractured outputs
hidden_states = self.feed_forward.forward(hidden_states, mode)
activation_dtype = self.args.decoders_optimizations.get_tensor_dtype(
decoder_id=self.layer_num, tensor=TensorGroup.ACTIVATION
)
if self.post_ff_norm is not None:
post_ff_norm_config = self.args.get_norm_config("ff", mode, self.prefetcher)
hidden_states = self.post_ff_norm(hidden_states, mode, norm_config=post_ff_norm_config) # Gathered
if self.num_devices > 1 and not self.args.is_distributed_norm(mode):
hidden_states = ttnn.mesh_partition(
hidden_states,
memory_config=hidden_states.memory_config(),
dim=3,
cluster_axis=1,
)
out = ttnn.add(
residual,
hidden_states,
memory_config=skip_mem_cfg,
dtype=self.args.ccl_dtype
if TG and not self.args.is_distributed_norm(mode)
else activation_dtype or ttnn.bfloat16,
)
return out # fractured across devices
|