| """Reference implementation of the dense decoder-only transformer behind the Standalone and Sampled standalone |
| models of "ShellElastic: Scalable Elastic Pretraining with Simultaneous Maximal Updates and Hyperparameter |
| Transfer". |
| |
| Each released model is an ordinary dense decoder. The released weights already include the constant forward |
| multipliers of the model (for these models, the FFN output multiplier 1/e, where e is the FFN expansion), so this |
| file contains no scaling constants. |
| |
| Architecture (one layer, pre-norm): |
| |
| h = x + Attention(RMSNorm(x)) |
| out = h + SwiGLU(RMSNorm(h)) |
| |
| * RMSNorm with eps 1e-5; this file computes the statistics in float32 and casts the result back to the input dtype |
| before the learned gain is applied. |
| * Attention: separate query/key/value projections without biases; queries and keys are RMS-normalized within |
| each head (the gains are stored at full width, one per channel); rotary position embeddings in the |
| rotate-half form with base 10,000; causal softmax attention with scale 1/sqrt(head_dim); the attention output |
| is multiplied elementwise by sigmoid(x W_gate^T) before the output projection. |
| * SwiGLU feed-forward network: down_proj(silu(gate_proj(x)) * up_proj(x)). |
| * Untied input embedding and output head, no biases anywhere, context length 2,048, and the 50,432-token GPT-NeoX |
| vocabulary (Table 8 of the paper; token ids of the GPT-NeoX-20B tokenizer, Section 4.1). |
| |
| Rotary tables: ``rope_precision`` is ``"bfloat16"`` (the default) or ``"float32"``. ``"bfloat16"`` rounds the |
| positions and inverse frequencies to bfloat16, their product to bfloat16, and cos/sin of those angles to bfloat16, |
| in every compute precision; ``"float32"`` computes the tables in float32. |
| |
| Dependencies: torch, and safetensors for loading the weights. |
| |
| Example: |
| |
| import torch |
| from modeling_dense import DenseDecoder |
| |
| model = DenseDecoder.from_pretrained("path/to/model_dir", device="cuda", dtype=torch.bfloat16) |
| input_ids = torch.tensor([[...]], device="cuda") # GPT-NeoX-20B token ids |
| logits = model(input_ids) # [batch, tokens, 50432], float32 |
| logits, loss = model(input_ids[:, :-1], labels=input_ids[:, 1:]) |
| """ |
|
|
| from __future__ import annotations |
|
|
| import json |
| import os |
| from dataclasses import asdict, dataclass, fields |
|
|
| import torch |
| import torch.nn.functional as F |
| from torch import nn |
|
|
| __all__ = ["DenseConfig", "DenseDecoder"] |
|
|
|
|
| @dataclass |
| class DenseConfig: |
| vocab_size: int = 50432 |
| hidden_size: int = 384 |
| intermediate_size: int = 1536 |
| num_hidden_layers: int = 32 |
| num_attention_heads: int = 6 |
| head_dim: int = 64 |
| max_position_embeddings: int = 2048 |
| rope_theta: float = 10000.0 |
| rope_precision: str = "bfloat16" |
| rms_norm_eps: float = 1e-5 |
|
|
| @classmethod |
| def from_dict(cls, values: dict) -> "DenseConfig": |
| """Build a config from a dict, ignoring keys that only describe the release (method, tier, ...).""" |
| known = {f.name for f in fields(cls)} |
| return cls(**{k: v for k, v in values.items() if k in known}) |
|
|
| @classmethod |
| def from_json_file(cls, path: str) -> "DenseConfig": |
| with open(path) as f: |
| return cls.from_dict(json.load(f)) |
|
|
| def to_dict(self) -> dict: |
| return asdict(self) |
|
|
|
|
| class RMSNorm(nn.Module): |
| """RMSNorm over the last dimension, or within consecutive groups of ``group_size`` channels. |
| |
| The normalization runs in float32; the result is cast back to the input dtype and then multiplied by the |
| per-channel gain (one gain per channel, also in grouped mode). |
| """ |
|
|
| def __init__(self, dim: int, eps: float, group_size: int | None = None) -> None: |
| super().__init__() |
| self.eps = eps |
| self.group_size = group_size |
| self.weight = nn.Parameter(torch.ones(dim)) |
|
|
| def forward(self, x: torch.Tensor) -> torch.Tensor: |
| h = x.float() |
| if self.group_size is not None: |
| h = h.unflatten(-1, (-1, self.group_size)) |
| h = h * torch.rsqrt(h.pow(2).mean(-1, keepdim=True) + self.eps) |
| if self.group_size is not None: |
| h = h.flatten(-2) |
| return h.to(x.dtype) * self.weight |
|
|
|
|
| def rotary_cos_sin(config: DenseConfig, num_positions: int, device: torch.device) -> tuple[torch.Tensor, torch.Tensor]: |
| """Rotary tables of shape [num_positions, head_dim] in float32 (rotate-half layout).""" |
| dim = config.head_dim |
| frequencies = 1.0 / (config.rope_theta ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim)) |
| frequencies = frequencies.to(device) |
| positions = torch.arange(num_positions, device=device, dtype=torch.float32) |
| if config.rope_precision == "bfloat16": |
| |
| |
| def bf16(t: torch.Tensor) -> torch.Tensor: |
| return t.to(torch.bfloat16).float() |
|
|
| angles = bf16(torch.outer(bf16(positions), bf16(frequencies))) |
| angles = torch.cat((angles, angles), dim=-1) |
| return bf16(angles.cos()), bf16(angles.sin()) |
| if config.rope_precision != "float32": |
| raise ValueError(f"rope_precision must be 'bfloat16' or 'float32', got {config.rope_precision!r}") |
| angles = torch.outer(positions, frequencies) |
| angles = torch.cat((angles, angles), dim=-1) |
| return angles.cos(), angles.sin() |
|
|
|
|
| def rotate_half(x: torch.Tensor) -> torch.Tensor: |
| """(a, b) -> (-b, a) over the two halves of the last dimension.""" |
| half = x.shape[-1] // 2 |
| return torch.cat((-x[..., half:], x[..., :half]), dim=-1) |
|
|
|
|
| def apply_rotary(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor: |
| """x: [batch, tokens, heads, head_dim]; cos, sin: [tokens, head_dim], already in x's dtype.""" |
| cos = cos[None, :, None, :] |
| sin = sin[None, :, None, :] |
| return x * cos + rotate_half(x) * sin |
|
|
|
|
| class Attention(nn.Module): |
| def __init__(self, config: DenseConfig) -> None: |
| super().__init__() |
| self.num_heads = config.num_attention_heads |
| self.head_dim = config.head_dim |
| width = config.num_attention_heads * config.head_dim |
| self.q_proj = nn.Linear(config.hidden_size, width, bias=False) |
| self.k_proj = nn.Linear(config.hidden_size, width, bias=False) |
| self.v_proj = nn.Linear(config.hidden_size, width, bias=False) |
| self.gate_proj = nn.Linear(config.hidden_size, width, bias=False) |
| self.o_proj = nn.Linear(width, config.hidden_size, bias=False) |
| self.q_norm = RMSNorm(width, config.rms_norm_eps, group_size=config.head_dim) |
| self.k_norm = RMSNorm(width, config.rms_norm_eps, group_size=config.head_dim) |
|
|
| def forward(self, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor: |
| batch, length, _ = x.shape |
| heads = (batch, length, self.num_heads, self.head_dim) |
| q = self.q_norm(self.q_proj(x)).view(heads) |
| k = self.k_norm(self.k_proj(x)).view(heads) |
| v = self.v_proj(x).view(heads) |
| q = apply_rotary(q, cos, sin) |
| k = apply_rotary(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 |
| ) |
| out = out.transpose(1, 2).reshape(batch, length, self.num_heads * self.head_dim) |
| out = out * torch.sigmoid(self.gate_proj(x)) |
| return self.o_proj(out) |
|
|
|
|
| class MLP(nn.Module): |
| def __init__(self, config: DenseConfig) -> None: |
| super().__init__() |
| self.gate_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False) |
| self.up_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False) |
| self.down_proj = nn.Linear(config.intermediate_size, config.hidden_size, bias=False) |
|
|
| def forward(self, x: torch.Tensor) -> torch.Tensor: |
| return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x)) |
|
|
|
|
| class DecoderLayer(nn.Module): |
| def __init__(self, config: DenseConfig) -> None: |
| super().__init__() |
| self.attn_norm = RMSNorm(config.hidden_size, config.rms_norm_eps) |
| self.attn = Attention(config) |
| self.mlp_norm = RMSNorm(config.hidden_size, config.rms_norm_eps) |
| self.mlp = MLP(config) |
|
|
| def forward(self, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor: |
| h = x + self.attn(self.attn_norm(x), cos, sin) |
| return h + self.mlp(self.mlp_norm(h)) |
|
|
|
|
| class DenseDecoder(nn.Module): |
| """Decoder-only language model. ``forward`` returns float32 logits (and the mean cross-entropy loss when |
| ``labels`` is given; label -100 is ignored).""" |
|
|
| def __init__(self, config: DenseConfig) -> None: |
| super().__init__() |
| self.config = config |
| self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size) |
| self.layers = nn.ModuleList(DecoderLayer(config) for _ in range(config.num_hidden_layers)) |
| self.norm = RMSNorm(config.hidden_size, config.rms_norm_eps) |
| self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) |
|
|
| def hidden_states(self, input_ids: torch.Tensor) -> torch.Tensor: |
| """Final hidden states after the last norm, [batch, tokens, hidden_size], in the model's dtype.""" |
| length = input_ids.shape[1] |
| if length > self.config.max_position_embeddings: |
| raise ValueError(f"{length} tokens exceed the context of {self.config.max_position_embeddings}") |
| x = self.embed_tokens(input_ids) |
| cos, sin = rotary_cos_sin(self.config, length, input_ids.device) |
| cos, sin = cos.to(x.dtype), sin.to(x.dtype) |
| for layer in self.layers: |
| x = layer(x, cos, sin) |
| return self.norm(x) |
|
|
| def forward(self, input_ids: torch.Tensor, labels: torch.Tensor | None = None): |
| logits = self.lm_head(self.hidden_states(input_ids)).float() |
| if labels is None: |
| return logits |
| loss = F.cross_entropy(logits.flatten(0, 1), labels.flatten(), ignore_index=-100) |
| return logits, loss |
|
|
| @torch.no_grad() |
| def generate(self, input_ids: torch.Tensor, max_new_tokens: int, temperature: float = 0.0) -> torch.Tensor: |
| """Simple sampling without a key/value cache (greedy when temperature is 0).""" |
| for _ in range(max_new_tokens): |
| context = input_ids[:, -self.config.max_position_embeddings:] |
| next_logits = self(context)[:, -1] |
| if temperature > 0: |
| probs = torch.softmax(next_logits / temperature, dim=-1) |
| next_ids = torch.multinomial(probs, num_samples=1) |
| else: |
| next_ids = next_logits.argmax(dim=-1, keepdim=True) |
| input_ids = torch.cat([input_ids, next_ids], dim=1) |
| return input_ids |
|
|
| @classmethod |
| def from_pretrained(cls, directory: str, device="cpu", dtype: torch.dtype = torch.float32) -> "DenseDecoder": |
| """Load ``config.json`` and ``model.safetensors`` from a released model directory. |
| |
| The weights are read and cast to ``dtype`` 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 = DenseConfig.from_json_file(os.path.join(directory, "config.json")) |
| state = load_file(os.path.join(directory, "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) |
| with torch.device("meta"): |
| model = cls(config) |
| model.load_state_dict(state, strict=True, assign=True) |
| return model.to(device).eval() |
|
|
