Spaces:
Paused
Paused
| """M10 -- Mixture-of-Depths (MoD): a per-block token router for fewer FLOPs/token. | |
| External basis: Raposo et al., "Mixture-of-Depths: Dynamically allocating compute in | |
| transformer-based language models" (DeepMind, 2024). | |
| Idea: not every token needs every block. A tiny fp router scores each token; the top | |
| `capacity` fraction (e.g. 0.5) are passed through the wrapped block (attention+FFN); | |
| the remaining tokens BYPASS the block entirely via the residual -- their compute is | |
| skipped. The compute budget per block is therefore fixed at `capacity * tokens`. | |
| So the router learns *which* tokens matter, we multiply the block's contribution for the | |
| selected tokens by the router's sigmoid weight (a weighted residual / straight-through | |
| style). Skipped tokens are returned bit-identically (pure residual identity). | |
| CAUSAL CAVEAT (important): the standard MoD top-k selection is a NON-CAUSAL operation -- | |
| it ranks tokens within a fixed window, so token i's keep/skip decision can depend on | |
| later tokens. For training a witness on a fixed block this is fine (and is exactly the | |
| "expert-choice"-style routing DeepMind use during training). At inference / for a strictly | |
| causal model you would instead train a per-token predictor of the routing decision (the | |
| "router predictor" in the paper) so no future information leaks. We do top-k over the | |
| batch*time tokens here for the witness; this is acceptable for measuring the FLOP/quality | |
| tradeoff but is NOT a drop-in causal inference path. | |
| Pure PyTorch, CPU/CUDA. | |
| """ | |
| from __future__ import annotations | |
| import torch | |
| import torch.nn as nn | |
| __all__ = ["MoDBlock"] | |
| class MoDBlock(nn.Module): | |
| """Wrap a block (any ``nn.Module`` mapping ``(B,T,D) -> (B,T,D)``) with token routing. | |
| Args: | |
| block: the inner block to (conditionally) apply. It MUST already be a residual block, | |
| i.e. ``block(x)`` returns ``x + delta`` (as the TinyGPT / models.py Block does). | |
| MoD then mixes ``block(x_sel)`` for selected tokens and passes the rest through | |
| unchanged. | |
| capacity: fraction in (0, 1] of tokens to route through the block. ``1.0`` means every | |
| token is processed (exact parity with the unwrapped block, modulo the router weight, | |
| see ``use_router_weight``). | |
| use_router_weight: if True (default, the MoD recipe), the block's *contribution* for a | |
| selected token is scaled by the router's sigmoid weight, so the router gets a | |
| gradient signal telling it which tokens benefit from the block. At capacity=1.0 this | |
| is disabled automatically so the wrapper is an exact identity over the block. | |
| The realized fraction of tokens actually processed is exposed via ``last_processed_fraction`` | |
| after each forward (and via ``realized_fraction()``). | |
| """ | |
| def __init__(self, block: nn.Module, capacity: float = 0.5, | |
| use_router_weight: bool = True) -> None: | |
| super().__init__() | |
| if not (0.0 < capacity <= 1.0): | |
| raise ValueError(f"capacity must be in (0, 1], got {capacity}") | |
| self.block = block | |
| self.capacity = float(capacity) | |
| self.use_router_weight = bool(use_router_weight) | |
| self._d = self._infer_dim(block) | |
| # tiny fp router: one score per token. AdamW owns this Parameter. | |
| self.router = nn.Linear(self._d, 1) | |
| nn.init.zeros_(self.router.bias) | |
| nn.init.normal_(self.router.weight, std=0.02) | |
| self.last_processed_fraction: float = 0.0 | |
| def _infer_dim(block: nn.Module) -> int: | |
| for attr in ("n_embd", "d", "dim"): | |
| if hasattr(block, attr) and isinstance(getattr(block, attr), int): | |
| return getattr(block, attr) | |
| # fall back to the first LayerNorm / Linear feature size we can find | |
| for m in block.modules(): | |
| if isinstance(m, nn.LayerNorm): | |
| return m.normalized_shape[0] | |
| if isinstance(m, nn.Linear): | |
| return m.in_features | |
| raise ValueError("could not infer block embedding dim; pass a block exposing .n_embd/.d") | |
| def realized_fraction(self) -> float: | |
| """Fraction of tokens actually pushed through the block on the last forward.""" | |
| return self.last_processed_fraction | |
| def forward(self, x: torch.Tensor) -> torch.Tensor: | |
| b, t, d = x.shape | |
| n = b * t | |
| flat = x.reshape(n, d) | |
| # capacity=1.0 -> process everything; this is EXACT parity with self.block(x). | |
| if self.capacity >= 1.0: | |
| self.last_processed_fraction = 1.0 | |
| return self.block(x) | |
| k = max(1, int(round(self.capacity * n))) | |
| k = min(k, n) | |
| scores = self.router(flat).squeeze(-1) # (n,) | |
| # top-k tokens by router score get the block; the rest bypass via the residual. | |
| top_val, top_idx = torch.topk(scores, k, sorted=False) | |
| self.last_processed_fraction = k / n | |
| sel = flat.index_select(0, top_idx).view(1, k, d) # (1,k,d) keeps block's (B,T,D) API | |
| block_out = self.block(sel).view(k, d) # = sel + delta_sel (residual block) | |
| if self.use_router_weight: | |
| # weighted residual: scale the block's *delta* by the router's sigmoid weight so the | |
| # router receives gradient. out = sel + w * (block_out - sel) = sel + w*delta. | |
| w = torch.sigmoid(top_val).unsqueeze(-1) # (k,1) | |
| new_sel = sel.view(k, d) + w * (block_out - sel.view(k, d)) | |
| else: | |
| new_sel = block_out | |
| # start from the pure residual (every token passed through unchanged) ... | |
| out = flat.clone() | |
| # ... then overwrite the selected rows with the block's (weighted) output. | |
| out = out.index_copy(0, top_idx, new_sel) | |
| return out.view(b, t, d) | |
| def block_flop_fraction(self) -> float: | |
| """The fraction of the unwrapped block's FLOPs this wrapper spends (== capacity). | |
| The router adds a tiny d-per-token term on top, negligible vs the block's ~d^2. | |
| """ | |
| return self.capacity | |