# SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc. # SPDX-License-Identifier: Apache-2.0 """Host-side weight pipeline of the GR00T N1.5 / N1.6 / N1.7 port: categorize -> transform -> *device plan*. Everything here is CPU torch + safetensors (through :class:`common.checkpoint.LazyCheckpoint`); no ``ttnn`` import. The module implements ``docs/research/weight-loader-spec.md`` (cited as *WL* below) and the layout facts of ``docs/research/host-layout-helpers-validation.md`` (*HL*): * **Key -> category rules** (WL §2): 15 / 14 / 15 first-match-wins regex rules per version with explicit ``skip:*`` categories (WL §3). :func:`categorize` partitions every key exactly once and raises on an unmatched key. * **Host transforms** (WL §4) as pure torch functions: ``nn.Linear`` ``[out,in] -> [in,out]`` orientation, SigLIP conv -> linear in ``(ph,pw,c)`` column order (N1.6's Linear needs no permutation, HL §3.1), N1.6 bilinear position table 16x16 -> 18x18, per-head zero padding (SigLIP 72 -> 96, DiT 48 -> 64) of Q/K/V columns + biases and of the output-projection rows, fused ``[Q | K | V]`` and ``[gate | up]`` projections, N1.6 pixel-unshuffle channel-order fold into ``mlp1`` (HL §3.3), N1.7 merger column permutation (HL §2.6), embodiment slot extraction with the exact tau-fold ``b2_step[t]`` and the ``b3_pos`` fold (WL §4.7), AdaLN / timestep precompute for the four buckets (WL §4.6; synthesis §0 fact 5: blocks ``(scale, shift)``, output head ``(shift, scale)``). * **Device plan** (WL §5): :func:`build_device_plan` returns ``{name: PlanTensor}`` whose *names are the contract* between this module and the ``tt/`` code (list below); :func:`enumerate_plan` gives the same entries without reading any weight (shapes, dtypes, layouts, source keys) so that byte accounting (WL §6) runs from headers alone. * **Host cache skeleton** (WL §8, Tier 0): :class:`HostPlanCache` saves / loads the plan as torch ``.pt`` files with a manifest and a ``.complete`` marker; the ttnn ``.tensorbin`` tier belongs to ``tt/``. Plan-tensor name contract (``{i}`` = layer / block index, ``{k}`` = denoising step 0..3, ``{d}`` = deepstack index) -------------------------------------------------------------------------------------------------------------------- Category ``vision`` (SigLIP N1.5 / N1.6: 27 layers, hidden 1152; Qwen3-VL N1.7: 24 blocks, hidden 1024):: vision.patch_embed.W [608,1152] (N1.5/N1.6: (ph,pw,c) columns, rows 588 -> 608) | [1536,1024] (N1.7) vision.patch_embed.b [1,1152] | [1,1024] vision.pos_embed [256,1152] (N1.5 raster) | [324,1152] (N1.6, host-resized 18x18 raster) -- absent N1.7 vision.post_ln.w / .b [1,1152] (N1.5/N1.6 only) vision.layer.{i}.ln1.w / .b [1,D] (N1.7: blocks.{i}.norm1) vision.layer.{i}.ln2.w / .b [1,D] (N1.7: blocks.{i}.norm2) vision.layer.{i}.Wqkv [1152,4608] (heads 72 -> 96, pad_heads) | [1152,3456] (no pad) | [1024,3072] (N1.7) vision.layer.{i}.bqkv [1,4608] | [1,3456] | [1,3072] vision.layer.{i}.Wo [1536,1152] (rows padded) | [1152,1152] | [1024,1024] (N1.7 attn.proj) vision.layer.{i}.bo [1,D] vision.layer.{i}.fc1.W / .b [1152,4304], [1,4304] | [1024,4096], [1,4096] vision.layer.{i}.fc2.W / .b [4304,1152], [1,1152] | [4096,1024], [1,1024] Category ``vision_merger`` (N1.7 only):: vision_merger.merger.norm.w / .b [1,1024] (pre-concat LN) vision_merger.merger.fc1.W / .b [4096,4096], [1,4096] (columns (mh,mw,c) unless n17_merger_order says otherwise) vision_merger.merger.fc2.W / .b [4096,2048], [1,2048] vision_merger.deepstack.{d}.norm.w / .b [1,4096] (post-shuffle LN; permuted with the merger order) vision_merger.deepstack.{d}.fc1.W / .b [4096,4096], [1,4096] vision_merger.deepstack.{d}.fc2.W / .b [4096,2048], [1,2048] Category ``vision_host`` (N1.7 only, never uploaded):: vision.pos_embed_table [2304,1024] fp32 -- consumed by qwen3vl_pos_embeds() per grid_thw Category ``projector``:: projector.fc1.W / .b [1152,2048], [1,2048] (N1.5 mlp1.0) projector.ln.w / .b [1,4608] (N1.6 mlp1.0, eps 1e-5; folded if n16_unshuffle_order == "(i,j,c)") projector.fc1.W / .b [4608,2048], [1,2048] (N1.6 mlp1.1, GELU-erf after; columns folded likewise) projector.fc2.W / .b [2048,2048], [1,2048] (N1.6 mlp1.3) Category ``llm`` (12 / 16 / 16 layers):: llm.layer.{i}.Wqkv [2048,4096] = [Q 2048 | K 1024 | V 1024] (rope_layout "hf": no row permutation) llm.layer.{i}.Wo [2048,2048] llm.layer.{i}.W_gate_up [2048,12288] = [gate 6144 | up 6144] ("plain") or 32-column tile pairs ("tile_interleaved") llm.layer.{i}.W_down [6144,2048] llm.layer.{i}.input_ln.w [1,2048] llm.layer.{i}.post_attn_ln.w [1,2048] llm.layer.{i}.q_norm.w / k_norm.w [1,128] llm.final_norm.w [1,2048] (N1.5 / N1.6 only; N1.7's backbone feature is pre-norm) Category ``llm_embed`` (only when ``TransformOptions.include_embed_table``):: llm.embed_tokens [151680|151936, 2048] ROW_MAJOR bf16 Category ``vlln``:: vlln.w / vlln.b [1,2048] Category ``vl_self_attention`` (N1.5 / N1.7, 4 blocks):: vlsa.block.{i}.Wqkv / .bqkv [2048,6144], [1,6144] vlsa.block.{i}.Wo / .bo [2048,2048], [1,2048] vlsa.block.{i}.norm1.w / .b [1,2048] vlsa.block.{i}.norm3.w / .b [1,2048] vlsa.block.{i}.ff1.W / .b [2048,8192], [1,8192] (GELU-tanh after) vlsa.block.{i}.ff2.W / .b [8192,2048], [1,2048] Category ``dit`` (16 / 32 / 32 blocks; even = cross-attention, odd = self-attention; 32 heads x 48 -> 64):: dit.block.{i}.Wqkv / .bqkv [1536,6144], [1,6144] (odd blocks; [1536,4608] without head padding) dit.block.{i}.Wq / .bq [1536,2048], [1,2048] (even blocks; [1536,1536] without padding) dit.block.{i}.Wkv / .bkv [2048,4096], [1,4096] (even blocks, in = vl_embs 2048; [2048,3072] without padding) dit.block.{i}.Wo / .bo [2048,1536], [1,1536] (rows padded; [1536,1536] without padding) dit.block.{i}.ff1.W / .b [1536,6144], [1,6144] (GELU-tanh after) dit.block.{i}.ff2.W / .b [6144,1536], [1,1536] dit.proj_out_2.W / .b [1536,1024], [1,1024] (always bf16) Category ``dit_precompute`` (host fp32 math, rounded once to bf16; WL §4.6):: dit.block.{i}.adaln.step{k}.one_plus_scale [1,1536] = 1 + scale, scale, shift = chunk(norm1.linear(SiLU(temb_k))) dit.block.{i}.adaln.step{k}.shift [1,1536] dit.out_head.adaln.step{k}.one_plus_scale [1,1536] shift, scale = chunk(proj_out_1(SiLU(temb_k))) (shift FIRST) dit.out_head.adaln.step{k}.shift [1,1536] Category ``encoders`` (one embodiment slot ``e``; ``W[e]`` is already ``[in,out]``, consumed as ``x @ W + b``):: enc.state.W1 / .b1 [64|128|160, 1024] (rows zero-padded from max_state_dim), [1,1024] enc.state.W2 / .b2 [1024,1536], [1,1536] enc.action.W1 / .b1 [32|128|160, 1536] (rows padded from max_action_dim), [1,1536] enc.action.W2a [1536,1536] first 1536 rows of W2 (the a_emb half) enc.action.b2_step{k} [1,1536] = b2 + tau_k @ W2[1536:] (tau half folded, exact) enc.action.W3 [1536,1536] enc.action.b3_pos [H,1536] = b3 + position_embedding[0:H] (H = 16 / 50 / 40) enc.future_tokens [32,1536] (N1.5 only) enc.decoder.W1 / .b1 [1024,1024], [1,1024] enc.decoder.W2 / .b2 [1024, 32|128|160] (columns padded), [1, 32|128|160] Every ``PlanTensor`` carries ``dtype`` in {"bf16", "bfp8_b", "fp32"} (the ttnn target; the torch tensor itself stays bf16 / fp32), ``layout`` in {"TILE", "ROW_MAJOR"}, ``memory`` in {"DRAM", "HOST"}, its ``category``, the checkpoint ``source_keys`` and a ``transform`` tag. The dtype policy (WL §7) is ``"bf16"`` (everything bf16) or ``"mixed"`` (``bfp8_b`` for the large matmul weights of the ViT / merger / N1.6 projector / LLM / VL-SA / DiT blocks, bf16 for vectors, ``proj_out_2``, the N1.5 projector and the encoders). """ from __future__ import annotations import hashlib import json import math import os import re from dataclasses import asdict, dataclass, field from pathlib import Path from typing import Any, Callable, Dict, Iterable, List, Mapping, Optional, Sequence, Tuple import torch import torch.nn.functional as F from models.experimental.gr00t.common.checkpoint import LazyCheckpoint from models.experimental.gr00t.common.configs import GR00T_VERSIONS, HF_SNAPSHOT_SHAS, TILE, GR00TConfig, get_config # Bump whenever a transform changes the *bytes* of a plan tensor (head padding, fusion layout, folds, RoPE layout). TRANSFORM_VERSION = "1" SKIP_PREFIX = "skip:" CATEGORIES: Tuple[str, ...] = ( "vision", "vision_merger", "vision_host", "projector", "llm", "llm_embed", "vlln", "vl_self_attention", "dit", "dit_precompute", "encoders", "encoders_slotted", "skip:vision_pool_head", "skip:lm_head", "skip:llm_final_norm", ) DTYPE_POLICIES: Tuple[str, ...] = ("bf16", "mixed") DEVICE_DTYPES: Tuple[str, ...] = ("bf16", "bfp8_b", "fp32") TIMESTEP_BUCKETS: Tuple[int, ...] = (0, 250, 500, 750) # -------------------------------------------------------------------------------------------------------------------- # 1. Key -> category rules (WL §2; first match wins) # -------------------------------------------------------------------------------------------------------------------- @dataclass(frozen=True) class Rule: """One row of the WL §2 tables: ``re.search(pattern, key)`` -> ``category``; ``strip`` is the container prefix removed to obtain the HF-relative canonical name (None for ``skip:*`` rows).""" pattern: str category: str strip: Optional[str] = None def matches(self, key: str) -> bool: return re.search(self.pattern, key) is not None def _rules_eagle(bb: str, lm: str, with_vlsa: bool) -> Tuple[Rule, ...]: """N1.5 / N1.6 share the Eagle layout; ``bb`` = backbone container, ``lm`` = language-model container.""" vis = f"{bb}vision_model." llm = f"{lm}model." rules = [ Rule(rf"^{re.escape(vis)}vision_model\.head\.", "skip:vision_pool_head"), Rule(rf"^{re.escape(lm)}lm_head\.weight$", "skip:lm_head"), Rule(rf"^{re.escape(vis)}", "vision", vis), Rule(rf"^{re.escape(bb)}mlp1\.", "projector", f"{bb}mlp1."), Rule(rf"^{re.escape(llm)}embed_tokens\.weight$", "llm_embed", llm), Rule(rf"^{re.escape(llm)}layers\.\d+\.", "llm", llm), Rule(rf"^{re.escape(llm)}norm\.weight$", "llm", llm), Rule(r"^action_head\.vlln\.", "vlln", "action_head."), ] if with_vlsa: rules.append(Rule(r"^action_head\.vl_self_attention\.", "vl_self_attention", "action_head.")) rules += [ Rule(r"^action_head\.model\.timestep_encoder\.", "dit_precompute", "action_head."), Rule(r"^action_head\.model\.transformer_blocks\.\d+\.norm1\.linear\.", "dit_precompute", "action_head."), Rule(r"^action_head\.model\.proj_out_1\.", "dit_precompute", "action_head."), Rule(r"^action_head\.model\.", "dit", "action_head."), Rule(r"^action_head\.(state_encoder|action_encoder|action_decoder)\.", "encoders_slotted", "action_head."), ] if with_vlsa: # N1.5 also has the learned future tokens rules.append(Rule(r"^action_head\.(position_embedding|future_tokens)\.weight$", "encoders", "action_head.")) else: rules.append(Rule(r"^action_head\.position_embedding\.weight$", "encoders", "action_head.")) return tuple(rules) def _rules_n17() -> Tuple[Rule, ...]: vis = "backbone.model.model.visual." llm = "backbone.model.model.language_model." return ( Rule(r"^backbone\.model\.lm_head\.weight$", "skip:lm_head"), Rule(rf"^{re.escape(llm)}norm\.weight$", "skip:llm_final_norm"), Rule(rf"^{re.escape(vis)}(merger|deepstack_merger_list)\.", "vision_merger", vis), Rule(rf"^{re.escape(vis)}pos_embed\.weight$", "vision_host", vis), Rule(rf"^{re.escape(vis)}", "vision", vis), Rule(rf"^{re.escape(llm)}embed_tokens\.weight$", "llm_embed", llm), Rule(rf"^{re.escape(llm)}layers\.\d+\.", "llm", llm), Rule(r"^action_head\.vlln\.", "vlln", "action_head."), Rule(r"^action_head\.vl_self_attention\.", "vl_self_attention", "action_head."), Rule(r"^action_head\.model\.timestep_encoder\.", "dit_precompute", "action_head."), Rule(r"^action_head\.model\.transformer_blocks\.\d+\.norm1\.linear\.", "dit_precompute", "action_head."), Rule(r"^action_head\.model\.proj_out_1\.", "dit_precompute", "action_head."), Rule(r"^action_head\.model\.", "dit", "action_head."), Rule(r"^action_head\.(state_encoder|action_encoder|action_decoder)\.", "encoders_slotted", "action_head."), Rule(r"^action_head\.position_embedding\.weight$", "encoders", "action_head."), ) RULES: Mapping[str, Tuple[Rule, ...]] = { "n15": _rules_eagle("backbone.eagle_model.", "backbone.eagle_model.language_model.", with_vlsa=True), "n16": _rules_eagle("backbone.model.", "backbone.model.language_model.", with_vlsa=False), "n17": _rules_n17(), } if tuple(len(RULES[v]) for v in GR00T_VERSIONS) != (15, 14, 15): raise RuntimeError("WL §2 rule counts must be 15 / 14 / 15") # WL §2.1-2.3 tensor counts per category (T1 fixture). EXPECTED_CATEGORY_COUNTS: Mapping[str, Mapping[str, int]] = { "n15": { "skip:vision_pool_head": 11, "skip:lm_head": 1, "vision": 437, "projector": 2, "llm_embed": 1, "llm": 133, "vlln": 2, "vl_self_attention": 64, "dit_precompute": 38, "dit": 194, "encoders_slotted": 14, "encoders": 2, }, "n16": { "skip:vision_pool_head": 11, "skip:lm_head": 1, "vision": 437, "projector": 6, "llm_embed": 1, "llm": 177, "vlln": 2, "dit_precompute": 70, "dit": 386, "encoders_slotted": 14, "encoders": 1, }, "n17": { "skip:lm_head": 1, "skip:llm_final_norm": 1, "vision_merger": 24, "vision_host": 1, "vision": 290, "llm_embed": 1, "llm": 176, "vlln": 2, "vl_self_attention": 64, "dit_precompute": 70, "dit": 386, "encoders_slotted": 14, "encoders": 1, }, } def _check_version(version: str) -> None: if version not in GR00T_VERSIONS: raise KeyError(f"unknown GR00T version {version!r}; expected one of {GR00T_VERSIONS}") def matching_rules(version: str, key: str) -> List[int]: """Indices of *all* rules matching ``key`` (the first one wins; the rest must be broader parents).""" _check_version(version) return [i for i, r in enumerate(RULES[version]) if r.matches(key)] def rule_of(version: str, key: str) -> Rule: """The first matching rule; raises ``KeyError`` for an unmatched key (no silent fallback).""" _check_version(version) for r in RULES[version]: if r.matches(key): return r raise KeyError(f"{version}: checkpoint key {key!r} matches no WL §2 rule") def category_of(version: str, key: str) -> str: return rule_of(version, key).category def is_skipped(version: str, key: str) -> bool: return category_of(version, key).startswith(SKIP_PREFIX) def canonical_name(version: str, key: str) -> str: """HF-relative canonical name (WL §2 'strip prefix -> canonical'); raises for ``skip:*`` keys.""" r = rule_of(version, key) if r.strip is None: raise KeyError(f"{version}: {key!r} is skipped ({r.category}) and has no canonical name") if not key.startswith(r.strip): raise KeyError(f"{version}: {key!r} does not start with its rule prefix {r.strip!r}") return key[len(r.strip) :] def categorize( version: str, keys: Optional[Iterable[str]] = None, ck: Optional[LazyCheckpoint] = None ) -> Dict[str, List[str]]: """Partition the checkpoint keys into WL §2 categories (first match wins). Every key must match at least one rule and lands in exactly one category; the returned lists are disjoint and their union is the input (asserted). ``keys`` defaults to all keys of the pinned snapshot. """ _check_version(version) if keys is None: ck = ck if ck is not None else LazyCheckpoint(version) keys = ck.keys() keys = list(keys) out: Dict[str, List[str]] = {} for k in keys: out.setdefault(category_of(version, k), []).append(k) assigned = [k for ks in out.values() for k in ks] if len(assigned) != len(keys) or set(assigned) != set(keys) or len(set(keys)) != len(keys): raise AssertionError(f"{version}: categorize() is not a partition of the key set") return out def skipped_keys(version: str, keys: Iterable[str]) -> List[str]: return [k for k in keys if is_skipped(version, k)] # -------------------------------------------------------------------------------------------------------------------- # 2. Byte accounting (WL §5 conventions: TILE pads the last two dims to 32; bfp8_b = 1088 B per 32x32 tile) # -------------------------------------------------------------------------------------------------------------------- BYTES_PER_ELEMENT: Mapping[str, float] = {"bf16": 2.0, "fp32": 4.0, "bfp8_b": 1088.0 / 1024.0} def pad32(n: int) -> int: return -(-n // TILE) * TILE def device_bytes(shape: Sequence[int], dtype: str, layout: str) -> float: """Bytes of one device tensor: TILE layout pads the last two dims (a 1-D shape is ``[1,N]``); ROW_MAJOR is exact.""" if dtype not in BYTES_PER_ELEMENT: raise ValueError(f"unknown device dtype {dtype!r}") s = [int(d) for d in shape] if not s: raise ValueError("empty shape") if layout == "TILE": if len(s) == 1: s = [1] + s s[-1] = pad32(s[-1]) s[-2] = pad32(s[-2]) elif layout != "ROW_MAJOR": raise ValueError(f"unknown layout {layout!r}") n = 1 for d in s: n *= d return n * BYTES_PER_ELEMENT[dtype] # -------------------------------------------------------------------------------------------------------------------- # 3. Pure host transforms (WL §4) # -------------------------------------------------------------------------------------------------------------------- def to_in_out(w: torch.Tensor) -> torch.Tensor: """``nn.Linear.weight [out,in] -> [in,out]`` so the device computes ``x @ W`` (WL §4.1 'Orientation').""" if w.ndim != 2: raise ValueError(f"expected a 2-D Linear weight, got {tuple(w.shape)}") return w.transpose(0, 1).contiguous() def to_row(v: torch.Tensor) -> torch.Tensor: """``[N] -> [1,N]`` (WL §4.1 'Biases / norm vectors').""" if v.ndim != 1: raise ValueError(f"expected a 1-D vector, got {tuple(v.shape)}") return v.reshape(1, -1).contiguous() def pad_dim(t: torch.Tensor, dim: int, to: int) -> torch.Tensor: """Zero-pad dimension ``dim`` of ``t`` up to ``to`` (no-op when already there; shrinking raises).""" cur = t.shape[dim] if cur == to: return t if cur > to: raise ValueError(f"cannot pad dim {dim} from {cur} down to {to}") pad = [0, 0] * t.ndim # F.pad pads from the last dim backwards: index of dim `dim` in the flat list is 2*(ndim-1-dim)+1 pad[2 * (t.ndim - 1 - dim) + 1] = to - cur return F.pad(t, pad) def siglip_conv_to_linear(w_conv: torch.Tensor) -> torch.Tensor: """N1.5 ``Conv2d`` kernel ``[1152,3,14,14] -> [1152,588]`` with ``(ph,pw,c)`` columns (pi0 unfold order). ``permute(0,2,3,1).reshape(out, ph*pw*c)`` (WL §4.2 N1.5 row; HL §4 check 3: 3.3e-6 vs golden with pos). """ if w_conv.ndim != 4: raise ValueError(f"expected a Conv2d kernel [out,c,kh,kw], got {tuple(w_conv.shape)}") out, c, kh, kw = w_conv.shape return w_conv.permute(0, 2, 3, 1).reshape(out, kh * kw * c).contiguous() def siglip_linear_patch_weight(w_lin: torch.Tensor, patch: int = 14, channels: int = 3) -> torch.Tensor: """N1.6 ``Linear(588 -> 1152)`` weight: already ``(ph,pw,c)`` columns -> **no permutation** (HL §3.1 correction).""" if w_lin.ndim != 2 or w_lin.shape[1] != patch * patch * channels: raise ValueError(f"expected [out, {patch * patch * channels}], got {tuple(w_lin.shape)}") return w_lin def permute_patch_columns( w_out_in: torch.Tensor, src: str, dst: str, patch: int = 14, channels: int = 3 ) -> torch.Tensor: """Reorder the ``patch*patch*channels`` input columns of a patch-embed Linear between ``"(ph,pw,c)"`` (pi0 / vendored SigLIP2 unfold, column ``ph*pw_n*c_n + pw*c_n + c``) and ``"(c,ph,pw)"`` (``F.unfold`` / conv order, column ``c*patch*patch + ph*patch + pw``) (HL §3.1 recipe iii).""" orders = ("(ph,pw,c)", "(c,ph,pw)") if src not in orders or dst not in orders: raise ValueError(f"orders must be in {orders}") if src == dst: return w_out_in out = w_out_in.shape[0] if src == "(ph,pw,c)": # columns (ph,pw,c) -> (c,ph,pw) return w_out_in.reshape(out, patch, patch, channels).permute(0, 3, 1, 2).reshape(out, -1).contiguous() return w_out_in.reshape(out, channels, patch, patch).permute(0, 2, 3, 1).reshape(out, -1).contiguous() def qwen3vl_conv3d_to_linear(w_conv3d: torch.Tensor) -> torch.Tensor: """N1.7 ``Conv3d`` kernel ``[1024,3,2,16,16] -> [1024,1536]``; column order ``(c,t,ph,pw)`` equals the processor's flatten, so a plain ``reshape`` is exact (WL §4.2 N1.7; HL §2.3(a): 1.75e-5 vs golden).""" if w_conv3d.ndim != 5: raise ValueError(f"expected a Conv3d kernel [out,c,t,kh,kw], got {tuple(w_conv3d.shape)}") return w_conv3d.reshape(w_conv3d.shape[0], -1).contiguous() def patch_embed_device_weight(w_out_in: torch.Tensor, in_pad: int) -> torch.Tensor: """``[out,in] -> [in_pad,out]``: transpose then zero-pad the K rows (588 -> 608); exact because K padding is matched by zero-padded patch vectors (WL §4.1 'Tile alignment').""" return pad_dim(to_in_out(w_out_in), 0, in_pad) def unfold_patches(pixels: torch.Tensor, patch: int, order: str = "(ph,pw,c)") -> torch.Tensor: """Reference/test helper: ``[B,C,H,W] -> [B*nh*nw, patch*patch*C]`` in raster token order with the given column order (pi0 ``_unfold_conv2d`` for ``"(ph,pw,c)"``; ``F.unfold`` order for ``"(c,ph,pw)"``).""" if pixels.ndim != 4: raise ValueError(f"expected [B,C,H,W], got {tuple(pixels.shape)}") b, c, h, w = pixels.shape if h % patch or w % patch: raise ValueError(f"image {h}x{w} is not a multiple of the patch size {patch}") nh, nw = h // patch, w // patch x = pixels.reshape(b, c, nh, patch, nw, patch) if order == "(ph,pw,c)": x = x.permute(0, 2, 4, 3, 5, 1) elif order == "(c,ph,pw)": x = x.permute(0, 2, 4, 1, 3, 5) else: raise ValueError(f"unknown patch column order {order!r}") return x.reshape(b * nh * nw, patch * patch * c).contiguous() def resize_pos_table_bilinear(pos: torch.Tensor, src_grid: int, dst_grid: int, antialias: bool = True) -> torch.Tensor: """N1.6 learned position table ``[src_grid^2, D] -> [dst_grid^2, D]`` (raster) in fp32 via ``F.interpolate(bilinear, align_corners=False, antialias)`` (``modeling_siglip2.py:593-646``; WL §4.2 N1.6: 9.5e-6 vs golden, antialias is a no-op for 16 -> 18 up-sampling, HL §3.2).""" if pos.ndim != 2 or pos.shape[0] != src_grid * src_grid: raise ValueError(f"expected [{src_grid * src_grid}, D], got {tuple(pos.shape)}") d = pos.shape[1] grid = pos.float().reshape(src_grid, src_grid, d).permute(2, 0, 1).unsqueeze(0) resized = F.interpolate(grid, size=(dst_grid, dst_grid), mode="bilinear", align_corners=False, antialias=antialias) return resized.reshape(d, dst_grid * dst_grid).transpose(0, 1).contiguous() def pad_heads_out(w_out_in: torch.Tensor, n_heads: int, head_dim: int, head_dim_pad: int) -> torch.Tensor: """Zero-pad every head of a Q/K/V projection on the **out** axis: ``[nh*hd, in] -> [nh*hd_pad, in]`` (WL §4.2 SigLIP 72 -> 96, §4.6 DiT 48 -> 64).""" if w_out_in.shape[0] != n_heads * head_dim: raise ValueError(f"out dim {w_out_in.shape[0]} != {n_heads} x {head_dim}") if head_dim_pad == head_dim: return w_out_in w = w_out_in.reshape(n_heads, head_dim, w_out_in.shape[1]) return F.pad(w, (0, 0, 0, head_dim_pad - head_dim)).reshape(n_heads * head_dim_pad, -1).contiguous() def pad_heads_vec(b: torch.Tensor, n_heads: int, head_dim: int, head_dim_pad: int) -> torch.Tensor: """Per-head zero padding of a Q/K/V bias ``[nh*hd] -> [nh*hd_pad]``.""" if b.ndim != 1 or b.shape[0] != n_heads * head_dim: raise ValueError(f"bias shape {tuple(b.shape)} != [{n_heads} x {head_dim}]") if head_dim_pad == head_dim: return b return F.pad(b.reshape(n_heads, head_dim), (0, head_dim_pad - head_dim)).reshape(-1).contiguous() def pad_heads_in(w_out_in: torch.Tensor, n_heads: int, head_dim: int, head_dim_pad: int) -> torch.Tensor: """Zero-pad the **in** axis of the output projection per head: ``[out, nh*hd] -> [out, nh*hd_pad]``; the zero rows (after transposition) absorb the zero V columns, so attention output is bit-exact (WL §4.2/§4.6).""" if w_out_in.shape[1] != n_heads * head_dim: raise ValueError(f"in dim {w_out_in.shape[1]} != {n_heads} x {head_dim}") if head_dim_pad == head_dim: return w_out_in w = w_out_in.reshape(w_out_in.shape[0], n_heads, head_dim) return F.pad(w, (0, head_dim_pad - head_dim)).reshape(w_out_in.shape[0], n_heads * head_dim_pad).contiguous() def fuse_projections(ws: Sequence[torch.Tensor], bs: Optional[Sequence[torch.Tensor]] = None): """Concatenate ``nn.Linear`` weights along the **out** axis (order as given, e.g. ``[Q | K | V]``) and transpose once: ``-> [in, sum(out)]``; biases ``-> [1, sum(out)]`` (WL §4.1 'Fused projections', the layout ``nlp_create_qkv_heads`` consumes). Head padding must be applied per tensor **before** calling this.""" in_dim = ws[0].shape[1] for w in ws: if w.ndim != 2 or w.shape[1] != in_dim: raise ValueError(f"all fused weights must share the in dim {in_dim}; got {[tuple(w.shape) for w in ws]}") w_fused = to_in_out(torch.cat(list(ws), dim=0)) if bs is None: return w_fused, None if len(bs) != len(ws) or any(b.shape[0] != w.shape[0] for b, w in zip(bs, ws)): raise ValueError("bias shapes do not match the fused weights") return w_fused, to_row(torch.cat(list(bs), dim=0)) def fuse_gate_up(w_gate: torch.Tensor, w_up: torch.Tensor, layout: str = "plain", tile: int = TILE) -> torch.Tensor: """SwiGLU ``[gate | up]`` fusion ``-> [in, 2*inter]`` (WL §4.1 'Fused gate/up'). ``"plain"``: consumer computes ``silu(y[:, :N]) * y[:, N:]``. ``"tile_interleaved"``: 32-column tile pairs ``[gate_t0, up_t0, gate_t1, up_t1, ...]`` == ``tt_dit.utils.tensor.prepare_for_fused_swiglu(W, ndev=1, gate_is_first=True)`` for the ``minimal_matmul(fuse_swiglu=True)`` epilogue. """ if w_gate.shape != w_up.shape: raise ValueError(f"gate {tuple(w_gate.shape)} and up {tuple(w_up.shape)} differ") w = to_in_out(torch.cat([w_gate, w_up], dim=0)) if layout == "plain": return w if layout == "tile_interleaved": return interleave_gate_up_tiles(w, tile) raise ValueError(f"unknown gate/up layout {layout!r}") def interleave_gate_up_tiles(w_gu: torch.Tensor, tile: int = TILE) -> torch.Tensor: """``[rows, 2N]`` with ``[gate | up]`` halves -> per-tile interleave ``[gate_t0, up_t0, gate_t1, up_t1, ...]`` (``prepare_for_fused_swiglu`` semantics, ``tt_dit/utils/tensor.py:20-45``, ``ndev=1``, ``gate_is_first=True``).""" rows, two_n = w_gu.shape n = two_n // 2 if n % tile: raise ValueError(f"half width {n} is not a multiple of the tile width {tile}") return w_gu.reshape(rows, 2, n // tile, tile).permute(0, 2, 1, 3).reshape(rows, two_n).contiguous() def deinterleave_gate_up_tiles(w_int: torch.Tensor, tile: int = TILE) -> torch.Tensor: """Inverse of :func:`interleave_gate_up_tiles`.""" rows, two_n = w_int.shape n = two_n // 2 return w_int.reshape(rows, n // tile, 2, tile).permute(0, 2, 1, 3).reshape(rows, two_n).contiguous() def reverse_permute(w: torch.Tensor, n_heads: int) -> torch.Tensor: """HF (rotate-half) -> Meta (pairwise) RoPE row layout for a q/k projection ``[nh*hd, in]``: ``view(nh, 2, hd/2, in).transpose(1, 2).reshape(out, in)`` (``tt_transformers/tt/load_checkpoints.py:891``).""" out, in_dim = w.shape return w.reshape(n_heads, 2, out // n_heads // 2, in_dim).transpose(1, 2).reshape(out, in_dim).contiguous() def reverse_permute_1d(v: torch.Tensor) -> torch.Tensor: """``(r0..r_{d/2-1}, i0..i_{d/2-1}) -> (r0, i0, r1, i1, ...)`` on the last dim (``load_checkpoints.py:899``); used for ``q_norm`` / ``k_norm`` under the ``meta`` RoPE layout.""" d = v.shape[-1] if d % 2: raise ValueError("last dim must be even") return torch.stack((v[..., : d // 2], v[..., d // 2 :]), dim=-1).flatten(start_dim=v.ndim - 1).contiguous() def rope_cos_sin_hf_to_meta(cos_hf: torch.Tensor, sin_hf: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: """HF half-duplicated ``[c0..c_{d/2-1}, c0..]`` -> Meta pairwise ``[c0, c0, c1, c1, ...]`` (``convert_rope_style_hf_to_meta``, ``load_checkpoints.py:921``).""" if cos_hf.shape != sin_hf.shape: raise ValueError("cos/sin shapes differ") half = cos_hf.shape[-1] // 2 return ( torch.repeat_interleave(cos_hf[..., :half], 2, dim=-1), torch.repeat_interleave(sin_hf[..., :half], 2, dim=-1), ) def unshuffle_column_permutation(hidden: int, merge: int = 2) -> Tuple[torch.Tensor, torch.Tensor]: """N1.6 pixel-unshuffle channel orders (HL §3.3). HF ``F.pixel_unshuffle`` column index is ``c*merge^2 + i*merge + j``; the device-friendly ``(i,j,c)`` order is ``i*merge*hidden + j*hidden + c``. Returns ``(perm, inv)`` with ``x_hf = x_dev[perm]`` and ``inv = argsort(perm)``. """ c = torch.arange(hidden).reshape(hidden, 1, 1) i = torch.arange(merge).reshape(1, merge, 1) j = torch.arange(merge).reshape(1, 1, merge) perm = (i * merge * hidden + j * hidden + c).reshape(-1) # indexed by hf position c*m^2 + i*m + j return perm, torch.argsort(perm) def fold_unshuffle_order_into_mlp1( ln_w: torch.Tensor, ln_b: torch.Tensor, fc1_w_out_in: torch.Tensor, hidden: int, merge: int = 2 ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: """Fold the ``(i,j,c)`` device unshuffle order into N1.6's ``mlp1``: ``LN(4608)`` weight/bias and the input columns of ``mlp1.1`` are indexed with ``inv`` (LayerNorm statistics are permutation-invariant) -- HL §3.3 / §1.2 #21 (7.6e-6 vs ``vision_projected_vit_embeds``).""" width = hidden * merge * merge if ln_w.shape != (width,) or ln_b.shape != (width,) or fc1_w_out_in.shape[1] != width: raise ValueError("mlp1 shapes do not match hidden * merge^2") _, inv = unshuffle_column_permutation(hidden, merge) return ln_w[inv].contiguous(), ln_b[inv].contiguous(), fc1_w_out_in[:, inv].contiguous() def pixel_unshuffle_tokens(x_raster: torch.Tensor, grid_h: int, grid_w: int, merge: int = 2, order: str = "(c,i,j)"): """Reference/test helper for the N1.6 2x2 merge: ``[grid_h*grid_w, D]`` raster -> ``[(gh/m)*(gw/m), m*m*D]`` with columns in HF ``"(c,i,j)"`` (== ``F.pixel_unshuffle``) or device ``"(i,j,c)"`` order (HL §3.3).""" n, d = x_raster.shape if n != grid_h * grid_w or grid_h % merge or grid_w % merge: raise ValueError(f"cannot unshuffle {tuple(x_raster.shape)} as {grid_h}x{grid_w} with merge {merge}") x = x_raster.reshape(grid_h // merge, merge, grid_w // merge, merge, d) if order == "(c,i,j)": x = x.permute(0, 2, 4, 1, 3) elif order == "(i,j,c)": x = x.permute(0, 2, 1, 3, 4) else: raise ValueError(f"unknown unshuffle order {order!r}") return x.reshape((grid_h // merge) * (grid_w // merge), merge * merge * d).contiguous() MERGER_ORDERS: Tuple[str, ...] = ("(mh,mw,c)", "(c,mh,mw)") def merger_column_permutation(hidden: int, merge: int, dst: str) -> Tuple[torch.Tensor, torch.Tensor]: """N1.7 merger input columns: HF concatenates the 2x2 window rows as ``(mh,mw,c)`` -> column ``(mh*m+mw)*hidden+c`` (HL §2.6). ``dst="(c,mh,mw)"`` is the pixel-unshuffle-like order ``c*m^2 + mh*m + mw``. Returns ``(perm, inv)`` with ``x_hf = x_dev[perm]``.""" if dst not in MERGER_ORDERS: raise ValueError(f"merger order must be in {MERGER_ORDERS}") if dst == "(mh,mw,c)": idx = torch.arange(hidden * merge * merge) return idx, idx mh = torch.arange(merge).reshape(merge, 1, 1) mw = torch.arange(merge).reshape(1, merge, 1) c = torch.arange(hidden).reshape(1, 1, hidden) perm = (c * merge * merge + mh * merge + mw).reshape(-1) # indexed by hf position (mh*m+mw)*hidden + c return perm, torch.argsort(perm) def permute_input_columns(w_out_in: torch.Tensor, inv: torch.Tensor) -> torch.Tensor: """``W[:, inv]`` for an ``[out,in]`` Linear whose input arrives in a permuted order (``x_hf = x_dev[perm]``).""" return w_out_in[:, inv].contiguous() def qwen3vl_pos_embeds( table: torch.Tensor, grid_thw: torch.Tensor, num_grid_per_side: int = 48, merge: int = 2 ) -> torch.Tensor: """N1.7 4-tap bilinear interpolation of the ``[side^2, D]`` position table per image ``(t,h,w)``, emitted in the tower's merge-window row order ``(bh, bw, mh, mw)`` (``fast_pos_embed_interpolate``, ``modeling_qwen3_vl.py:642``; WL §4.2 N1.7; HL §2.4). Returns ``[sum(t*h*w), D]`` in the table's dtype (call with fp32).""" if table.ndim != 2 or table.shape[0] != num_grid_per_side * num_grid_per_side: raise ValueError(f"expected [{num_grid_per_side ** 2}, D] table, got {tuple(table.shape)}") side = num_grid_per_side outs: List[torch.Tensor] = [] for t, h, w in grid_thw.reshape(-1, 3).tolist(): t, h, w = int(t), int(h), int(w) if h % merge or w % merge: raise ValueError(f"grid {h}x{w} not divisible by merge {merge}") h_grid = torch.linspace(0, side - 1, h) w_grid = torch.linspace(0, side - 1, w) h_floor = h_grid.int() w_floor = w_grid.int() h_ceil = (h_floor + 1).clamp(max=side - 1) w_ceil = (w_floor + 1).clamp(max=side - 1) dh = h_grid - h_floor dw = w_grid - w_floor corners = [ ((h_floor * side)[:, None] + w_floor[None, :], (1 - dh)[:, None] * (1 - dw)[None, :]), ((h_floor * side)[:, None] + w_ceil[None, :], (1 - dh)[:, None] * dw[None, :]), ((h_ceil * side)[:, None] + w_floor[None, :], dh[:, None] * (1 - dw)[None, :]), ((h_ceil * side)[:, None] + w_ceil[None, :], dh[:, None] * dw[None, :]), ] pos = torch.zeros(h * w, table.shape[1], dtype=table.dtype) for idx, wgt in corners: pos = pos + table[idx.reshape(-1).long()] * wgt.reshape(-1, 1).to(table.dtype) h_idx = torch.arange(h).reshape(h // merge, merge) w_idx = torch.arange(w).reshape(w // merge, merge) reorder = (h_idx[:, :, None, None] * w + w_idx[None, None, :, :]).transpose(1, 2).reshape(-1) outs.append(pos[reorder].repeat(t, 1)) return torch.cat(outs, dim=0).contiguous() def qwen3vl_vision_rope( grid_thw: torch.Tensor, head_dim: int = 64, theta: float = 10000.0, merge: int = 2 ) -> Tuple[torch.Tensor, torch.Tensor]: """N1.7 2-D vision RoPE tables ``cos, sin [N_patches, head_dim]`` in rotate-half layout for the merge-window row order: ``freqs = [row * inv_freq (16) | col * inv_freq (16)]`` duplicated to 64 (``rot_pos_emb``, ``modeling_qwen3_vl.py:603-640``; HL §2.5 exact).""" dim = head_dim // 2 inv_freq = 1.0 / (theta ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim)) rows: List[torch.Tensor] = [] cols: List[torch.Tensor] = [] for t, h, w in grid_thw.reshape(-1, 3).tolist(): t, h, w = int(t), int(h), int(w) hpos = torch.arange(h).unsqueeze(1).expand(h, w).reshape(h // merge, merge, w // merge, merge) wpos = torch.arange(w).unsqueeze(0).expand(h, w).reshape(h // merge, merge, w // merge, merge) rows.append(hpos.transpose(1, 2).reshape(-1).repeat(t)) cols.append(wpos.transpose(1, 2).reshape(-1).repeat(t)) r = torch.cat(rows).float() c = torch.cat(cols).float() freqs = torch.cat([torch.outer(r, inv_freq), torch.outer(c, inv_freq)], dim=-1) emb = torch.cat([freqs, freqs], dim=-1) return emb.cos(), emb.sin() def embed_rows(table: torch.Tensor, input_ids: torch.Tensor) -> torch.Tensor: """Host gather of token-embedding rows (WL §4.4 option A): ``table[input_ids]`` with a range check.""" ids = input_ids.reshape(-1).long() if ids.numel() and (ids.min() < 0 or ids.max() >= table.shape[0]): raise IndexError(f"input id out of range [0, {table.shape[0]})") return table[ids].reshape(*input_ids.shape, table.shape[1]) # -------------------------------------------------------------------------------------------------------------------- # 4. Sinusoids, timestep / AdaLN precompute, embodiment slots (WL §4.6-4.7) # -------------------------------------------------------------------------------------------------------------------- def tau_sinusoid(t: float, dim: int = 1536, max_period: float = 10000.0) -> torch.Tensor: """Action-encoder tau embedding ``[dim]`` fp32: ``f_i = exp(-i ln(max_period) / (dim/2))``, ``cat[sin(t f), cos(t f)]`` (sin first; arch-n15 §3.3, WL §4.7 ``b2_step``).""" half = dim // 2 freqs = torch.exp(-torch.arange(half, dtype=torch.float32) * math.log(max_period) / half) e = float(t) * freqs return torch.cat([torch.sin(e), torch.cos(e)]) def timestep_proj( t: float, dim: int = 256, flip_sin_to_cos: bool = True, downscale_freq_shift: int = 1, max_period: float = 10000.0, ) -> torch.Tensor: """diffusers ``Timesteps(dim, flip_sin_to_cos, downscale_freq_shift)`` projection ``[1, dim]`` fp32: ``e = t * exp(-ln(max_period) * arange(dim/2) / (dim/2 - shift))``, ``cat[sin, cos]`` flipped to ``cat[cos, sin]`` (WL §4.6; T8 sincos golden).""" half = dim // 2 exponent = -math.log(max_period) * torch.arange(half, dtype=torch.float32) / (half - downscale_freq_shift) e = float(t) * torch.exp(exponent) emb = torch.cat([torch.sin(e), torch.cos(e)]) if flip_sin_to_cos: emb = torch.cat([emb[half:], emb[:half]]) return emb.reshape(1, dim) def timestep_embedding( t: float, w1: torch.Tensor, b1: torch.Tensor, w2: torch.Tensor, b2: torch.Tensor, cfg: GR00TConfig ) -> torch.Tensor: """``temb = L2(SiLU(L1(Timesteps(t))))`` ``[1,1536]`` in fp32 from bf16 weights (WL §4.6; T8: <= 1.1e-6).""" d = cfg.dit proj = timestep_proj( t, d.timestep_embed_dim, d.timestep_flip_sin_to_cos, d.timestep_downscale_freq_shift, d.timestep_max_period ) h = F.silu(proj @ w1.float().T + b1.float()) return h @ w2.float().T + b2.float() def adaln_vectors( silu_temb: torch.Tensor, w: torch.Tensor, b: torch.Tensor, order: Tuple[str, str] ) -> Tuple[torch.Tensor, torch.Tensor]: """``(one_plus_scale, shift)`` ``[1,1536]`` fp32 from a 2-chunk AdaLN Linear ``[3072,1536]`` applied to ``SiLU(temb)``; ``order`` names the chunks: blocks ``("scale","shift")``, output head ``("shift","scale")`` (synthesis §0 fact 5; WL §4.6: swapped order -> PCC 0.61).""" if order not in (("scale", "shift"), ("shift", "scale")): raise ValueError(f"unknown AdaLN chunk order {order}") out = silu_temb @ w.float().T + b.float() first, second = out.chunk(2, dim=-1) scale, shift = (first, second) if order[0] == "scale" else (second, first) return (1.0 + scale).contiguous(), shift.contiguous() def fold_tau_bias(b2: torch.Tensor, w2_tau_rows: torch.Tensor, t: float, tau_dim: int = 1536) -> torch.Tensor: """``b2_step = b2 + tau_t @ W2[1536:]`` ``[1,1536]`` fp32 -- exact removal of the tau half of ``W2`` (WL §4.7).""" if w2_tau_rows.shape[0] != tau_dim: raise ValueError(f"expected the {tau_dim} tau rows of W2, got {tuple(w2_tau_rows.shape)}") return (b2.float() + tau_sinusoid(t, tau_dim) @ w2_tau_rows.float()).reshape(1, -1).contiguous() def fold_position_bias(b3: torch.Tensor, pos_rows: torch.Tensor) -> torch.Tensor: """``b3_pos = b3 + position_embedding[0:H]`` ``[H,1536]`` fp32 (WL §4.7; the golden ``action_encoder_out`` taps are *before* this add).""" return (b3.float().reshape(1, -1) + pos_rows.float()).contiguous() # -------------------------------------------------------------------------------------------------------------------- # 5. Device plan # -------------------------------------------------------------------------------------------------------------------- ROPE_LAYOUTS: Tuple[str, ...] = ("hf", "meta") GATE_UP_LAYOUTS: Tuple[str, ...] = ("plain", "tile_interleaved") UNSHUFFLE_ORDERS: Tuple[str, ...] = ("(c,i,j)", "(i,j,c)") @dataclass(frozen=True) class TransformOptions: """Byte-changing knobs of the host transforms (part of the cache identity, WL §8.2 item 2).""" pad_heads: bool = True # SigLIP 72 -> 96 and DiT 48 -> 64 per-head zero padding rope_layout: str = "hf" # "hf" rotate-half (own kernels / tt_dit) | "meta" pairwise (tt_transformers) gate_up_layout: str = "plain" # "plain" [gate | up] | "tile_interleaved" (fused-SwiGLU epilogue) n16_unshuffle_order: str = "(c,i,j)" # "(c,i,j)" HF pixel_unshuffle | "(i,j,c)" device order folded into mlp1 n17_merger_order: str = "(mh,mw,c)" # merger fc1 input columns: HF "(mh,mw,c)" | "(c,mh,mw)" precompute_dtype: str = "bf16" # AdaLN / b2_step / b3_pos vectors: "bf16" (rounded once) | "fp32" include_embed_table: bool = False # add `llm.embed_tokens` (ROW_MAJOR) to the plan def __post_init__(self) -> None: if self.rope_layout not in ROPE_LAYOUTS: raise ValueError(f"rope_layout must be in {ROPE_LAYOUTS}") if self.gate_up_layout not in GATE_UP_LAYOUTS: raise ValueError(f"gate_up_layout must be in {GATE_UP_LAYOUTS}") if self.n16_unshuffle_order not in UNSHUFFLE_ORDERS: raise ValueError(f"n16_unshuffle_order must be in {UNSHUFFLE_ORDERS}") if self.n17_merger_order not in MERGER_ORDERS: raise ValueError(f"n17_merger_order must be in {MERGER_ORDERS}") if self.precompute_dtype not in ("bf16", "fp32"): raise ValueError("precompute_dtype must be 'bf16' or 'fp32'") def tag(self) -> str: """Short deterministic string for cache directories.""" parts = [ "pad1" if self.pad_heads else "pad0", self.rope_layout, "gu-" + ("plain" if self.gate_up_layout == "plain" else "tiles"), "us-" + ("cij" if self.n16_unshuffle_order == "(c,i,j)" else "ijc"), "mg-" + ("hwc" if self.n17_merger_order == "(mh,mw,c)" else "chw"), "pc-" + self.precompute_dtype, "emb1" if self.include_embed_table else "emb0", ] return "-".join(parts) @dataclass(frozen=True) class PlanEntry: """Metadata of one device tensor, known from headers alone (WL §5 rows).""" name: str category: str shape: Tuple[int, ...] # logical device shape after the host transform (unpadded to tiles) dtype: str # ttnn target: bf16 | bfp8_b | fp32 layout: str # TILE | ROW_MAJOR memory: str # DRAM | HOST source_keys: Tuple[str, ...] transform: str bfp8_candidate: bool = False builder: Optional[Callable[["_Ctx"], torch.Tensor]] = field(default=None, repr=False, compare=False) @property def nbytes(self) -> float: return 0.0 if self.memory == "HOST" else device_bytes(self.shape, self.dtype, self.layout) @dataclass class PlanTensor: """One built device-plan tensor: the torch tensor plus the metadata the ``tt/`` uploader needs.""" name: str tensor: torch.Tensor dtype: str layout: str memory: str category: str source_keys: Tuple[str, ...] transform: str logical_shape: Tuple[int, ...] @property def nbytes(self) -> float: return 0.0 if self.memory == "HOST" else device_bytes(self.logical_shape, self.dtype, self.layout) def meta(self) -> Dict[str, Any]: return { "name": self.name, "category": self.category, "dtype": self.dtype, "layout": self.layout, "memory": self.memory, "logical_shape": list(self.logical_shape), "torch_dtype": str(self.tensor.dtype).replace("torch.", ""), "source_keys": list(self.source_keys), "transform": self.transform, "device_bytes": self.nbytes, } class DevicePlan(dict): """``{name: PlanTensor | PlanEntry}`` plus the identity of the plan.""" def __init__( self, version: str, embodiment_id: int, dtype_policy: str, options: TransformOptions, items: Iterable = () ): super().__init__(items) self.version = version self.embodiment_id = embodiment_id self.dtype_policy = dtype_policy self.options = options def by_category(self) -> Dict[str, List[str]]: out: Dict[str, List[str]] = {} for name, item in self.items(): out.setdefault(item.category, []).append(name) return out def summary(self) -> Dict[str, Any]: return byte_summary(self) class _Ctx: """Lazy reader shared by all builders of one plan: memoizes the embodiment slots and the fp32 precompute.""" def __init__(self, ck: LazyCheckpoint, cfg: GR00TConfig, embodiment_id: int, options: TransformOptions): self.ck = ck self.cfg = cfg self.e = embodiment_id self.opts = options self._slots: Dict[str, torch.Tensor] = {} self._silu_temb: Dict[int, torch.Tensor] = {} self._adaln: Dict[int, Dict[int, Tuple[torch.Tensor, torch.Tensor]]] = {} self._head: Optional[Dict[int, Tuple[torch.Tensor, torch.Tensor]]] = None self._pos_rows: Optional[torch.Tensor] = None # ---- raw reads def get(self, key: str) -> torch.Tensor: return self.ck.get(key) def slot(self, key: str) -> torch.Tensor: t = self._slots.get(key) if t is None: t = self.ck.slot(key, self.e) self._slots[key] = t return t def pos_rows(self) -> torch.Tensor: if self._pos_rows is None: h = self.cfg.encoders.action_horizon self._pos_rows = self.ck.rows(f"{self.cfg.checkpoint.action_head_prefix}position_embedding.weight", 0, h) return self._pos_rows # ---- precompute (fp32) def silu_temb(self, step: int) -> torch.Tensor: if step not in self._silu_temb: pre = f"{self.cfg.checkpoint.dit_prefix}timestep_encoder.timestep_embedder." temb = timestep_embedding( TIMESTEP_BUCKETS[step], self.get(pre + "linear_1.weight"), self.get(pre + "linear_1.bias"), self.get(pre + "linear_2.weight"), self.get(pre + "linear_2.bias"), self.cfg, ) self._silu_temb[step] = F.silu(temb) return self._silu_temb[step] def adaln(self, block: int, step: int) -> Tuple[torch.Tensor, torch.Tensor]: if block not in self._adaln: pre = f"{self.cfg.checkpoint.dit_prefix}transformer_blocks.{block}.norm1.linear." w = self.get(pre + "weight") b = self.get(pre + "bias") self._adaln[block] = { k: adaln_vectors(self.silu_temb(k), w, b, self.cfg.dit.adaln_order) for k in range(len(TIMESTEP_BUCKETS)) } return self._adaln[block][step] def out_head(self, step: int) -> Tuple[torch.Tensor, torch.Tensor]: if self._head is None: pre = f"{self.cfg.checkpoint.dit_prefix}proj_out_1." w = self.get(pre + "weight") b = self.get(pre + "bias") self._head = { k: adaln_vectors(self.silu_temb(k), w, b, self.cfg.dit.proj_out_1_order) for k in range(len(TIMESTEP_BUCKETS)) } return self._head[step] def precompute_cast(self, t: torch.Tensor) -> torch.Tensor: return t.to(torch.bfloat16) if self.opts.precompute_dtype == "bf16" else t.float() def _mixed_dtype(policy: str, bfp8_candidate: bool) -> str: return "bfp8_b" if (policy == "mixed" and bfp8_candidate) else "bf16" def _own(t: torch.Tensor) -> torch.Tensor: """Return ``t`` as a contiguous tensor that owns exactly its bytes. ``safe_open(...).get_slice(key)[e]`` hands back a view whose storage spans the whole ``[32, ...]`` tensor (e.g. 302 MB for a 9.4 MB slot); keeping such views alive pins the mmap and makes ``torch.save`` write the whole storage, so every plan tensor is compacted here (a no-op copy for tensors that already own their storage). """ if t.is_contiguous() and t.untyped_storage().nbytes() == t.numel() * t.element_size(): return t return t.clone(memory_format=torch.contiguous_format) class _PlanBuilder: """Enumerates the WL §5 device tensors of one version; every entry carries a builder closure.""" def __init__(self, version: str, embodiment_id: int, dtype_policy: str, options: TransformOptions): _check_version(version) if dtype_policy not in DTYPE_POLICIES: raise ValueError(f"dtype_policy must be in {DTYPE_POLICIES}") self.version = version self.cfg = get_config(version) self.e = int(embodiment_id) if not 0 <= self.e < self.cfg.encoders.max_num_embodiments: raise IndexError(f"embodiment id {self.e} out of range [0, {self.cfg.encoders.max_num_embodiments})") self.policy = dtype_policy self.opts = options self.entries: Dict[str, PlanEntry] = {} # ---- registration def add( self, name: str, category: str, shape: Sequence[int], source_keys: Sequence[str], transform: str, builder: Callable[[_Ctx], torch.Tensor], big: bool = False, layout: str = "TILE", memory: str = "DRAM", dtype: Optional[str] = None, ) -> None: if name in self.entries: raise KeyError(f"duplicate plan tensor name {name!r}") if category not in CATEGORIES: raise ValueError(f"unknown category {category!r}") entry = PlanEntry( name=name, category=category, shape=tuple(int(d) for d in shape), dtype=dtype if dtype is not None else _mixed_dtype(self.policy, big), layout=layout, memory=memory, source_keys=tuple(source_keys), transform=transform, bfp8_candidate=big, builder=builder, ) self.entries[name] = entry def add_linear(self, name: str, category: str, key_w: str, in_dim: int, out_dim: int, big: bool) -> None: """``key_w [out,in] -> name [in,out]``.""" self.add(name, category, (in_dim, out_dim), (key_w,), "T", lambda c, k=key_w: to_in_out(c.get(k)), big=big) def add_vec(self, name: str, category: str, key: str, n: int) -> None: """``key [n] -> name [1,n]``.""" self.add(name, category, (1, n), (key,), "row", lambda c, k=key: to_row(c.get(k))) def add_ln(self, name: str, category: str, key_prefix: str, n: int) -> None: self.add_vec(f"{name}.w", category, key_prefix + "weight", n) self.add_vec(f"{name}.b", category, key_prefix + "bias", n) # ---- sections def build(self) -> Dict[str, PlanEntry]: self._vision() self._projector() self._llm() self._vlln() self._vl_self_attention() self._dit() self._dit_precompute() self._encoders() return self.entries def _vision(self) -> None: cfg = self.cfg v = cfg.vision vp = cfg.checkpoint.vision_prefix pad = self.opts.pad_heads if v.family in ("siglip_v1", "siglip2_naflex"): base = vp + "vision_model." d = v.hidden hd, hdp = v.head_dim, (v.head_dim_pad if pad else v.head_dim) nh = v.n_heads kw = base + "embeddings.patch_embedding.weight" if v.family == "siglip_v1": self.add( "vision.patch_embed.W", "vision", (v.patch_in_features_pad, d), (kw,), "conv2d->(ph,pw,c) linear, T, pad K 588->608", lambda c, k=kw, p=v.patch_in_features_pad: patch_embed_device_weight( siglip_conv_to_linear(c.get(k)), p ), ) kp = base + "embeddings.position_embedding.weight" self.add( "vision.pos_embed", "vision", (v.pos_table_rows, d), (kp,), "as-is (raster)", lambda c, k=kp: c.get(k).contiguous(), ) else: self.add( "vision.patch_embed.W", "vision", (v.patch_in_features_pad, d), (kw,), "linear (ph,pw,c) as-is, T, pad K 588->608", lambda c, k=kw, p=v.patch_in_features_pad: patch_embed_device_weight( siglip_linear_patch_weight(c.get(k), v.patch_size), p ), ) kp = base + "embeddings.position_embedding.weight" dst = int(round(math.sqrt(v.patches_per_image))) self.add( "vision.pos_embed", "vision", (v.patches_per_image, d), (kp,), f"bilinear {v.pos_grid}x{v.pos_grid}->{dst}x{dst} (fp32, align_corners=False), bf16", lambda c, k=kp, s=v.pos_grid, t=dst: resize_pos_table_bilinear(c.get(k), s, t).to(torch.bfloat16), ) self.add_vec("vision.patch_embed.b", "vision", base + "embeddings.patch_embedding.bias", d) self.add_ln("vision.post_ln", "vision", base + "post_layernorm.", d) for i in range(v.n_layers): lp = f"{base}encoder.layers.{i}." self.add_ln(f"vision.layer.{i}.ln1", "vision", lp + "layer_norm1.", d) self.add_ln(f"vision.layer.{i}.ln2", "vision", lp + "layer_norm2.", d) wk = [lp + f"self_attn.{n}_proj.weight" for n in ("q", "k", "v")] bk = [lp + f"self_attn.{n}_proj.bias" for n in ("q", "k", "v")] self.add( f"vision.layer.{i}.Wqkv", "vision", (d, 3 * nh * hdp), tuple(wk), f"per-head pad {hd}->{hdp} (out), cat[Q|K|V], T", lambda c, ks=wk: fuse_projections([pad_heads_out(c.get(k), nh, hd, hdp) for k in ks])[0], big=True, ) self.add( f"vision.layer.{i}.bqkv", "vision", (1, 3 * nh * hdp), tuple(bk), f"per-head pad {hd}->{hdp}, cat, row", lambda c, ks=bk: to_row(torch.cat([pad_heads_vec(c.get(k), nh, hd, hdp) for k in ks])), ) ko = lp + "self_attn.out_proj.weight" self.add( f"vision.layer.{i}.Wo", "vision", (nh * hdp, d), (ko,), f"per-head pad {hd}->{hdp} (in), T", lambda c, k=ko: to_in_out(pad_heads_in(c.get(k), nh, hd, hdp)), big=True, ) self.add_vec(f"vision.layer.{i}.bo", "vision", lp + "self_attn.out_proj.bias", d) self.add_linear(f"vision.layer.{i}.fc1.W", "vision", lp + "mlp.fc1.weight", d, v.mlp_dim, big=True) self.add_vec(f"vision.layer.{i}.fc1.b", "vision", lp + "mlp.fc1.bias", v.mlp_dim) self.add_linear(f"vision.layer.{i}.fc2.W", "vision", lp + "mlp.fc2.weight", v.mlp_dim, d, big=True) self.add_vec(f"vision.layer.{i}.fc2.b", "vision", lp + "mlp.fc2.bias", d) elif v.family == "qwen3vl": d = v.hidden kw = vp + "patch_embed.proj.weight" self.add( "vision.patch_embed.W", "vision", (v.patch_in_features, d), (kw,), "conv3d reshape (c,t,ph,pw), T", lambda c, k=kw: to_in_out(qwen3vl_conv3d_to_linear(c.get(k))), ) self.add_vec("vision.patch_embed.b", "vision", vp + "patch_embed.proj.bias", d) kp = vp + "pos_embed.weight" self.add( "vision.pos_embed_table", "vision_host", (v.pos_table_rows, d), (kp,), "host fp32 table (qwen3vl_pos_embeds per grid_thw)", lambda c, k=kp: c.get(k).float(), layout="ROW_MAJOR", memory="HOST", dtype="fp32", ) for i in range(v.n_layers): bp = f"{vp}blocks.{i}." self.add_ln(f"vision.layer.{i}.ln1", "vision", bp + "norm1.", d) self.add_ln(f"vision.layer.{i}.ln2", "vision", bp + "norm2.", d) self.add_linear(f"vision.layer.{i}.Wqkv", "vision", bp + "attn.qkv.weight", d, 3 * d, big=True) self.add_vec(f"vision.layer.{i}.bqkv", "vision", bp + "attn.qkv.bias", 3 * d) self.add_linear(f"vision.layer.{i}.Wo", "vision", bp + "attn.proj.weight", d, d, big=True) self.add_vec(f"vision.layer.{i}.bo", "vision", bp + "attn.proj.bias", d) self.add_linear( f"vision.layer.{i}.fc1.W", "vision", bp + "mlp.linear_fc1.weight", d, v.mlp_dim, big=True ) self.add_vec(f"vision.layer.{i}.fc1.b", "vision", bp + "mlp.linear_fc1.bias", v.mlp_dim) self.add_linear( f"vision.layer.{i}.fc2.W", "vision", bp + "mlp.linear_fc2.weight", v.mlp_dim, d, big=True ) self.add_vec(f"vision.layer.{i}.fc2.b", "vision", bp + "mlp.linear_fc2.bias", d) # mergers (category vision_merger) pj = cfg.projector width = pj.in_dim # 4096 merge = v.merge_size _, inv = merger_column_permutation(d, merge, self.opts.n17_merger_order) permuted = self.opts.n17_merger_order != "(mh,mw,c)" tag = f"columns {self.opts.n17_merger_order}" mp = vp + "merger." self.add_ln("vision_merger.merger.norm", "vision_merger", mp + "norm.", d) self.add( "vision_merger.merger.fc1.W", "vision_merger", (width, width), (mp + "linear_fc1.weight",), f"T, {tag}", lambda c, k=mp + "linear_fc1.weight": to_in_out( permute_input_columns(c.get(k), inv) if permuted else c.get(k) ), big=True, ) self.add_vec("vision_merger.merger.fc1.b", "vision_merger", mp + "linear_fc1.bias", width) self.add_linear( "vision_merger.merger.fc2.W", "vision_merger", mp + "linear_fc2.weight", width, pj.out_dim, big=True ) self.add_vec("vision_merger.merger.fc2.b", "vision_merger", mp + "linear_fc2.bias", pj.out_dim) for k in range(pj.n_deepstack): dp = f"{vp}deepstack_merger_list.{k}." self.add( f"vision_merger.deepstack.{k}.norm.w", "vision_merger", (1, width), (dp + "norm.weight",), f"row, {tag}", lambda c, kk=dp + "norm.weight": to_row(c.get(kk)[inv] if permuted else c.get(kk)), ) self.add( f"vision_merger.deepstack.{k}.norm.b", "vision_merger", (1, width), (dp + "norm.bias",), f"row, {tag}", lambda c, kk=dp + "norm.bias": to_row(c.get(kk)[inv] if permuted else c.get(kk)), ) self.add( f"vision_merger.deepstack.{k}.fc1.W", "vision_merger", (width, width), (dp + "linear_fc1.weight",), f"T, {tag}", lambda c, kk=dp + "linear_fc1.weight": to_in_out( permute_input_columns(c.get(kk), inv) if permuted else c.get(kk) ), big=True, ) self.add_vec(f"vision_merger.deepstack.{k}.fc1.b", "vision_merger", dp + "linear_fc1.bias", width) self.add_linear( f"vision_merger.deepstack.{k}.fc2.W", "vision_merger", dp + "linear_fc2.weight", width, pj.out_dim, big=True, ) self.add_vec(f"vision_merger.deepstack.{k}.fc2.b", "vision_merger", dp + "linear_fc2.bias", pj.out_dim) else: raise ValueError(f"unknown vision family {v.family!r}") def _projector(self) -> None: cfg = self.cfg pj = cfg.projector pp = cfg.checkpoint.projector_prefix if pj.kind == "linear": self.add_linear("projector.fc1.W", "projector", pp + "0.weight", pj.in_dim, pj.out_dim, big=False) self.add_vec("projector.fc1.b", "projector", pp + "0.bias", pj.out_dim) elif pj.kind == "unshuffle_ln_mlp": fold = self.opts.n16_unshuffle_order == "(i,j,c)" hidden = cfg.vision.hidden merge = cfg.vision.merge_size tag = "folded (i,j,c)" if fold else "as-is (c,i,j)" k_lnw, k_lnb, k_fc1 = pp + "0.weight", pp + "0.bias", pp + "1.weight" def folded(c: _Ctx) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: return fold_unshuffle_order_into_mlp1(c.get(k_lnw), c.get(k_lnb), c.get(k_fc1), hidden, merge) self.add( "projector.ln.w", "projector", (1, pj.ln_dim), (k_lnw,), f"row, {tag}", (lambda c: to_row(folded(c)[0])) if fold else (lambda c: to_row(c.get(k_lnw))), ) self.add( "projector.ln.b", "projector", (1, pj.ln_dim), (k_lnb,), f"row, {tag}", (lambda c: to_row(folded(c)[1])) if fold else (lambda c: to_row(c.get(k_lnb))), ) self.add( "projector.fc1.W", "projector", (pj.in_dim, pj.hidden_dim), (k_fc1,), f"T, {tag}", (lambda c: to_in_out(folded(c)[2])) if fold else (lambda c: to_in_out(c.get(k_fc1))), big=True, ) self.add_vec("projector.fc1.b", "projector", pp + "1.bias", pj.hidden_dim) self.add_linear("projector.fc2.W", "projector", pp + "3.weight", pj.hidden_dim, pj.out_dim, big=True) self.add_vec("projector.fc2.b", "projector", pp + "3.bias", pj.out_dim) elif pj.kind == "qwen3vl_merger": return # lives in `vision_merger` else: raise ValueError(f"unknown projector kind {pj.kind!r}") def _llm(self) -> None: cfg = self.cfg l = cfg.llm lp = cfg.checkpoint.llm_prefix meta = self.opts.rope_layout == "meta" d, nh, nkv, hd = l.hidden, l.n_heads, l.n_kv_heads, l.head_dim for i in range(l.n_layers_executed): p = f"{lp}layers.{i}." wk = [p + f"self_attn.{n}_proj.weight" for n in ("q", "k", "v")] def qkv(c: _Ctx, ks: List[str] = wk) -> torch.Tensor: q, k, v = (c.get(x) for x in ks) if meta: q, k = reverse_permute(q, nh), reverse_permute(k, nkv) return fuse_projections([q, k, v])[0] self.add( f"llm.layer.{i}.Wqkv", "llm", (d, (nh + 2 * nkv) * hd), tuple(wk), "cat[Q|K|V], T" + (", meta rope rows" if meta else ""), qkv, big=True, ) self.add_linear(f"llm.layer.{i}.Wo", "llm", p + "self_attn.o_proj.weight", nh * hd, d, big=True) gk, uk = p + "mlp.gate_proj.weight", p + "mlp.up_proj.weight" self.add( f"llm.layer.{i}.W_gate_up", "llm", (d, 2 * l.intermediate), (gk, uk), f"cat[gate|up], T, {self.opts.gate_up_layout}", lambda c, g=gk, u=uk: fuse_gate_up(c.get(g), c.get(u), self.opts.gate_up_layout), big=True, ) self.add_linear(f"llm.layer.{i}.W_down", "llm", p + "mlp.down_proj.weight", l.intermediate, d, big=True) self.add_vec(f"llm.layer.{i}.input_ln.w", "llm", p + "input_layernorm.weight", d) self.add_vec(f"llm.layer.{i}.post_attn_ln.w", "llm", p + "post_attention_layernorm.weight", d) for n in ("q_norm", "k_norm"): key = p + f"self_attn.{n}.weight" self.add( f"llm.layer.{i}.{n}.w", "llm", (1, hd), (key,), "row" + (", reverse_permute_1d" if meta else ""), lambda c, k=key: to_row(reverse_permute_1d(c.get(k)) if meta else c.get(k)), ) if l.final_norm_applied: self.add_vec("llm.final_norm.w", "llm", lp + "norm.weight", d) if self.opts.include_embed_table: ke = lp + "embed_tokens.weight" self.add( "llm.embed_tokens", "llm_embed", (l.vocab, d), (ke,), "as-is (ROW_MAJOR table)", lambda c, k=ke: c.get(k).contiguous(), layout="ROW_MAJOR", ) def _vlln(self) -> None: vp = self.cfg.checkpoint.vlln_prefix self.add_vec("vlln.w", "vlln", vp + "weight", self.cfg.vl_adapter.dim) self.add_vec("vlln.b", "vlln", vp + "bias", self.cfg.vl_adapter.dim) def _vl_self_attention(self) -> None: cfg = self.cfg a = cfg.vl_adapter if a.n_blocks == 0: return pre = cfg.checkpoint.vl_self_attention_prefix if pre is None: raise ValueError(f"{self.version}: vl_adapter.n_blocks={a.n_blocks} but no vl_self_attention prefix") d = a.dim for i in range(a.n_blocks): bp = f"{pre}transformer_blocks.{i}." wk = [bp + f"attn1.to_{n}.weight" for n in ("q", "k", "v")] bk = [bp + f"attn1.to_{n}.bias" for n in ("q", "k", "v")] self.add( f"vlsa.block.{i}.Wqkv", "vl_self_attention", (d, 3 * d), tuple(wk), "cat[Q|K|V], T", lambda c, ks=wk: fuse_projections([c.get(k) for k in ks])[0], big=True, ) self.add( f"vlsa.block.{i}.bqkv", "vl_self_attention", (1, 3 * d), tuple(bk), "cat, row", lambda c, ks=bk: to_row(torch.cat([c.get(k) for k in ks])), ) self.add_linear(f"vlsa.block.{i}.Wo", "vl_self_attention", bp + "attn1.to_out.0.weight", d, d, big=True) self.add_vec(f"vlsa.block.{i}.bo", "vl_self_attention", bp + "attn1.to_out.0.bias", d) self.add_ln(f"vlsa.block.{i}.norm1", "vl_self_attention", bp + "norm1.", d) self.add_ln(f"vlsa.block.{i}.norm3", "vl_self_attention", bp + "norm3.", d) self.add_linear( f"vlsa.block.{i}.ff1.W", "vl_self_attention", bp + "ff.net.0.proj.weight", d, a.ff_dim, big=True ) self.add_vec(f"vlsa.block.{i}.ff1.b", "vl_self_attention", bp + "ff.net.0.proj.bias", a.ff_dim) self.add_linear(f"vlsa.block.{i}.ff2.W", "vl_self_attention", bp + "ff.net.2.weight", a.ff_dim, d, big=True) self.add_vec(f"vlsa.block.{i}.ff2.b", "vl_self_attention", bp + "ff.net.2.bias", d) def _dit(self) -> None: cfg = self.cfg dt = cfg.dit dp = cfg.checkpoint.dit_prefix w, nh, hd = dt.width, dt.n_heads, dt.head_dim hdp = dt.head_dim_pad if self.opts.pad_heads else hd inner = nh * hdp cross = set(dt.cross_block_indices) tag = f"per-head pad {hd}->{hdp}" for i in range(dt.n_blocks): bp = f"{dp}transformer_blocks.{i}." kq, kk, kv = (bp + f"attn1.to_{n}.weight" for n in ("q", "k", "v")) bq, bk, bv = (bp + f"attn1.to_{n}.bias" for n in ("q", "k", "v")) if i in cross: cd = dt.cross_attention_dim self.add( f"dit.block.{i}.Wq", "dit", (w, inner), (kq,), f"{tag} (out), T", lambda c, k=kq: to_in_out(pad_heads_out(c.get(k), nh, hd, hdp)), big=True, ) self.add( f"dit.block.{i}.bq", "dit", (1, inner), (bq,), f"{tag}, row", lambda c, k=bq: to_row(pad_heads_vec(c.get(k), nh, hd, hdp)), ) self.add( f"dit.block.{i}.Wkv", "dit", (cd, 2 * inner), (kk, kv), f"{tag} (out), cat[K|V], T", lambda c, ks=(kk, kv): fuse_projections([pad_heads_out(c.get(k), nh, hd, hdp) for k in ks])[0], big=True, ) self.add( f"dit.block.{i}.bkv", "dit", (1, 2 * inner), (bk, bv), f"{tag}, cat, row", lambda c, ks=(bk, bv): to_row(torch.cat([pad_heads_vec(c.get(k), nh, hd, hdp) for k in ks])), ) else: self.add( f"dit.block.{i}.Wqkv", "dit", (w, 3 * inner), (kq, kk, kv), f"{tag} (out), cat[Q|K|V], T", lambda c, ks=(kq, kk, kv): fuse_projections([pad_heads_out(c.get(k), nh, hd, hdp) for k in ks])[0], big=True, ) self.add( f"dit.block.{i}.bqkv", "dit", (1, 3 * inner), (bq, bk, bv), f"{tag}, cat, row", lambda c, ks=(bq, bk, bv): to_row(torch.cat([pad_heads_vec(c.get(k), nh, hd, hdp) for k in ks])), ) ko = bp + "attn1.to_out.0.weight" self.add( f"dit.block.{i}.Wo", "dit", (inner, w), (ko,), f"{tag} (in), T", lambda c, k=ko: to_in_out(pad_heads_in(c.get(k), nh, hd, hdp)), big=True, ) self.add_vec(f"dit.block.{i}.bo", "dit", bp + "attn1.to_out.0.bias", w) self.add_linear(f"dit.block.{i}.ff1.W", "dit", bp + "ff.net.0.proj.weight", w, dt.ff_dim, big=True) self.add_vec(f"dit.block.{i}.ff1.b", "dit", bp + "ff.net.0.proj.bias", dt.ff_dim) self.add_linear(f"dit.block.{i}.ff2.W", "dit", bp + "ff.net.2.weight", dt.ff_dim, w, big=True) self.add_vec(f"dit.block.{i}.ff2.b", "dit", bp + "ff.net.2.bias", w) self.add_linear("dit.proj_out_2.W", "dit", dp + "proj_out_2.weight", w, dt.out_dim, big=False) self.add_vec("dit.proj_out_2.b", "dit", dp + "proj_out_2.bias", dt.out_dim) def _dit_precompute(self) -> None: cfg = self.cfg dt = cfg.dit dp = cfg.checkpoint.dit_prefix te = [dp + f"timestep_encoder.timestep_embedder.linear_{j}.{n}" for j in (1, 2) for n in ("weight", "bias")] pdt = self.opts.precompute_dtype for i in range(dt.n_blocks): src = tuple(te + [dp + f"transformer_blocks.{i}.norm1.linear.{n}" for n in ("weight", "bias")]) for k in range(len(TIMESTEP_BUCKETS)): self.add( f"dit.block.{i}.adaln.step{k}.one_plus_scale", "dit_precompute", (1, dt.width), src, f"1 + scale, (scale, shift) = chunk(norm1.linear(SiLU(temb_t{TIMESTEP_BUCKETS[k]}))), fp32->{pdt}", lambda c, b=i, s=k: c.precompute_cast(c.adaln(b, s)[0]), dtype=pdt, ) self.add( f"dit.block.{i}.adaln.step{k}.shift", "dit_precompute", (1, dt.width), src, f"shift of norm1.linear(SiLU(temb_t{TIMESTEP_BUCKETS[k]})), fp32->{pdt}", lambda c, b=i, s=k: c.precompute_cast(c.adaln(b, s)[1]), dtype=pdt, ) src = tuple(te + [dp + f"proj_out_1.{n}" for n in ("weight", "bias")]) for k in range(len(TIMESTEP_BUCKETS)): self.add( f"dit.out_head.adaln.step{k}.one_plus_scale", "dit_precompute", (1, dt.width), src, f"1 + scale, (shift, scale) = chunk(proj_out_1(SiLU(temb_t{TIMESTEP_BUCKETS[k]}))), fp32->{pdt}", lambda c, s=k: c.precompute_cast(c.out_head(s)[0]), dtype=pdt, ) self.add( f"dit.out_head.adaln.step{k}.shift", "dit_precompute", (1, dt.width), src, f"shift of proj_out_1(SiLU(temb_t{TIMESTEP_BUCKETS[k]})), fp32->{pdt}", lambda c, s=k: c.precompute_cast(c.out_head(s)[1]), dtype=pdt, ) def _encoders(self) -> None: cfg = self.cfg en = cfg.encoders ap = cfg.checkpoint.action_head_prefix e = self.e pdt = self.opts.precompute_dtype s_in, s_pad = en.max_state_dim, en.state_dim_pad a_in, a_pad = en.max_action_dim, en.action_dim_pad hid, emb, h = en.hidden_size, en.input_embedding_dim, en.action_horizon sp = ap + "state_encoder." self.add( "enc.state.W1", "encoders", (s_pad, hid), (sp + "layer1.W",), f"slot[{e}] [in,out], pad rows {s_in}->{s_pad}", lambda c: pad_dim(c.slot(sp + "layer1.W"), 0, s_pad).contiguous(), ) self.add( "enc.state.b1", "encoders", (1, hid), (sp + "layer1.b",), f"slot[{e}], row", lambda c: to_row(c.slot(sp + "layer1.b")), ) self.add( "enc.state.W2", "encoders", (hid, emb), (sp + "layer2.W",), f"slot[{e}] [in,out]", lambda c: c.slot(sp + "layer2.W").contiguous(), ) self.add( "enc.state.b2", "encoders", (1, emb), (sp + "layer2.b",), f"slot[{e}], row", lambda c: to_row(c.slot(sp + "layer2.b")), ) acp = ap + "action_encoder." self.add( "enc.action.W1", "encoders", (a_pad, emb), (acp + "W1.W",), f"slot[{e}] [in,out], pad rows {a_in}->{a_pad}", lambda c: pad_dim(c.slot(acp + "W1.W"), 0, a_pad).contiguous(), ) self.add( "enc.action.b1", "encoders", (1, emb), (acp + "W1.b",), f"slot[{e}], row", lambda c: to_row(c.slot(acp + "W1.b")), ) self.add( "enc.action.W2a", "encoders", (emb, emb), (acp + "W2.W",), f"slot[{e}] rows [0:{emb}] (a_emb half)", lambda c: c.slot(acp + "W2.W")[:emb].contiguous(), ) for k, t in enumerate(TIMESTEP_BUCKETS): self.add( f"enc.action.b2_step{k}", "encoders", (1, emb), (acp + "W2.W", acp + "W2.b"), f"b2 + tau_{t} @ W2[{emb}:] (fp32) -> {pdt}", lambda c, tt=t: c.precompute_cast( fold_tau_bias(c.slot(acp + "W2.b"), c.slot(acp + "W2.W")[emb:], tt, en.tau_dim) ), dtype=pdt, ) self.add( "enc.action.W3", "encoders", (emb, emb), (acp + "W3.W",), f"slot[{e}] [in,out]", lambda c: c.slot(acp + "W3.W").contiguous(), ) self.add( "enc.action.b3_pos", "encoders", (h, emb), (acp + "W3.b", ap + "position_embedding.weight"), f"b3 + position_embedding[0:{h}] (fp32) -> {pdt}", lambda c: c.precompute_cast(fold_position_bias(c.slot(acp + "W3.b"), c.pos_rows())), dtype=pdt, ) if cfg.dit.n_future_tokens: kf = ap + "future_tokens.weight" self.add( "enc.future_tokens", "encoders", (cfg.dit.n_future_tokens, emb), (kf,), "as-is", lambda c: c.get(kf).contiguous(), ) dcp = ap + "action_decoder." self.add( "enc.decoder.W1", "encoders", (hid, hid), (dcp + "layer1.W",), f"slot[{e}] [in,out]", lambda c: c.slot(dcp + "layer1.W").contiguous(), ) self.add( "enc.decoder.b1", "encoders", (1, hid), (dcp + "layer1.b",), f"slot[{e}], row", lambda c: to_row(c.slot(dcp + "layer1.b")), ) self.add( "enc.decoder.W2", "encoders", (hid, a_pad), (dcp + "layer2.W",), f"slot[{e}] [in,out], pad cols {a_in}->{a_pad}", lambda c: pad_dim(c.slot(dcp + "layer2.W"), 1, a_pad).contiguous(), ) self.add( "enc.decoder.b2", "encoders", (1, a_pad), (dcp + "layer2.b",), f"slot[{e}], pad {a_in}->{a_pad}, row", lambda c: to_row(pad_dim(c.slot(dcp + "layer2.b"), 0, a_pad)), ) def enumerate_plan( version: str, embodiment_id: Optional[int] = None, dtype_policy: str = "bf16", options: Optional[TransformOptions] = None, pad_heads: Optional[bool] = None, ) -> DevicePlan: """All WL §5 device tensors of ``version`` as :class:`PlanEntry` (no weight is read). ``embodiment_id`` defaults to the config's canonical embodiment; ``pad_heads`` overrides ``options.pad_heads``. """ _check_version(version) cfg = get_config(version) e = cfg.canonical_embodiment_id if embodiment_id is None else int(embodiment_id) opts = options if options is not None else TransformOptions() if pad_heads is not None and pad_heads != opts.pad_heads: opts = TransformOptions(**{**asdict(opts), "pad_heads": pad_heads}) b = _PlanBuilder(version, e, dtype_policy, opts) return DevicePlan(version, e, dtype_policy, opts, b.build().items()) def build_device_plan( version: str, embodiment_id: Optional[int] = None, dtype_policy: str = "bf16", pad_heads: bool = True, options: Optional[TransformOptions] = None, ck: Optional[LazyCheckpoint] = None, categories: Optional[Iterable[str]] = None, names: Optional[Iterable[str]] = None, ) -> DevicePlan: """Read the checkpoint and produce every plan tensor (``{name: PlanTensor}``; names per the module docstring). ``categories`` / ``names`` restrict the build (tests, partial loads). Every built tensor is checked against the header-only :func:`enumerate_plan` shape; the torch dtype is bf16 for weights (as stored, zero-padded) and ``options.precompute_dtype`` for the host-computed vectors. Source keys are read in shard order lazily; nothing but the requested tensors is materialised. """ plan_meta = enumerate_plan(version, embodiment_id, dtype_policy, options, pad_heads) cfg = get_config(version) ck = ck if ck is not None else LazyCheckpoint(version) if ck.version is not None and ck.version != version: raise ValueError(f"checkpoint is {ck.version}, plan is {version}") ctx = _Ctx(ck, cfg, plan_meta.embodiment_id, plan_meta.options) wanted = set(plan_meta) if categories is not None: cats = set(categories) unknown = cats - set(CATEGORIES) if unknown: raise KeyError(f"unknown categories {sorted(unknown)}") wanted = {n for n in wanted if plan_meta[n].category in cats} if names is not None: req = set(names) missing = req - set(plan_meta) if missing: raise KeyError(f"unknown plan tensor names {sorted(missing)[:5]}") wanted &= req # shard-sorted build order keeps the reads streaming (WL §1.2) order = sorted(wanted, key=lambda n: (min(ck.shard_of(k) for k in plan_meta[n].source_keys), n)) out = DevicePlan(version, plan_meta.embodiment_id, dtype_policy, plan_meta.options) for name in order: entry = plan_meta[name] if entry.builder is None: raise RuntimeError(f"{name}: entry has no builder") t = entry.builder(ctx) if tuple(t.shape) != entry.shape: raise ValueError(f"{name}: built shape {tuple(t.shape)} != planned {entry.shape}") t = _own(t) out[name] = PlanTensor( name=name, tensor=t, dtype=entry.dtype, layout=entry.layout, memory=entry.memory, category=entry.category, source_keys=entry.source_keys, transform=entry.transform, logical_shape=entry.shape, ) return out # -------------------------------------------------------------------------------------------------------------------- # 6. Byte accounting summary (WL §6) # -------------------------------------------------------------------------------------------------------------------- MIB = float(1 << 20) # WL §6.1 device-byte totals (one embodiment, token table excluded) in MiB. SPEC_TOTALS_MIB: Mapping[str, Mapping[str, float]] = { "n15": {"bf16": 3493.9, "mixed": 1897.8, "embed_table": 592.5}, "n16": {"bf16": 4530.5, "mixed": 2457.2, "embed_table": 592.5}, "n17": {"bf16": 4782.5, "mixed": 2592.5, "embed_table": 593.5}, } # WL §5 per-category device MiB (bf16, mixed) and device tensor counts (token table excluded). SPEC_CATEGORY_MIB: Mapping[str, Mapping[str, Tuple[float, float, int]]] = { "n15": { "dit": (1010.1, 545.1, 146), "dit_precompute": (12.8, 12.8, 136), "encoders": (15.2, 15.2, 18), "llm": (1155.3, 615.3, 97), "projector": (4.6, 4.6, 2), "vision": (905.3, 494.2, 329), "vl_self_attention": (390.5, 210.5, 48), "vlln": (0.2, 0.2, 2), }, "n16": { "dit": (2017.1, 1087.1, 290), "dit_precompute": (24.8, 24.8, 264), "encoders": (15.8, 15.8, 17), "llm": (1540.4, 820.4, 129), "projector": (26.8, 14.6, 6), "vision": (905.5, 494.4, 329), "vlln": (0.2, 0.2, 2), }, "n17": { "dit": (2017.1, 1087.1, 290), "dit_precompute": (24.8, 24.8, 264), "encoders": (16.0, 16.0, 17), "llm": (1540.2, 820.2, 128), "vision": (598.6, 328.6, 290), "vision_merger": (195.1, 105.1, 24), "vl_self_attention": (390.5, 210.5, 48), "vlln": (0.2, 0.2, 2), }, } def byte_summary(plan: Mapping[str, Any]) -> Dict[str, Any]: """Per-category device bytes of a plan (entries or built tensors): ``{"categories": {cat: {"count", "bytes", "mib"}}, "total_bytes", "total_mib", "embed_table_mib", "n_tensors"}``; the token table (``llm_embed``) and host- only tensors are reported separately and excluded from the total (WL §6.1 convention).""" cats: Dict[str, Dict[str, float]] = {} total = 0.0 embed = 0.0 n = 0 for item in plan.values(): b = item.nbytes c = cats.setdefault(item.category, {"count": 0, "bytes": 0.0, "mib": 0.0}) c["count"] += 1 c["bytes"] += b c["mib"] = c["bytes"] / MIB if item.category == "llm_embed": embed += b elif item.memory != "HOST": total += b n += 1 version = getattr(plan, "version", None) if version is not None and embed == 0.0: cfg = get_config(version) embed = device_bytes((cfg.llm.vocab, cfg.llm.hidden), "bf16", "ROW_MAJOR") return { "categories": cats, "total_bytes": total, "total_mib": total / MIB, "embed_table_mib": embed / MIB, "n_tensors": n, } def assert_byte_totals(version: str, rel_tol: float = 0.01) -> Dict[str, Any]: """Recompute WL §6.1 totals from headers for both policies and assert they agree with the spec within ``rel_tol`` (T23). Returns the two summaries.""" out: Dict[str, Any] = {} for policy in DTYPE_POLICIES: s = byte_summary(enumerate_plan(version, dtype_policy=policy)) spec = SPEC_TOTALS_MIB[version][policy] if abs(s["total_mib"] - spec) > rel_tol * spec: raise AssertionError(f"{version}/{policy}: {s['total_mib']:.1f} MiB != spec {spec} MiB") emb_spec = SPEC_TOTALS_MIB[version]["embed_table"] if abs(s["embed_table_mib"] - emb_spec) > rel_tol * emb_spec: raise AssertionError(f"{version}: embed table {s['embed_table_mib']:.1f} MiB != spec {emb_spec}") out[policy] = s return out # -------------------------------------------------------------------------------------------------------------------- # 7. On-disk host cache skeleton (WL §8; Tier 0 torch tensors -- the ttnn .tensorbin tier lives in tt/) # -------------------------------------------------------------------------------------------------------------------- CACHE_ENV = "GR00T_TT_CACHE" CACHE_MARKER = ".complete" CACHE_MANIFEST = "manifest.json" CACHE_FORMAT = "gr00t-host-plan-v1" def default_cache_root() -> Path: return Path(os.environ.get(CACHE_ENV, str(Path.home() / ".cache" / "gr00t-tt"))) def _sha256_file(path: Path, chunk: int = 1 << 22) -> str: h = hashlib.sha256() with open(path, "rb") as fh: while True: b = fh.read(chunk) if not b: break h.update(b) return h.hexdigest() class HostPlanCache: """Per-(version, snapshot sha, embodiment, dtype policy, arch, transform version + options) directory of torch ``.pt`` plan tensors with a manifest and a ``.complete`` marker written last (WL §8.2 items 1-2, v3 semantics of ``models/common/weight_cache.py``: a missing marker, a manifest mismatch or a wrong file size => cold load).""" def __init__(self, root: Optional[Path] = None, arch: str = "p150a"): self.root = Path(root) if root is not None else default_cache_root() self.arch = arch def directory( self, version: str, embodiment_id: int, dtype_policy: str, options: Optional[TransformOptions] = None, snapshot_sha: Optional[str] = None, ) -> Path: _check_version(version) opts = options if options is not None else TransformOptions() sha = (snapshot_sha or HF_SNAPSHOT_SHAS[version])[:12] return ( self.root / version / sha / f"emb{int(embodiment_id)}" / dtype_policy / self.arch / f"xf-{TRANSFORM_VERSION}-{opts.tag()}" / "host" ) def directory_for(self, plan: DevicePlan) -> Path: return self.directory(plan.version, plan.embodiment_id, plan.dtype_policy, plan.options) @staticmethod def file_name(name: str) -> str: return name + ".pt" def is_complete(self, directory: Path) -> bool: directory = Path(directory) marker = directory / CACHE_MARKER if not marker.is_file() or not (directory / CACHE_MANIFEST).is_file(): return False try: return marker.read_text().strip() == CACHE_FORMAT except OSError: return False def save(self, plan: DevicePlan, directory: Optional[Path] = None, hash_files: bool = True) -> Dict[str, Any]: """Write every tensor as ``.pt`` (``torch.save`` of the bare tensor), then the manifest, then the marker. An existing marker is removed first so a crash mid-write leaves an incomplete (= cold) directory.""" directory = Path(directory) if directory is not None else self.directory_for(plan) directory.mkdir(parents=True, exist_ok=True) marker = directory / CACHE_MARKER if marker.exists(): marker.unlink() tensors: Dict[str, Any] = {} total = 0 for name in sorted(plan): pt = plan[name] if not isinstance(pt, PlanTensor): raise TypeError(f"{name}: save() needs built PlanTensors, got {type(pt).__name__}") path = directory / self.file_name(name) torch.save(pt.tensor.contiguous(), path) size = path.stat().st_size total += size meta = pt.meta() meta["file"] = path.name meta["file_bytes"] = size meta["sha256"] = _sha256_file(path) if hash_files else None tensors[name] = meta manifest = { "format": CACHE_FORMAT, "version": plan.version, "hf_sha": HF_SNAPSHOT_SHAS[plan.version], "embodiment_id": plan.embodiment_id, "dtype_policy": plan.dtype_policy, "transform_version": TRANSFORM_VERSION, "options": asdict(plan.options), "arch": self.arch, "n_tensors": len(tensors), "file_bytes_total": total, "device_bytes_total": byte_summary(plan)["total_bytes"], "tensors": tensors, } with open(directory / CACHE_MANIFEST, "w") as fh: json.dump(manifest, fh, indent=1, sort_keys=True) marker.write_text(CACHE_FORMAT + "\n") return manifest def load_manifest(self, directory: Path) -> Dict[str, Any]: directory = Path(directory) if not self.is_complete(directory): raise FileNotFoundError(f"{directory}: no complete host plan cache") with open(directory / CACHE_MANIFEST) as fh: m = json.load(fh) if m.get("format") != CACHE_FORMAT or m.get("transform_version") != TRANSFORM_VERSION: raise ValueError( f"{directory}: cache format/transform version mismatch ({m.get('format')}, {m.get('transform_version')})" ) return m def load(self, directory: Path, names: Optional[Iterable[str]] = None, verify_hashes: bool = False) -> DevicePlan: """Rebuild the ``DevicePlan`` from a complete cache directory (file sizes always checked; sha256 optional).""" directory = Path(directory) m = self.load_manifest(directory) opts = TransformOptions(**m["options"]) plan = DevicePlan(m["version"], int(m["embodiment_id"]), m["dtype_policy"], opts) wanted = set(m["tensors"]) if names is None else set(names) missing = wanted - set(m["tensors"]) if missing: raise KeyError(f"{directory}: tensors not in the cache manifest: {sorted(missing)[:5]}") for name in sorted(wanted): meta = m["tensors"][name] path = directory / meta["file"] if not path.is_file() or path.stat().st_size != meta["file_bytes"]: raise ValueError(f"{path}: missing or size mismatch -> cold load required") if verify_hashes and meta.get("sha256") and _sha256_file(path) != meta["sha256"]: raise ValueError(f"{path}: sha256 mismatch") t = torch.load(path, map_location="cpu", weights_only=True) if tuple(t.shape) != tuple(meta["logical_shape"]): raise ValueError(f"{name}: cached shape {tuple(t.shape)} != manifest {meta['logical_shape']}") plan[name] = PlanTensor( name=name, tensor=t, dtype=meta["dtype"], layout=meta["layout"], memory=meta["memory"], category=meta["category"], source_keys=tuple(meta["source_keys"]), transform=meta["transform"], logical_shape=tuple(meta["logical_shape"]), ) return plan def load_or_build( self, version: str, embodiment_id: Optional[int] = None, dtype_policy: str = "bf16", options: Optional[TransformOptions] = None, force_convert: bool = False, ) -> DevicePlan: """Warm path: load a complete cache; cold path (or ``force_convert`` / ``GR00T_TT_FORCE_CONVERT=1``): build from the checkpoint and save.""" cfg = get_config(version) e = cfg.canonical_embodiment_id if embodiment_id is None else int(embodiment_id) opts = options if options is not None else TransformOptions() d = self.directory(version, e, dtype_policy, opts) force = force_convert or os.environ.get("GR00T_TT_FORCE_CONVERT", "0") == "1" if not force and self.is_complete(d): return self.load(d) plan = build_device_plan(version, e, dtype_policy, options=opts) self.save(plan, d) return plan def plan_manifest(plan: DevicePlan) -> Dict[str, Any]: """Metadata-only description of a plan (built or enumerated): the WL §8.2 manifest content without file facts.""" items = {} for name in sorted(plan): it = plan[name] if isinstance(it, PlanTensor): items[name] = it.meta() else: items[name] = { "name": it.name, "category": it.category, "dtype": it.dtype, "layout": it.layout, "memory": it.memory, "logical_shape": list(it.shape), "source_keys": list(it.source_keys), "transform": it.transform, "device_bytes": it.nbytes, } return { "version": plan.version, "hf_sha": HF_SNAPSHOT_SHAS[plan.version], "embodiment_id": plan.embodiment_id, "dtype_policy": plan.dtype_policy, "transform_version": TRANSFORM_VERSION, "options": asdict(plan.options), "tensors": items, } __all__ = [ "TRANSFORM_VERSION", "CATEGORIES", "DTYPE_POLICIES", "TIMESTEP_BUCKETS", "Rule", "RULES", "EXPECTED_CATEGORY_COUNTS", "matching_rules", "rule_of", "category_of", "is_skipped", "canonical_name", "categorize", "skipped_keys", "device_bytes", "pad32", "to_in_out", "to_row", "pad_dim", "siglip_conv_to_linear", "siglip_linear_patch_weight", "permute_patch_columns", "qwen3vl_conv3d_to_linear", "patch_embed_device_weight", "unfold_patches", "resize_pos_table_bilinear", "pad_heads_out", "pad_heads_vec", "pad_heads_in", "fuse_projections", "fuse_gate_up", "interleave_gate_up_tiles", "deinterleave_gate_up_tiles", "reverse_permute", "reverse_permute_1d", "rope_cos_sin_hf_to_meta", "unshuffle_column_permutation", "fold_unshuffle_order_into_mlp1", "pixel_unshuffle_tokens", "merger_column_permutation", "permute_input_columns", "qwen3vl_pos_embeds", "qwen3vl_vision_rope", "embed_rows", "tau_sinusoid", "timestep_proj", "timestep_embedding", "adaln_vectors", "fold_tau_bias", "fold_position_bias", "TransformOptions", "PlanEntry", "PlanTensor", "DevicePlan", "enumerate_plan", "build_device_plan", "byte_summary", "assert_byte_totals", "SPEC_TOTALS_MIB", "SPEC_CATEGORY_MIB", "HostPlanCache", "plan_manifest", ]