| |
| """Reference PyTorch implementation of the ShellElastic supernet. |
| |
| This file implements the model of "ShellElastic: Scalable Elastic Pretraining with Simultaneous Maximal Updates and |
| Hyperparameter Transfer" (Jin et al.). It depends only on ``torch``; ``safetensors`` is needed to load released |
| weights. |
| |
| Notation follows Section 3.1 of the paper. |
| |
| * Width. With base width ``g`` and supported set ``S = {s_1 < ... < s_K = M}`` (stored in ``config.widths``, in |
| units of ``g``), the residual stream has ``M * g`` features split into contiguous groups |
| ``G_k = [s_{k-1} g, s_k g)``. ``G_1`` is the dense core, ``G_2 .. G_K`` are the shells. The sub-network of tier |
| ``k`` (Low, Medium, High, Extra, Max) is the core plus the first ``k - 1`` shells, i.e. the first ``s_k g`` |
| features. |
| * Structured projections (Eq. 2). Every hidden matrix is block lower triangular: the block row ``W_k`` that |
| produces group ``G_k`` reads only the first ``s_k g`` input features and has shape ``[|G_k|, s_k g]``. Each |
| matrix is stored by block row: ``<proj>.core.weight`` is ``W_1`` (the rows of the core ``G_1``), and |
| ``<proj>.shell.<j>.weight`` is ``W_{j+1}`` (the rows of the shell ``G_{j+1}``, ``j = 1 .. K-1``), of shape |
| ``[|G_{j+1}|, s_{j+1} g]``. The zero blocks above the diagonal are never stored. |
| * Group-local normalization (Eq. 3). The two pre-norms of every layer are RMSNorms whose statistics are computed |
| separately over the core and over each shell (one gain per feature). |
| * Attention (Eq. 4). Tier ``k`` uses the first ``s_k g / d_head`` heads; heads never straddle a group. Queries and |
| keys get a per-head RMSNorm and rotary position embeddings, and the attention output is multiplied elementwise |
| by ``sigmoid(x W_G^T)`` before the output projection. |
| * FFN (Eq. 5). SwiGLU with hidden width ``e`` times the model width; ``gate_proj`` / ``up_proj`` have block |
| boundaries ``s_k g`` on the input axis and ``e s_k g`` on the output axis, ``down_proj`` the reverse. |
| * Readout (Eq. 6). Every tier has its own terminal RMSNorm over all of its ``s_k g`` features (``norm.<tier>``) and |
| its own readout ``U_s`` (``lm_head.<tier>``). In the released weights the learned logit scale ``alpha_s`` is |
| already multiplied into ``lm_head.<tier>.weight`` and the FFN multiplier ``1/e`` into ``down_proj``, so inference |
| needs no extra scalars. |
| |
| Because no operation lets the features of a tier depend on the groups outside it, one forward pass at width |
| ``M g`` contains the exact hidden states of every smaller tier in its leading features (Proposition 1). |
| ``forward(..., tiers=[...])`` uses this to return several tiers' logits from one trunk pass, and ``extract(tier)`` |
| returns a standalone smaller model with identical outputs. |
| |
| Numerics. This file computes the RMSNorm statistics in float32 and casts the result back to the activation dtype |
| before the gain is applied. ``config.rope_precision`` selects the rotary tables: ``"bfloat16"`` (the default) |
| rounds the positions, the frequencies, their product and the cosine and sine to bfloat16; ``"float32"`` computes |
| them in float32. The released weights are float32; run the model in float32, or cast it to bfloat16 with |
| ``from_pretrained(..., dtype=torch.bfloat16)``. |
| |
| Example:: |
| |
| model = ShellElasticForCausalLM.from_pretrained("path/to/checkpoint", device="cuda") |
| logits = model(input_ids, tier="medium") # [batch, seq, vocab] |
| all_logits = model(input_ids, tiers=model.config.tiers) # dict: tier -> logits |
| logits, loss = model(input_ids[:, :-1], tier="low", labels=input_ids[:, 1:]) |
| low = model.extract("low") # standalone Low sub-network |
| """ |
|
|
| from __future__ import annotations |
|
|
| import json |
| import os |
| from dataclasses import asdict, dataclass, field |
| from typing import Dict, List, Optional, Sequence, Union |
|
|
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
|
|
|
|
| DEFAULT_TIERS = ["low", "medium", "high", "extra", "max"] |
|
|
|
|
| |
| |
| |
|
|
|
|
| @dataclass |
| class ShellElasticConfig: |
| """Hyperparameters of a ShellElastic supernet (or of a sub-network extracted from it). |
| |
| Attributes: |
| vocab_size: Size of the vocabulary. |
| num_hidden_layers: Number of transformer layers ``L``. |
| base_width: Base width ``g``; every group is a multiple of it. |
| widths: Supported set ``S`` in units of ``g``, strictly increasing. The last entry is ``M`` (the full width |
| is ``M * g``). Group ``G_1`` (the core) has ``widths[0] * g`` features, every later group |
| ``(widths[k] - widths[k-1]) * g``. |
| tiers: One name per supported width (Low, Medium, High, Extra, Max). |
| readout_tiers: Tiers that have their own terminal norm and readout in this checkpoint (default: all tiers; |
| a sub-network returned by ``extract`` keeps only its own). |
| head_dim: Attention head dimension ``d_head`` (must divide ``base_width``). |
| mlp_base_width: FFN hidden features per base width, i.e. ``e * g`` with FFN expansion ``e``. |
| rope_theta: Base of the rotary position embedding. |
| rope_precision: Rotary tables, ``"bfloat16"`` (default: position, frequency and their product rounded to |
| bfloat16, and the cosine and sine rounded to bfloat16, in any compute precision) or ``"float32"`` |
| (computed in float32). |
| rms_norm_eps: Epsilon of every RMSNorm. |
| max_position_embeddings: Context length (the sequence length of 2,048 tokens). |
| tie_word_embeddings: Always false (the embedding and the readouts are separate). |
| tokenizer: Tokenizer the token ids refer to. |
| """ |
|
|
| vocab_size: int = 50432 |
| num_hidden_layers: int = 32 |
| base_width: int = 128 |
| widths: List[int] = field(default_factory=lambda: [4, 5, 6, 7, 8]) |
| tiers: List[str] = field(default_factory=lambda: list(DEFAULT_TIERS)) |
| readout_tiers: Optional[List[str]] = None |
| head_dim: int = 128 |
| mlp_base_width: int = 128 |
| rope_theta: float = 10000.0 |
| rope_precision: str = "bfloat16" |
| rms_norm_eps: float = 1e-5 |
| max_position_embeddings: int = 2048 |
| tie_word_embeddings: bool = False |
| tokenizer: str = "EleutherAI/gpt-neox-20b" |
| architectures: List[str] = field(default_factory=lambda: ["ShellElasticForCausalLM"]) |
| model_type: str = "shellelastic" |
|
|
| def __post_init__(self) -> None: |
| self.widths = [int(w) for w in self.widths] |
| self.tiers = [str(t) for t in self.tiers] |
| if self.readout_tiers is None: |
| self.readout_tiers = list(self.tiers) |
| if any(b <= a for a, b in zip([0] + self.widths, self.widths)): |
| raise ValueError(f"widths must be positive and strictly increasing, got {self.widths}") |
| if len(self.tiers) != len(self.widths): |
| raise ValueError("tiers and widths must have the same length") |
| if not self.readout_tiers or any(t not in self.tiers for t in self.readout_tiers): |
| raise ValueError(f"readout_tiers {self.readout_tiers} must be a non-empty subset of tiers {self.tiers}") |
| self.readout_tiers = sorted(set(self.readout_tiers), key=self.tiers.index) |
| if self.base_width % self.head_dim != 0: |
| raise ValueError("head_dim must divide base_width so that heads never straddle a group") |
| if self.rope_precision not in ("bfloat16", "float32"): |
| raise ValueError("rope_precision must be 'bfloat16' or 'float32'") |
| if self.tie_word_embeddings: |
| raise ValueError("ShellElastic models do not tie the embedding and the readouts") |
|
|
| |
|
|
| @property |
| def hidden_size(self) -> int: |
| """Full model width ``M * g``.""" |
| return self.widths[-1] * self.base_width |
|
|
| @property |
| def boundaries(self) -> List[int]: |
| """Right edges of the residual groups, ``[s_1 g, ..., s_K g]``.""" |
| return [s * self.base_width for s in self.widths] |
|
|
| @property |
| def mlp_boundaries(self) -> List[int]: |
| """Right edges of the FFN hidden groups, ``[e s_1 g, ..., e s_K g]``.""" |
| return [s * self.mlp_base_width for s in self.widths] |
|
|
| @property |
| def group_sizes(self) -> List[int]: |
| """Sizes of the core and of every shell.""" |
| b = self.boundaries |
| return [b[0]] + [hi - lo for lo, hi in zip(b, b[1:])] |
|
|
| def num_groups(self, tier: str) -> int: |
| """Number of groups (core + shells) used by ``tier``.""" |
| return self.tiers.index(tier) + 1 |
|
|
| def width(self, tier: str) -> int: |
| """Model width of ``tier``.""" |
| return self.boundaries[self.tiers.index(tier)] |
|
|
| def for_tier(self, tier: str) -> "ShellElasticConfig": |
| """Configuration of the standalone sub-network ``tier`` (see ``extract``).""" |
| k = self.num_groups(tier) |
| values = asdict(self) |
| values.update(widths=self.widths[:k], tiers=self.tiers[:k], readout_tiers=[tier]) |
| return ShellElasticConfig(**values) |
|
|
| |
|
|
| def to_dict(self) -> dict: |
| return asdict(self) |
|
|
| @classmethod |
| def from_dict(cls, values: dict) -> "ShellElasticConfig": |
| known = {f for f in cls.__dataclass_fields__} |
| return cls(**{k: v for k, v in values.items() if k in known}) |
|
|
| @classmethod |
| def from_json_file(cls, path: str) -> "ShellElasticConfig": |
| with open(path, "r", encoding="utf-8") as f: |
| return cls.from_dict(json.load(f)) |
|
|
|
|
| |
| |
| |
|
|
|
|
| def rms_norm(x: torch.Tensor, weight: torch.Tensor, eps: float, |
| group_sizes: Optional[Sequence[int]] = None) -> torch.Tensor: |
| """RMSNorm over the last dimension, with separate statistics per feature group (Eq. 3). |
| |
| The statistics are computed in float32 and the result is cast back to the input dtype before the gain is |
| applied. ``group_sizes=None`` normalizes over the whole last dimension. |
| """ |
| xf = x.float() |
| chunks = xf.split(list(group_sizes), dim=-1) if group_sizes is not None else (xf,) |
| normed = [c * torch.rsqrt(c.pow(2).mean(dim=-1, keepdim=True) + eps) for c in chunks] |
| return torch.cat(normed, dim=-1).to(x.dtype) * weight |
|
|
|
|
| class RMSNorm(nn.Module): |
| """RMSNorm with one gain per feature. ``group_sizes`` selects group-local statistics. |
| |
| The gain has the full width of the module; a narrower input uses its leading entries. |
| """ |
|
|
| def __init__(self, width: int, eps: float) -> None: |
| super().__init__() |
| self.eps = eps |
| self.weight = nn.Parameter(torch.ones(width)) |
|
|
| def forward(self, x: torch.Tensor, group_sizes: Optional[Sequence[int]] = None) -> torch.Tensor: |
| return rms_norm(x, self.weight[: x.shape[-1]], self.eps, group_sizes) |
|
|
|
|
| class ShellLinear(nn.Module): |
| """Block-lower-triangular linear map, stored as one weight per output group (block row). |
| |
| ``core.weight`` (the block row of the core ``G_1``) has shape ``[out_bounds[0], in_bounds[0]]``; |
| ``shell[str(j)].weight`` (the block row of the shell ``G_{j+1}``, ``j = 1 .. K-1``) has shape |
| ``[out_bounds[j] - out_bounds[j-1], in_bounds[j]]``: its rows read only the first ``in_bounds[j]`` input |
| features. There is no bias. |
| """ |
|
|
| def __init__(self, in_bounds: Sequence[int], out_bounds: Sequence[int]) -> None: |
| super().__init__() |
| self.in_bounds = list(in_bounds) |
| self.out_bounds = list(out_bounds) |
| self.core = nn.Linear(self.in_bounds[0], self.out_bounds[0], bias=False) |
| self.shell = nn.ModuleDict({ |
| str(k): nn.Linear(self.in_bounds[k], self.out_bounds[k] - self.out_bounds[k - 1], bias=False) |
| for k in range(1, len(self.out_bounds)) |
| }) |
|
|
| def block_rows(self, num_groups: int) -> List[torch.Tensor]: |
| """Weights of the first ``num_groups`` block rows (core first).""" |
| return [self.core.weight] + [self.shell[str(k)].weight for k in range(1, num_groups)] |
|
|
| def forward(self, x: torch.Tensor, num_groups: int) -> torch.Tensor: |
| |
| return torch.cat([F.linear(x[..., : w.shape[1]], w) for w in self.block_rows(num_groups)], dim=-1) |
|
|
| def dense_weight(self, num_groups: int) -> torch.Tensor: |
| """The equivalent dense ``[out, in]`` matrix of the first ``num_groups`` groups, with the blocks above the |
| diagonal filled with zeros.""" |
| n_in = self.in_bounds[num_groups - 1] |
| return torch.cat([F.pad(w, (0, n_in - w.shape[1])) for w in self.block_rows(num_groups)], dim=0) |
|
|
|
|
| def rotary_cos_sin(length: int, head_dim: int, theta: float, precision: str, |
| device: torch.device, dtype: torch.dtype): |
| """Cosine and sine tables of shape ``[length, head_dim]`` for rotate-half RoPE.""" |
| frequency = 1.0 / (theta ** (torch.arange(0, head_dim, 2).float() / head_dim)) |
| position = torch.arange(length, dtype=torch.float32) |
| frequency, position = frequency.to(device), position.to(device) |
| if precision == "bfloat16": |
| |
| |
| angles = torch.outer(position.bfloat16().float(), frequency.bfloat16().float()).bfloat16().float() |
| angles = torch.cat([angles, angles], dim=-1) |
| return angles.cos().bfloat16().to(dtype), angles.sin().bfloat16().to(dtype) |
| angles = torch.outer(position, frequency) |
| angles = torch.cat([angles, angles], dim=-1) |
| return angles.cos().to(dtype), angles.sin().to(dtype) |
|
|
|
|
| def apply_rotary(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor: |
| """Rotate-half RoPE for ``x`` of shape ``[batch, seq, heads, head_dim]``.""" |
| x1, x2 = x.chunk(2, dim=-1) |
| rotated = torch.cat([-x2, x1], dim=-1) |
| return x * cos[None, :, None, :] + rotated * sin[None, :, None, :] |
|
|
|
|
| class ShellAttention(nn.Module): |
| """Causal self-attention with nested heads, QK-norm, RoPE and an elementwise sigmoid gate.""" |
|
|
| def __init__(self, config: ShellElasticConfig) -> None: |
| super().__init__() |
| b, width = config.boundaries, config.hidden_size |
| self.head_dim = config.head_dim |
| self.q_proj = ShellLinear(b, b) |
| self.k_proj = ShellLinear(b, b) |
| self.v_proj = ShellLinear(b, b) |
| self.gate_proj = ShellLinear(b, b) |
| self.o_proj = ShellLinear(b, b) |
| |
| self.q_norm = RMSNorm(width, config.rms_norm_eps) |
| self.k_norm = RMSNorm(width, config.rms_norm_eps) |
|
|
| def _per_head_norm(self, norm: RMSNorm, x: torch.Tensor) -> torch.Tensor: |
| heads = x.shape[-2] |
| gain = norm.weight[: heads * self.head_dim].view(heads, self.head_dim) |
| return rms_norm(x, gain, norm.eps) |
|
|
| def forward(self, x: torch.Tensor, num_groups: int, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor: |
| batch, seq, width = x.shape |
| heads = width // self.head_dim |
| shape = (batch, seq, heads, self.head_dim) |
| q = self.q_proj(x, num_groups).view(shape) |
| k = self.k_proj(x, num_groups).view(shape) |
| v = self.v_proj(x, num_groups).view(shape) |
| q = apply_rotary(self._per_head_norm(self.q_norm, q), cos, sin) |
| k = apply_rotary(self._per_head_norm(self.k_norm, k), cos, sin) |
| out = F.scaled_dot_product_attention( |
| q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2), |
| is_causal=True, scale=self.head_dim ** -0.5, |
| ).transpose(1, 2) |
| out = out * torch.sigmoid(self.gate_proj(x, num_groups).view(shape)) |
| return self.o_proj(out.reshape(batch, seq, width), num_groups) |
|
|
|
|
| class ShellMLP(nn.Module): |
| """SwiGLU feed-forward network with block-lower-triangular projections.""" |
|
|
| def __init__(self, config: ShellElasticConfig) -> None: |
| super().__init__() |
| b, h = config.boundaries, config.mlp_boundaries |
| self.gate_proj = ShellLinear(b, h) |
| self.up_proj = ShellLinear(b, h) |
| self.down_proj = ShellLinear(h, b) |
|
|
| def forward(self, x: torch.Tensor, num_groups: int) -> torch.Tensor: |
| hidden = F.silu(self.gate_proj(x, num_groups)) * self.up_proj(x, num_groups) |
| return self.down_proj(hidden, num_groups) |
|
|
|
|
| class ShellElasticLayer(nn.Module): |
| """Pre-norm transformer layer with group-local branch normalization.""" |
|
|
| def __init__(self, config: ShellElasticConfig) -> None: |
| super().__init__() |
| self.group_sizes = config.group_sizes |
| self.attn_norm = RMSNorm(config.hidden_size, config.rms_norm_eps) |
| self.attn = ShellAttention(config) |
| self.mlp_norm = RMSNorm(config.hidden_size, config.rms_norm_eps) |
| self.mlp = ShellMLP(config) |
|
|
| def forward(self, x: torch.Tensor, num_groups: int, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor: |
| groups = self.group_sizes[:num_groups] |
| x = x + self.attn(self.attn_norm(x, groups), num_groups, cos, sin) |
| x = x + self.mlp(self.mlp_norm(x, groups), num_groups) |
| return x |
|
|
|
|
| |
| |
| |
|
|
|
|
| class ShellElasticForCausalLM(nn.Module): |
| """ShellElastic supernet: a shared trunk plus one terminal norm and readout per tier.""" |
|
|
| def __init__(self, config: ShellElasticConfig) -> None: |
| super().__init__() |
| self.config = config |
| self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size) |
| self.layers = nn.ModuleList([ShellElasticLayer(config) for _ in range(config.num_hidden_layers)]) |
| self.norm = nn.ModuleDict({t: RMSNorm(config.width(t), config.rms_norm_eps) for t in config.readout_tiers}) |
| self.lm_head = nn.ModuleDict({ |
| t: nn.Linear(config.width(t), config.vocab_size, bias=False) for t in config.readout_tiers |
| }) |
|
|
| |
|
|
| @classmethod |
| def from_pretrained(cls, path: str, device: Union[str, torch.device] = "cpu", |
| dtype: Optional[torch.dtype] = None) -> "ShellElasticForCausalLM": |
| """Load ``config.json`` and ``model.safetensors`` from the directory ``path``; ``dtype`` optionally casts |
| the (float32) weights, e.g. to ``torch.bfloat16``. The weights are read and cast on the CPU and then moved |
| to ``device``, so the device holds only one copy of the model, in ``dtype``.""" |
| from safetensors.torch import load_file |
|
|
| config = ShellElasticConfig.from_json_file(os.path.join(path, "config.json")) |
| with torch.device("meta"): |
| model = cls(config) |
| state = load_file(os.path.join(path, "model.safetensors"), device="cpu") |
| if dtype is not None: |
| for name in list(state): |
| if state[name].is_floating_point(): |
| state[name] = state[name].to(dtype) |
| model.load_state_dict(state, strict=True, assign=True) |
| return model.to(device).eval() |
|
|
| def save_pretrained(self, path: str) -> None: |
| """Write ``config.json`` and ``model.safetensors`` to the directory ``path`` (for example, a sub-network |
| returned by ``extract``); ``from_pretrained`` reads them back.""" |
| from safetensors.torch import save_file |
|
|
| os.makedirs(path, exist_ok=True) |
| values = self.config.to_dict() |
| values["torch_dtype"] = str(self.embed_tokens.weight.dtype).replace("torch.", "") |
| with open(os.path.join(path, "config.json"), "w", encoding="utf-8") as f: |
| json.dump(values, f, indent=2) |
| f.write("\n") |
| state = {name: t.detach().contiguous() for name, t in self.state_dict().items()} |
| save_file(state, os.path.join(path, "model.safetensors"), metadata={"format": "pt"}) |
|
|
| |
|
|
| def hidden_states(self, input_ids: torch.Tensor, tier: Optional[str] = None) -> torch.Tensor: |
| """Trunk output (before the terminal norm) of width ``config.width(tier)``. |
| |
| Its first ``config.width(t)`` features are the hidden states of every smaller tier ``t``. |
| """ |
| tier = tier or self.config.tiers[-1] |
| num_groups = self.config.num_groups(tier) |
| width = self.config.width(tier) |
| seq = input_ids.shape[1] |
| if seq > self.config.max_position_embeddings: |
| raise ValueError(f"{seq} tokens exceed the context of {self.config.max_position_embeddings}") |
| x = self.embed_tokens(input_ids)[..., :width] |
| cos, sin = rotary_cos_sin(seq, self.config.head_dim, self.config.rope_theta, |
| self.config.rope_precision, x.device, x.dtype) |
| for layer in self.layers: |
| x = layer(x, num_groups, cos, sin) |
| return x |
|
|
| def logits(self, hidden: torch.Tensor, tier: str) -> torch.Tensor: |
| """Terminal norm and readout of ``tier`` (Eq. 6) applied to a trunk output of width >= its own.""" |
| if tier not in self.lm_head: |
| raise KeyError(f"tier {tier!r} has no readout in this model (available: {self.config.readout_tiers})") |
| h = hidden[..., : self.config.width(tier)] |
| return self.lm_head[tier](self.norm[tier](h)) |
|
|
| def forward(self, input_ids: torch.Tensor, tier: Optional[str] = None, tiers: Optional[Sequence[str]] = None, |
| labels: Optional[torch.Tensor] = None): |
| """Next-token logits (in the model's dtype). |
| |
| Args: |
| input_ids: ``[batch, seq]`` token ids (no padding; causal attention over the sequence). |
| tier: Tier to evaluate (default: the widest tier with a readout). |
| tiers: If given, evaluate all of these tiers with one trunk pass at the widest of them and return a |
| ``{tier: logits}`` dict. |
| labels: Optional ``[batch, seq]`` targets (-100 is ignored). The mean cross-entropy (computed in float32) |
| is returned as well: ``(logits, loss)``, or ``({tier: logits}, {tier: loss})`` with ``tiers``. |
| """ |
| if tiers is None: |
| tier = tier or self.config.readout_tiers[-1] |
| logits = self.logits(self.hidden_states(input_ids, tier), tier) |
| return logits if labels is None else (logits, self.loss(logits, labels)) |
| widest = max(tiers, key=self.config.tiers.index) |
| hidden = self.hidden_states(input_ids, widest) |
| out = {t: self.logits(hidden, t) for t in tiers} |
| return out if labels is None else (out, {t: self.loss(v, labels) for t, v in out.items()}) |
|
|
| @staticmethod |
| def loss(logits: torch.Tensor, labels: torch.Tensor) -> torch.Tensor: |
| """Mean next-token cross-entropy in float32 (label -100 is ignored).""" |
| return F.cross_entropy(logits.float().flatten(0, 1), labels.flatten(), ignore_index=-100) |
|
|
| @torch.no_grad() |
| def generate(self, input_ids: torch.Tensor, max_new_tokens: int, tier: Optional[str] = None, |
| temperature: float = 0.0, top_k: Optional[int] = None) -> torch.Tensor: |
| """Simple sampling loop without a key/value cache (greedy when ``temperature == 0``).""" |
| for _ in range(max_new_tokens): |
| context = input_ids[:, -self.config.max_position_embeddings:] |
| logits = self(context, tier=tier)[:, -1].float() |
| if temperature <= 0: |
| next_ids = logits.argmax(dim=-1, keepdim=True) |
| else: |
| logits = logits / temperature |
| if top_k is not None: |
| kth = torch.topk(logits, top_k, dim=-1).values[:, -1:] |
| logits = logits.masked_fill(logits < kth, float("-inf")) |
| next_ids = torch.multinomial(torch.softmax(logits, dim=-1), num_samples=1) |
| input_ids = torch.cat([input_ids, next_ids], dim=1) |
| return input_ids |
|
|
| |
|
|
| def subnetwork_state_dict(self, tier: str) -> Dict[str, torch.Tensor]: |
| """Weights of the standalone sub-network ``tier``: the embedding columns, the core and the shells it |
| contains, the leading entries of every norm gain, and its own terminal norm and readout. Loads into |
| ``ShellElasticForCausalLM(config.for_tier(tier))``.""" |
| if tier not in self.lm_head: |
| raise KeyError(f"tier {tier!r} has no readout in this model (available: {self.config.readout_tiers})") |
| k, width = self.config.num_groups(tier), self.config.width(tier) |
| out = {} |
| for name, tensor in self.state_dict().items(): |
| if name.startswith(("norm.", "lm_head.")): |
| if name.split(".")[1] == tier: |
| out[name] = tensor |
| elif name == "embed_tokens.weight": |
| out[name] = tensor[:, :width] |
| elif ".shell." in name: |
| if int(name.split(".shell.")[1].split(".")[0]) < k: |
| out[name] = tensor |
| elif name.endswith("norm.weight"): |
| out[name] = tensor[:width] |
| else: |
| out[name] = tensor |
| return out |
|
|
| def extract(self, tier: str) -> "ShellElasticForCausalLM": |
| """A standalone copy of the sub-network ``tier``; its logits equal ``self(x, tier=tier)``.""" |
| with torch.device("meta"): |
| model = ShellElasticForCausalLM(self.config.for_tier(tier)) |
| state = {n: t.clone() for n, t in self.subnetwork_state_dict(tier).items()} |
| model.load_state_dict(state, strict=True, assign=True) |
| return model.train(self.training) |
|
|
| def dense_state_dict(self, tier: str) -> Dict[str, torch.Tensor]: |
| """Returns the weights of the sub-network ``tier`` as a state dict (nothing is written to disk), with every |
| structured projection as one dense matrix (``...<proj>.weight``, zeros above the block diagonal). The |
| group-local norms are still required to run it; this format is provided for inspection and porting.""" |
| k = self.config.num_groups(tier) |
| out = {} |
| for name, tensor in self.subnetwork_state_dict(tier).items(): |
| if ".core." in name or ".shell." in name: |
| continue |
| out[name] = tensor |
| for name, module in self.named_modules(): |
| if isinstance(module, ShellLinear): |
| out[f"{name}.weight"] = module.dense_weight(k) |
| return out |
|
|
| def num_parameters(self, tier: Optional[str] = None) -> int: |
| """Parameters stored for the whole model, or used by ``tier`` (non-zero weights it reads, its embedding |
| columns, and its own terminal norm and readout).""" |
| tensors = self.state_dict() if tier is None else self.subnetwork_state_dict(tier) |
| return sum(t.numel() for t in tensors.values()) |
|
|
|
|
| __all__ = [ |
| "ShellElasticConfig", |
| "ShellElasticForCausalLM", |
| "ShellLinear", |
| "RMSNorm", |
| "rms_norm", |
| "rotary_cos_sin", |
| ] |
|
|
