| |
|
|
| |
|
|
| import json |
| import math |
| from dataclasses import asdict, dataclass, replace |
| from pathlib import Path |
|
|
| import torch |
|
|
| import ttnn |
| from models.common.lightweightmodule import LightweightModule |
|
|
| TILE = 32 |
| SEQ_ALIGN = 128 |
| SEQ_ALIGN_LARGE = 256 |
| LARGE_PATCHES = 2048 |
| CU_WINDOW_ENTRIES = TILE * TILE |
| MASK_NEG = -1e30 |
| VISION_PREFIX = "vision_tower." |
| EMBEDDER_PREFIX = "embed_vision." |
|
|
|
|
| @dataclass(frozen=True) |
| class VisionConfig: |
| hidden_size: int |
| intermediate_size: int |
| num_hidden_layers: int |
| num_attention_heads: int |
| head_dim: int |
| patch_size: int |
| pooling_kernel_size: int |
| position_embedding_size: int |
| rope_theta: float |
| rms_norm_eps: float |
| text_hidden_size: int |
|
|
| @classmethod |
| def from_hf(cls, hf_config): |
| if isinstance(hf_config, VisionConfig): |
| return hf_config |
| if isinstance(hf_config, (str, Path)): |
| hf_config = json.loads(Path(hf_config).read_text()) |
| if not isinstance(hf_config, dict): |
| hf_config = hf_config.to_dict() |
| v = hf_config["vision_config"] |
| if v["rope_parameters"]["rope_type"] != "axial": |
| raise ValueError(f"unsupported vision rope {v['rope_parameters']}") |
| if v.get("use_clipped_linears") or v.get("standardize"): |
| raise ValueError("clipped linears and standardize are not part of this checkpoint") |
| if v["num_key_value_heads"] != v["num_attention_heads"]: |
| raise ValueError("vision attention is MHA") |
| return cls( |
| hidden_size=v["hidden_size"], |
| intermediate_size=v["intermediate_size"], |
| num_hidden_layers=v["num_hidden_layers"], |
| num_attention_heads=v["num_attention_heads"], |
| head_dim=v["head_dim"], |
| patch_size=v["patch_size"], |
| pooling_kernel_size=v["pooling_kernel_size"], |
| position_embedding_size=v["position_embedding_size"], |
| rope_theta=float(v["rope_parameters"]["rope_theta"]), |
| rms_norm_eps=v["rms_norm_eps"], |
| text_hidden_size=hf_config["text_config"]["hidden_size"], |
| ) |
|
|
|
|
| WEIGHT_GROUPS = ("qkv", "o", "gate", "up", "down") |
| DTYPES = {"bf16": ttnn.bfloat16, "bfp8": ttnn.bfloat8_b, "bfp4": ttnn.bfloat4_b, "fp32": ttnn.float32} |
| FIDELITIES = {"lofi": ttnn.MathFidelity.LoFi, "hifi2": ttnn.MathFidelity.HiFi2, "hifi4": ttnn.MathFidelity.HiFi4} |
|
|
|
|
| @dataclass(frozen=True) |
| class VisionPolicy: |
| weight_dtype: tuple = (("qkv", "bf16"), ("o", "bf16"), ("gate", "bf16"), ("up", "bf16"), ("down", "bf16")) |
| weight_fidelity: tuple = (("qkv", "hifi2"), ("o", "hifi2"), ("gate", "hifi2"), ("up", "hifi2"), ("down", "hifi2")) |
| residual: str = "fp32" |
| sdpa_fidelity: str = "hifi4" |
| sdpa_fp32_acc: bool = True |
| sdpa_exp_approx: bool = True |
| single_image_mask: bool = True |
| mask_patches: tuple = (640, 1280, 5120) |
| mask_dtype: str = "bfp4" |
| windowed_chunks: tuple = (128, 256) |
| windowed_chunks_short: tuple = (64, 128) |
| short_patches: int = 2048 |
| mask_chunks: tuple = (128, 256) |
| minimal_matmul: bool = True |
| gelu_in_matmul: bool = False |
| packed_gate_up: bool = True |
| l1_max_rows: int = 1280 |
| l1_wide_max_rows: int = 640 |
| attention_impl: str = "sdpa" |
|
|
| def dtype(self, group): |
| return DTYPES[dict(self.weight_dtype)[group]] |
|
|
| def fidelity(self, group): |
| return FIDELITIES[dict(self.weight_fidelity)[group]] |
|
|
| def with_weights(self, dtype=None, fidelity=None, groups=WEIGHT_GROUPS): |
| wd, wf = dict(self.weight_dtype), dict(self.weight_fidelity) |
| for g in groups: |
| if dtype is not None: |
| wd[g] = dtype |
| if fidelity is not None: |
| wf[g] = fidelity |
| return replace( |
| self, |
| weight_dtype=tuple((g, wd[g]) for g in WEIGHT_GROUPS), |
| weight_fidelity=tuple((g, wf[g]) for g in WEIGHT_GROUPS), |
| ) |
|
|
| def describe(self): |
| d = asdict(self) |
| d["weight_dtype"] = dict(self.weight_dtype) |
| d["weight_fidelity"] = dict(self.weight_fidelity) |
| return d |
|
|
|
|
| def head_channel_order(head_dim): |
| q = head_dim // 4 |
| return torch.cat( |
| [torch.arange(0, q), torch.arange(2 * q, 3 * q), torch.arange(q, 2 * q), torch.arange(3 * q, 4 * q)] |
| ) |
|
|
|
|
| def axial_inv_freq(config): |
| spatial = config.head_dim // 2 |
| return 1.0 / (config.rope_theta ** (torch.arange(0, spatial, 2, dtype=torch.float) / spatial)) |
|
|
|
|
| def rope_tables(config): |
| inv = axial_inv_freq(config) |
| n = config.position_embedding_size |
| ang = torch.arange(n, dtype=torch.float32)[:, None] * inv[None, :] |
| q = inv.numel() |
| zero = torch.zeros(n, q) |
| tables = [] |
| for axis in range(2): |
| parts = [] |
| for fn in (torch.cos, torch.sin): |
| v = fn(ang) |
| parts.append(torch.cat([v, zero, v, zero] if axis == 0 else [zero, v, zero, v], dim=-1)) |
| t = torch.cat(parts, dim=-1) |
| tables.append(torch.cat([t, torch.zeros(1, t.shape[-1])], dim=0)) |
| return tables |
|
|
|
|
| def padded_patches(num_patches): |
| align = SEQ_ALIGN_LARGE if num_patches >= LARGE_PATCHES else SEQ_ALIGN |
| return align * math.ceil(num_patches / align) |
|
|
|
|
| def padded_slots(num_slots): |
| return TILE * math.ceil(num_slots / TILE) |
|
|
|
|
| @dataclass |
| class VisionHostInputs: |
| pixels: torch.Tensor |
| x_idx: torch.Tensor |
| y_idx: torch.Tensor |
| cu_window: torch.Tensor |
| pool: torch.Tensor |
| gather_idx: torch.Tensor |
| key_bias: torch.Tensor |
| num_images: int |
| num_patches: int |
| patches_padded: int |
| slots: int |
| slots_padded: int |
| valid_patches: tuple |
| tokens_per_image: tuple |
| tokens_padded: int |
|
|
|
|
| @dataclass |
| class VisionDeviceInputs: |
| pixels: ttnn.Tensor |
| x_idx: ttnn.Tensor |
| y_idx: ttnn.Tensor |
| cu_window: ttnn.Tensor |
| pool: ttnn.Tensor |
| gather_idx: ttnn.Tensor |
| key_bias: ttnn.Tensor |
| host: VisionHostInputs |
|
|
|
|
| def prepare_host_inputs(config, pixel_values, position_ids, tokens_padded=None): |
| n, p, d = pixel_values.shape |
| k = config.pooling_kernel_size |
| slots = p // (k * k) |
| if slots * k * k != p: |
| raise ValueError(f"patch count {p} is not a multiple of {k * k}") |
| pp = padded_patches(p) |
| sp = padded_slots(slots) |
| pad_row = config.position_embedding_size |
| pixels = torch.zeros(1, 1, n * pp, d) |
| x_idx = torch.full((n, pp), pad_row, dtype=torch.int32) |
| y_idx = torch.full((n, pp), pad_row, dtype=torch.int32) |
| pool = torch.zeros(n, sp, pp) |
| key_bias = torch.full((n, 1, TILE, pp), MASK_NEG) |
| scale = config.hidden_size**0.5 / (k * k) |
| cu = [0] |
| valid_counts, token_counts, gather = [], [], [] |
| for i in range(n): |
| pos = position_ids[i].long() |
| valid = (pos != -1).all(dim=-1) |
| nv = int(valid.sum()) |
| if not bool(valid[:nv].all()): |
| raise ValueError("padding patches must follow valid patches") |
| pixels[0, 0, i * pp : i * pp + p] = pixel_values[i].float() |
| key_bias[i, 0, :, :nv] = 0.0 |
| x_idx[i, :nv] = pos[:nv, 0].int() |
| y_idx[i, :nv] = pos[:nv, 1].int() |
| xs, ys = pos[:nv, 0], pos[:nv, 1] |
| max_x = int(xs.max()) + 1 |
| slot = torch.div(xs, k, rounding_mode="floor") + (max_x // k) * torch.div(ys, k, rounding_mode="floor") |
| pool[i, slot, torch.arange(nv)] = scale |
| used = torch.unique(slot).tolist() |
| gather.extend(i * sp + s for s in used) |
| valid_counts.append(nv) |
| token_counts.append(len(used)) |
| cu.extend([i * pp + nv, (i + 1) * pp]) |
| total = len(gather) |
| tp = tokens_padded if tokens_padded is not None else padded_slots(n * slots) |
| if total > tp: |
| raise ValueError(f"{total} soft tokens exceed tokens_padded {tp}") |
| gather_idx = torch.full((1, tp), n * sp, dtype=torch.int32) |
| gather_idx[0, :total] = torch.tensor(gather, dtype=torch.int32) |
| if len(cu) > CU_WINDOW_ENTRIES: |
| raise ValueError(f"{n} images need {len(cu)} window boundaries, limit {CU_WINDOW_ENTRIES}") |
| cu = cu + [cu[-1]] * (CU_WINDOW_ENTRIES - len(cu)) |
| return VisionHostInputs( |
| pixels=pixels, |
| x_idx=x_idx.reshape(1, n * pp), |
| y_idx=y_idx.reshape(1, n * pp), |
| cu_window=torch.tensor(cu, dtype=torch.int32), |
| pool=pool, |
| gather_idx=gather_idx, |
| key_bias=key_bias, |
| num_images=n, |
| num_patches=p, |
| patches_padded=pp, |
| slots=slots, |
| slots_padded=sp, |
| valid_patches=tuple(valid_counts), |
| tokens_per_image=tuple(token_counts), |
| tokens_padded=tp, |
| ) |
|
|
|
|
| INPUT_SPEC = { |
| "pixels": (ttnn.bfloat16, ttnn.TILE_LAYOUT), |
| "x_idx": (ttnn.uint32, ttnn.ROW_MAJOR_LAYOUT), |
| "y_idx": (ttnn.uint32, ttnn.ROW_MAJOR_LAYOUT), |
| "cu_window": (ttnn.uint32, ttnn.ROW_MAJOR_LAYOUT), |
| "pool": (ttnn.bfloat16, ttnn.TILE_LAYOUT), |
| "gather_idx": (ttnn.uint32, ttnn.ROW_MAJOR_LAYOUT), |
| "key_bias": (ttnn.bfloat16, ttnn.TILE_LAYOUT), |
| } |
|
|
|
|
| def host_inputs_to_device(host, mesh_device): |
| tensors = { |
| name: ttnn.from_torch( |
| getattr(host, name).contiguous(), |
| dtype=dt, |
| layout=layout, |
| device=mesh_device, |
| memory_config=ttnn.DRAM_MEMORY_CONFIG, |
| ) |
| for name, (dt, layout) in INPUT_SPEC.items() |
| } |
| return VisionDeviceInputs(host=host, **tensors) |
|
|
|
|
| def copy_host_inputs(host, device_inputs): |
| old = device_inputs.host |
| for f in ("num_images", "patches_padded", "slots_padded", "tokens_padded"): |
| if getattr(old, f) != getattr(host, f): |
| raise ValueError(f"{f} differs from the captured bucket: {getattr(host, f)} vs {getattr(old, f)}") |
| for name, (dt, layout) in INPUT_SPEC.items(): |
| src = ttnn.from_torch(getattr(host, name).contiguous(), dtype=dt, layout=layout) |
| ttnn.copy_host_to_device_tensor(src, getattr(device_inputs, name)) |
| device_inputs.host = host |
|
|
|
|
| def _compute_config(mesh_device, fidelity=ttnn.MathFidelity.HiFi4, fp32_acc=True): |
| return ttnn.init_device_compute_kernel_config( |
| mesh_device.arch(), |
| math_fidelity=fidelity, |
| math_approx_mode=False, |
| fp32_dest_acc_en=fp32_acc, |
| packer_l1_acc=True, |
| ) |
|
|
|
|
| def _put_weight(t, mesh_device, dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT): |
| return ttnn.from_torch( |
| t.contiguous(), dtype=dtype, layout=layout, device=mesh_device, memory_config=ttnn.DRAM_MEMORY_CONFIG |
| ) |
|
|
|
|
| def _norm_weight(w, mesh_device): |
| return _put_weight(w.float().reshape(1, 1, w.numel() // TILE, TILE), mesh_device, layout=ttnn.ROW_MAJOR_LAYOUT) |
|
|
|
|
| MATMUL_BLOCKS = { |
| "qkv": ((8, 4, 8, 2, 2), (16, 4, 8, 1, 4)), |
| "o": ((8, 4, 8, 1, 4), (16, 8, 4, 1, 4)), |
| "gate": ((8, 8, 12, 2, 2), (16, 8, 4, 1, 4)), |
| "up": ((8, 8, 12, 2, 2), (16, 8, 4, 1, 4)), |
| "down": ((8, 4, 8, 1, 4), (8, 12, 8, 1, 4)), |
| "gate_up": ((4, 4, 24, 2, 2), (16, 8, 4, 1, 4)), |
| } |
| LARGE_M = 8192 |
|
|
|
|
| def matmul_config(group, m, grid): |
| mb, kb, nb, sh, sw = MATMUL_BLOCKS[group][1 if m >= LARGE_M else 0] |
| return ttnn.MinimalMatmulConfig( |
| M_block_size=mb, |
| K_block_size=kb, |
| N_block_size=nb, |
| subblock_h=sh, |
| subblock_w=sw, |
| compute_with_storage_grid_size=grid, |
| ) |
|
|
|
|
| @dataclass |
| class AttentionPlan: |
| mode: str |
| l1_activations: bool = False |
| cu_window: ttnn.Tensor = None |
| mask: ttnn.Tensor = None |
| program_config: object = None |
| key_bias: ttnn.Tensor = None |
| num_images: int = 1 |
| patches_padded: int = 0 |
|
|
|
|
| class VisionEncoderLayer(LightweightModule): |
| def __init__(self, config, layer_idx, mesh_device, weights, policy): |
| super().__init__() |
| self.policy = policy |
| self.residual_dtype = DTYPES[policy.residual] |
| self.config = config |
| self.layer_idx = layer_idx |
| self.mesh_device = mesh_device |
| self.w = weights |
| self.eps = config.rms_norm_eps |
| self.num_heads = config.num_attention_heads |
| self.norm_cfg = _compute_config(mesh_device) |
| self.mm_cfg = {g: _compute_config(mesh_device, fidelity=policy.fidelity(g)) for g in WEIGHT_GROUPS} |
| self.sdpa_compute_cfg = _compute_config( |
| mesh_device, fidelity=FIDELITIES[policy.sdpa_fidelity], fp32_acc=policy.sdpa_fp32_acc |
| ) |
| self.grid = mesh_device.compute_with_storage_grid_size() |
| self._l1_activations = False |
|
|
| @classmethod |
| def from_state_dict(cls, state_dict, *, hf_config, layer_idx, mesh_device, prefix=None, policy=None): |
| config = VisionConfig.from_hf(hf_config) |
| policy = policy or VisionPolicy() |
| if prefix is None: |
| prefix = f"{VISION_PREFIX}encoder.layers.{layer_idx}." |
| if prefix + "input_layernorm.weight" not in state_dict: |
| prefix = "" |
| order = head_channel_order(config.head_dim) |
| h = config.num_attention_heads |
|
|
| def g(name): |
| return state_dict[prefix + name].float() |
|
|
| def per_head_rows(w): |
| return w.reshape(h, config.head_dim, -1)[:, order].reshape(h * config.head_dim, -1) |
|
|
| def linear_w(w, group): |
| return _put_weight(w.T, mesh_device, dtype=policy.dtype(group)) |
|
|
| wq = per_head_rows(g("self_attn.q_proj.linear.weight")) |
| wk = per_head_rows(g("self_attn.k_proj.linear.weight")) |
| wv = g("self_attn.v_proj.linear.weight") |
| weights = { |
| "qkv": linear_w(torch.cat([wq, wk, wv], dim=0), "qkv"), |
| "o": linear_w(g("self_attn.o_proj.linear.weight"), "o"), |
| "q_norm": _norm_weight(g("self_attn.q_norm.weight")[order], mesh_device), |
| "k_norm": _norm_weight(g("self_attn.k_norm.weight")[order], mesh_device), |
| "input_norm": _norm_weight(g("input_layernorm.weight"), mesh_device), |
| "post_attn_norm": _norm_weight(g("post_attention_layernorm.weight"), mesh_device), |
| "pre_ff_norm": _norm_weight(g("pre_feedforward_layernorm.weight"), mesh_device), |
| "post_ff_norm": _norm_weight(g("post_feedforward_layernorm.weight"), mesh_device), |
| "gate": linear_w(g("mlp.gate_proj.linear.weight"), "gate"), |
| "up": linear_w(g("mlp.up_proj.linear.weight"), "up"), |
| "down": linear_w(g("mlp.down_proj.linear.weight"), "down"), |
| } |
| if policy.packed_gate_up: |
| weights["gate_up"] = _put_weight( |
| torch.cat([g("mlp.gate_proj.linear.weight"), g("mlp.up_proj.linear.weight")], dim=0).T, |
| mesh_device, |
| dtype=policy.dtype("gate"), |
| ) |
| return cls(config, layer_idx, mesh_device, weights, policy) |
|
|
| def _norm(self, x, weight=None, dtype=None): |
| mem = self._mem(x) |
| y = ttnn.rms_norm(x, epsilon=self.eps, weight=weight, compute_kernel_config=self.norm_cfg, memory_config=mem) |
| if dtype is not None and y.dtype != dtype: |
| z = ttnn.typecast(y, dtype, memory_config=mem) |
| ttnn.deallocate(y) |
| return z |
| return y |
|
|
| def _mem(self, x, wide=False): |
| limit = self.policy.l1_wide_max_rows if wide else self.policy.l1_max_rows |
| use_l1 = self._l1_activations and x.padded_shape[-2] <= limit |
| return ttnn.L1_MEMORY_CONFIG if use_l1 else ttnn.DRAM_MEMORY_CONFIG |
|
|
| def _linear(self, x, group, activation=None, wide=False): |
| if self.policy.minimal_matmul: |
| return ttnn.experimental.minimal_matmul( |
| input_tensor=x, |
| weight_tensor=self.w[group], |
| fused_activation=activation, |
| config=matmul_config(group, x.padded_shape[-2], self.grid), |
| compute_kernel_config=self.mm_cfg[group], |
| dtype=ttnn.bfloat16, |
| memory_config=self._mem(x, wide), |
| ) |
| return ttnn.linear( |
| x, |
| self.w[group], |
| activation=activation, |
| dtype=ttnn.bfloat16, |
| compute_kernel_config=self.mm_cfg[group], |
| memory_config=self._mem(x, wide), |
| ) |
|
|
| def attention(self, x, cos, sin, plan): |
| qkv = self._linear(x, "qkv") |
| q, k, v = ttnn.experimental.nlp_create_qkv_heads( |
| qkv, |
| num_heads=self.num_heads, |
| num_kv_heads=self.num_heads, |
| transpose_k_heads=False, |
| memory_config=self._mem(x), |
| ) |
| ttnn.deallocate(qkv) |
| qn = self._norm(q, self.w["q_norm"]) |
| kn = self._norm(k, self.w["k_norm"]) |
| vn = self._norm(v) |
| ttnn.deallocate(q) |
| ttnn.deallocate(k) |
| ttnn.deallocate(v) |
| qr = ttnn.experimental.rotary_embedding_hf(qn, cos, sin, is_decode_mode=False, memory_config=self._mem(x)) |
| kr = ttnn.experimental.rotary_embedding_hf(kn, cos, sin, is_decode_mode=False, memory_config=self._mem(x)) |
| ttnn.deallocate(qn) |
| ttnn.deallocate(kn) |
| if self.policy.attention_impl == "fp32_reference": |
| ctx = self._fp32_attention(qr, kr, vn, plan) |
| ttnn.deallocate(qr) |
| ttnn.deallocate(kr) |
| ttnn.deallocate(vn) |
| out = self._linear(ctx, "o") |
| ttnn.deallocate(ctx) |
| return out |
| extra = {"cu_window_seqlens": plan.cu_window} if plan.mode == "windowed" else {"attn_mask": plan.mask} |
| ctx = ttnn.transformer.scaled_dot_product_attention( |
| qr, |
| kr, |
| vn, |
| is_causal=False, |
| scale=1.0, |
| program_config=plan.program_config, |
| compute_kernel_config=self.sdpa_compute_cfg, |
| output_concat_heads=True, |
| memory_config=self._mem(x), |
| **extra, |
| ) |
| ttnn.deallocate(qr) |
| ttnn.deallocate(kr) |
| ttnn.deallocate(vn) |
| out = self._linear(ctx, "o") |
| ttnn.deallocate(ctx) |
| return out |
|
|
| def _fp32_attention(self, q, k, v, plan): |
| n, pp, d = plan.num_images, plan.patches_padded, self.config.head_dim |
| shape = (self.num_heads, n, pp, d) |
| q4, k4, v4 = [ttnn.reshape(ttnn.typecast(t, ttnn.float32), shape) for t in (q, k, v)] |
| scores = ttnn.matmul(q4, k4, transpose_b=True, dtype=ttnn.float32, compute_kernel_config=self.norm_cfg) |
| bias = ttnn.reshape(ttnn.slice(plan.key_bias, (0, 0, 0, 0), (n, 1, 1, pp)), (1, n, 1, pp)) |
| scores = ttnn.add(scores, ttnn.typecast(bias, ttnn.float32)) |
| probs = ttnn.softmax(scores, dim=-1, numeric_stable=True, compute_kernel_config=self.norm_cfg) |
| ctx = ttnn.matmul(probs, v4, dtype=ttnn.float32, compute_kernel_config=self.norm_cfg) |
| ctx = ttnn.reshape(ctx, (1, self.num_heads, n * pp, d)) |
| merged = ttnn.experimental.nlp_concat_heads( |
| ttnn.typecast(ctx, ttnn.bfloat16), memory_config=ttnn.DRAM_MEMORY_CONFIG |
| ) |
| for t in (q4, k4, v4, scores, probs, ctx): |
| ttnn.deallocate(t) |
| return merged |
|
|
| def mlp(self, x): |
| gelu = ttnn.UnaryWithParam(ttnn.UnaryOpType.GELU_TANH) |
| if self.policy.packed_gate_up: |
| gate, up = ttnn.experimental.minimal_matmul_split( |
| x, |
| self.w["gate_up"], |
| chunks=2, |
| dim=-1, |
| config=matmul_config("gate_up", x.padded_shape[-2], self.grid), |
| compute_kernel_config=self.mm_cfg["gate"], |
| dtype=ttnn.bfloat16, |
| memory_config=self._mem(x, wide=True), |
| ) |
| h = ttnn.multiply(gate, up, input_tensor_a_activations=[gelu], memory_config=self._mem(x, wide=True)) |
| elif self.policy.gelu_in_matmul: |
| gate = self._linear(x, "gate", activation=gelu, wide=True) |
| up = self._linear(x, "up", wide=True) |
| h = ttnn.multiply(gate, up, memory_config=self._mem(x, wide=True)) |
| else: |
| gate = self._linear(x, "gate", wide=True) |
| up = self._linear(x, "up", wide=True) |
| h = ttnn.multiply(gate, up, input_tensor_a_activations=[gelu], memory_config=self._mem(x, wide=True)) |
| ttnn.deallocate(gate) |
| ttnn.deallocate(up) |
| out = self._linear(h, "down") |
| ttnn.deallocate(h) |
| return out |
|
|
| def forward(self, hidden, cos, sin, plan): |
| self._l1_activations = plan.l1_activations |
| rd = self.residual_dtype |
| h = self._norm(hidden, self.w["input_norm"], dtype=ttnn.bfloat16) |
| a = self.attention(h, cos, sin, plan) |
| ttnn.deallocate(h) |
| h = self._norm(a, self.w["post_attn_norm"]) |
| ttnn.deallocate(a) |
| x = ttnn.add(hidden, h, dtype=rd, memory_config=self._mem(hidden)) |
| ttnn.deallocate(h) |
| h = self._norm(x, self.w["pre_ff_norm"], dtype=ttnn.bfloat16) |
| m = self.mlp(h) |
| ttnn.deallocate(h) |
| h = self._norm(m, self.w["post_ff_norm"]) |
| ttnn.deallocate(m) |
| out = ttnn.add(x, h, dtype=rd, memory_config=self._mem(x)) |
| ttnn.deallocate(h) |
| ttnn.deallocate(x) |
| return out |
|
|
|
|
| class VisionEncoder(LightweightModule): |
| def __init__(self, config, mesh_device, weights, layers, policy): |
| super().__init__() |
| self.policy = policy |
| self.residual_dtype = DTYPES[policy.residual] |
| self.config = config |
| self.mesh_device = mesh_device |
| self.w = weights |
| self.layers = layers |
| self.eps = config.rms_norm_eps |
| self.compute_cfg = _compute_config(mesh_device) |
| self.grid = mesh_device.compute_with_storage_grid_size() |
| self.zero_rows = ttnn.zeros( |
| (1, 1, TILE, config.text_hidden_size), |
| dtype=ttnn.bfloat16, |
| layout=ttnn.TILE_LAYOUT, |
| device=mesh_device, |
| memory_config=ttnn.DRAM_MEMORY_CONFIG, |
| ) |
|
|
| @classmethod |
| def from_state_dict(cls, state_dict, *, hf_config, mesh_device, policy=None, num_layers=None): |
| config = VisionConfig.from_hf(hf_config) |
| policy = policy or VisionPolicy() |
|
|
| def g(name): |
| return state_dict[name].float() |
|
|
| w_in = g(VISION_PREFIX + "patch_embedder.input_proj.weight") |
| table = g(VISION_PREFIX + "patch_embedder.position_embedding_table") |
| zero = torch.zeros(1, config.hidden_size) |
| rope_x, rope_y = rope_tables(config) |
| weights = { |
| "w_patch": _put_weight((2.0 * w_in).T, mesh_device), |
| "b_patch": _put_weight(-w_in.sum(dim=1).reshape(1, -1), mesh_device), |
| "pos_x": _put_weight(torch.cat([table[0], zero]), mesh_device, layout=ttnn.ROW_MAJOR_LAYOUT), |
| "pos_y": _put_weight(torch.cat([table[1], zero]), mesh_device, layout=ttnn.ROW_MAJOR_LAYOUT), |
| "rope_x": _put_weight(rope_x, mesh_device, layout=ttnn.ROW_MAJOR_LAYOUT), |
| "rope_y": _put_weight(rope_y, mesh_device, layout=ttnn.ROW_MAJOR_LAYOUT), |
| "w_embed": _put_weight(g(EMBEDDER_PREFIX + "embedding_projection.weight").T, mesh_device), |
| } |
| n = config.num_hidden_layers if num_layers is None else num_layers |
| layers = [ |
| VisionEncoderLayer.from_state_dict( |
| state_dict, |
| hf_config=config, |
| layer_idx=i, |
| mesh_device=mesh_device, |
| prefix=f"{VISION_PREFIX}encoder.layers.{i}.", |
| policy=policy, |
| ) |
| for i in range(n) |
| ] |
| return cls(config, mesh_device, weights, layers, policy) |
|
|
| def attention_mode(self, num_images, patches_padded): |
| p = self.policy |
| return "mask" if p.single_image_mask and num_images == 1 and patches_padded in p.mask_patches else "windowed" |
|
|
| def _sdpa_program_config(self, chunks, seq_len): |
| q, k = chunks |
| while seq_len % q: |
| q //= 2 |
| while seq_len % k: |
| k //= 2 |
| return ttnn.SDPAProgramConfig( |
| compute_with_storage_grid_size=self.grid, |
| q_chunk_size=q, |
| k_chunk_size=k, |
| exp_approx_mode=self.policy.sdpa_exp_approx, |
| ) |
|
|
| def attention_plan(self, inputs): |
| host = inputs.host |
| if self.attention_mode(host.num_images, host.patches_padded) == "mask": |
| pp = host.patches_padded |
| full = ttnn.repeat(inputs.key_bias, (1, 1, pp // TILE, 1)) |
| mask = ttnn.typecast(full, DTYPES[self.policy.mask_dtype]) |
| ttnn.deallocate(full) |
| return AttentionPlan( |
| "mask", |
| l1_activations=True, |
| mask=mask, |
| program_config=self._sdpa_program_config(self.policy.mask_chunks, pp), |
| key_bias=inputs.key_bias, |
| num_images=host.num_images, |
| patches_padded=pp, |
| ) |
| seq = host.num_images * host.patches_padded |
| short = host.patches_padded < self.policy.short_patches |
| chunks = self.policy.windowed_chunks_short if short else self.policy.windowed_chunks |
| return AttentionPlan( |
| "windowed", |
| cu_window=inputs.cu_window, |
| program_config=self._sdpa_program_config(chunks, seq), |
| key_bias=inputs.key_bias, |
| num_images=host.num_images, |
| patches_padded=host.patches_padded, |
| ) |
|
|
| def _linear(self, x, w, bias=None): |
| return ttnn.linear( |
| x, w, bias=bias, compute_kernel_config=self.compute_cfg, memory_config=ttnn.DRAM_MEMORY_CONFIG |
| ) |
|
|
| def _lookup(self, idx, table): |
| out = ttnn.embedding(idx, table, layout=ttnn.TILE_LAYOUT, memory_config=ttnn.DRAM_MEMORY_CONFIG) |
| return ttnn.reshape(out, (1, 1, out.shape[-2], out.shape[-1])) |
|
|
| def rope(self, inputs): |
| cs = ttnn.add(self._lookup(inputs.x_idx, self.w["rope_x"]), self._lookup(inputs.y_idx, self.w["rope_y"])) |
| d = self.config.head_dim |
| cos = ttnn.slice(cs, (0, 0, 0, 0), (1, 1, cs.shape[-2], d)) |
| sin = ttnn.slice(cs, (0, 0, 0, d), (1, 1, cs.shape[-2], 2 * d)) |
| ttnn.deallocate(cs) |
| return cos, sin |
|
|
| def patch_embed(self, inputs): |
| if self.policy.minimal_matmul: |
| h = ttnn.experimental.minimal_matmul( |
| input_tensor=inputs.pixels, |
| weight_tensor=self.w["w_patch"], |
| bias_tensor=self.w["b_patch"], |
| config=matmul_config("o", inputs.pixels.padded_shape[-2], self.grid), |
| compute_kernel_config=self.compute_cfg, |
| dtype=ttnn.bfloat16, |
| memory_config=ttnn.DRAM_MEMORY_CONFIG, |
| ) |
| else: |
| h = self._linear(inputs.pixels, self.w["w_patch"], bias=self.w["b_patch"]) |
| px = self._lookup(inputs.x_idx, self.w["pos_x"]) |
| py = self._lookup(inputs.y_idx, self.w["pos_y"]) |
| pos = ttnn.add(px, py) |
| ttnn.deallocate(px) |
| ttnn.deallocate(py) |
| out = ttnn.add(h, pos, dtype=self.residual_dtype) |
| ttnn.deallocate(h) |
| ttnn.deallocate(pos) |
| return out |
|
|
| def pool_and_embed(self, hidden, inputs): |
| host = inputs.host |
| h = ttnn.reshape(hidden, (host.num_images, host.patches_padded, self.config.hidden_size)) |
| pooled = ttnn.matmul(inputs.pool, h, compute_kernel_config=self.compute_cfg, dtype=ttnn.bfloat16) |
| flat = ttnn.reshape(pooled, (1, 1, host.num_images * host.slots_padded, self.config.hidden_size)) |
| normed = ttnn.rms_norm(flat, epsilon=self.eps, compute_kernel_config=self.compute_cfg) |
| proj = self._linear(normed, self.w["w_embed"]) |
| ttnn.deallocate(normed) |
| table = ttnn.concat([proj, self.zero_rows], dim=2) |
| ttnn.deallocate(proj) |
| table_rm = ttnn.to_layout(ttnn.reshape(table, (table.shape[-2], table.shape[-1])), ttnn.ROW_MAJOR_LAYOUT) |
| ttnn.deallocate(table) |
| tokens = self._lookup(inputs.gather_idx, table_rm) |
| ttnn.deallocate(table_rm) |
| return tokens, flat |
|
|
| def forward(self, inputs, return_intermediates=False): |
| cos, sin = self.rope(inputs) |
| plan = self.attention_plan(inputs) |
| h = self.patch_embed(inputs) |
| inter = {"patch_embed": ttnn.clone(h, memory_config=ttnn.DRAM_MEMORY_CONFIG)} if return_intermediates else None |
| for i, layer in enumerate(self.layers): |
| nxt = layer(h, cos, sin, plan) |
| if return_intermediates: |
| inter[f"layer{i:02d}"] = ttnn.clone(nxt, memory_config=ttnn.DRAM_MEMORY_CONFIG) |
| ttnn.deallocate(h) |
| h = nxt |
| ttnn.deallocate(cos) |
| ttnn.deallocate(sin) |
| if plan.mask is not None: |
| ttnn.deallocate(plan.mask) |
| tokens, pooled = self.pool_and_embed(h, inputs) |
| if not return_intermediates: |
| ttnn.deallocate(h) |
| ttnn.deallocate(pooled) |
| return {"tokens": tokens} |
| inter["pooled"] = pooled |
| inter["tokens"] = tokens |
| return inter |
|
|
