← Walkthroughmapa/models/towers.py
Source snapshot · source SHA-256 b5950e0090f0
1"""Twelve transformer blocks, each preserving token count and model width.23Main path: Encoder.forward_flat().45Input: flat tokens (B,total,d) and their metadata.6Output: requested width-d layer features, plus the normalized multi-layer output.7Next: attention.L1Block explains one block; encoder.py reshapes the returned taps.89Region identity is added once at the stack input. Attention stays within each10array/clip block. Blocks share the architecture and RoPE coordinates, not weights.1112With deep_sup enabled, the main output concatenates separately normalized taps133/6/9/12: width d -> 4d. Requested tap_blocks are returned separately as raw14width-d features; MapaEncoder.forward uses those for inference.15"""1617from __future__ import annotations1819import math20from typing import TYPE_CHECKING2122import torch23from torch import Tensor, nn2425if TYPE_CHECKING:26 from mapa.models.pack_r4 import R4Grid2728from mapa.models.attention import LN_EPS, L1Block29from mapa.models.pe import RegionIdentityEmbed, init_transformer_weights3031DEPTH = 12 # blocks; not a knob32HEAD_DIM = 64 # held across widths, the ViT sizing convention33ENC_D_MODEL = 256 # the width the ablation ladder was tuned at34VIT_SMALL = 384 # the released width35ENC_SUP_TAPS: tuple[int, ...] = (3, 6, 9, 12)36N_LEVELS = len(ENC_SUP_TAPS)373839def n_heads_for(d_model: int) -> int:40 """Head count for one width. Off-ladder widths fail here, not at the first matmul."""41 if d_model <= 0 or d_model % HEAD_DIM != 0:42 raise ValueError(f"d_model={d_model} must be a positive multiple of {HEAD_DIM}")43 return d_model // HEAD_DIM444546class Encoder(nn.Module):47 """Pre-norm stack of within-array attention blocks."""4849 def __init__(50 self,51 *,52 d_model: int,53 n_heads: int,54 n_regions: int,55 deep_sup: bool = True,56 sup_taps: tuple[int, ...] = ENC_SUP_TAPS,57 region_embed: bool = True,58 space_rope: bool = True,59 ) -> None:60 super().__init__()61 self.region_embed = RegionIdentityEmbed(n_regions, d_model, enabled=region_embed)62 self.blocks = nn.ModuleList(63 [L1Block(d_model, n_heads, space_rope=space_rope) for _ in range(DEPTH)]64 )65 self.deep_sup = bool(deep_sup)66 self.sup_taps = tuple(int(b) for b in sup_taps) if self.deep_sup else ()67 if self.deep_sup:68 if not self.sup_taps:69 raise ValueError("deep_sup=True requires a non-empty sup_taps")70 # Each tap carries its own affine LayerNorm and the four are concatenated. There71 # is no separate terminal norm: the deepest tap's norm is the terminal norm.72 self.norms_block = nn.ModuleList(73 [nn.LayerNorm(d_model, eps=LN_EPS) for _ in self.sup_taps]74 )75 self.norm_out = None76 else:77 self.norms_block = None78 self.norm_out = nn.LayerNorm(d_model, eps=LN_EPS)79 # V-JEPA 2 init: Linear trunc_normal(0.02) with zero bias, LayerNorm weight 1 bias 0,80 # then the depth-scaled residual rescale below. The region embedding is an81 # nn.Embedding, which this rule does not match, so it keeps its own near-zero init.82 self.apply(init_transformer_weights)83 self._rescale_blocks()8485 def _rescale_blocks(self) -> None:86 # Divide the attention output projection and the second MLP layer by sqrt(2*layer),87 # one-indexed, so residual branch variance stays flat with depth.88 for layer_id, block in enumerate(self.blocks):89 scale = math.sqrt(2.0 * (layer_id + 1))90 block.out.weight.data.div_(scale)91 block.mlp.fc2.weight.data.div_(scale)9293 def forward_flat(94 self,95 x_flat: Tensor,96 grid: R4Grid,97 region_packed: Tensor,98 *,99 tap_blocks: tuple[int, ...] = (),100 ) -> Tensor | tuple[Tensor, dict[int, Tensor]]:101 """Run the stack over ``(B, total, d)`` tokens laid out by the pack plan.102103 The grid coordinates are the same for every window in the batch, so they are tiled104 rather than rebuilt. Window ``b`` occupies its own set of array blocks, offset by105 ``b * total``, which makes attention block-diagonal per window and per array without106 the kernel knowing about windows at all.107 """108 b_size, m, d = x_flat.shape109 if m != grid.total:110 raise ValueError(f"x_flat has {m} tokens, grid.total={grid.total}")111 device = x_flat.device112113 # 1. Add the atlas-region embedding to each token.114 region_bm = region_packed[None, :].expand(b_size, m)115 x_flat = x_flat + self.region_embed(region_bm).to(x_flat.dtype)116117 # 2. Give every clip its own array blocks before flattening the batch.118 # E.g. [0,1040,2080] per clip -> [0,1040,2080,3120,4160] for two clips.119 offsets = torch.arange(b_size, device=device, dtype=torch.int64) * m120 array_starts = grid.cu_seqlens[:-1].to(torch.int64)121 batch_starts = array_starts[None, :] + offsets[:, None] # (B,n_arrays)122 final_end = torch.tensor([b_size * m], dtype=torch.int64, device=device)123 cu_b = torch.cat([batch_starts.reshape(-1), final_end])124 keep = torch.ones_like(cu_b, dtype=torch.bool)125 keep[1:] = cu_b[1:] != cu_b[:-1] # drop empty arrays and boundary coincidences126 cu_drop_b = cu_b[keep]127 cu_b = cu_b.to(grid.cu_seqlens.dtype)128129 depth_b = grid.depth[None, :].expand(b_size, m).reshape(b_size * m)130 time_b = grid.time_pos[None, :].expand(b_size, m).reshape(b_size * m)131 # 3. Build contact/time rotations once; reuse them in all 12 blocks.132 rope_cs = self.blocks[0].rope.cos_sin(depth_b, time_b)133134 # 4. Flatten to (B*total,d). Each block updates the same token positions.135 # Read attention.L1Block.forward_flat() for one block, then return here.136 xf = x_flat.reshape(b_size * m, d)137 taps: dict[int, Tensor] = {}138 levels: list[Tensor] = []139 for i, blk in enumerate(self.blocks):140 xf = blk.forward_flat(141 xf, depth_b, time_b, cu_b, cu_drop_b, grid.max_seqlen, rope_cs=rope_cs,142 )143 # Inference taps keep width d; the main multi-level output has width 4d.144 if (i + 1) in tap_blocks:145 taps[i + 1] = xf.reshape(b_size, m, d)146 if self.deep_sup and (i + 1) in self.sup_taps:147 level_index = self.sup_taps.index(i + 1)148 normalized_level = self.norms_block[level_index](xf)149 levels.append(normalized_level)150 # Inference uses raw taps (d each); pretraining used normalized taps (4*d).151 if self.deep_sup:152 out = torch.cat(levels, dim=-1) # (B*total,4*d)153 else:154 out = self.norm_out(xf) # (B*total,d)155 out = out.reshape(b_size, m, out.shape[-1])156 if tap_blocks:157 return out, taps158 return out159160161def build_encoder(162 *, n_regions: int, deep_sup: bool = True, region_embed: bool = True,163 space_rope: bool = True, d_model: int = ENC_D_MODEL,164) -> Encoder:165 return Encoder(166 d_model=d_model, n_heads=n_heads_for(d_model), n_regions=n_regions,167 deep_sup=deep_sup, sup_taps=ENC_SUP_TAPS if deep_sup else (),168 region_embed=region_embed, space_rope=space_rope,169 )