from __future__ import annotations from typing import Optional import torch import torch.nn.functional as F from transformers.modeling_utils import PreTrainedModel from transformers.generation import GenerationMixin from transformers.modeling_outputs import CausalLMOutput from .configuration_ksbyte import KsByteConfig # ---- Inlined ksbyte.config.py ------------------------------------------------- """Single source of truth for every tunable knob in ks_byte_lm. Mirrors the philosophy of the ks_diacritizer DiacConfig: nothing else in the codebase hard-codes a hyper-parameter. Override fields from the CLI / Modal / run_local by passing a dict to `.merge(...)`. A frozen copy of this config is embedded in every checkpoint, so a model can be rebuilt and resumed on any machine without the original launch command. """ from dataclasses import asdict, dataclass, field, replace from typing import List, Optional # ----- byte vocabulary ------------------------------------------------------- # 0..255 raw UTF-8 bytes, then three control ids. Stored as uint16 on disk. BYTE_VOCAB = 256 BOS_ID = 256 EOS_ID = 257 PAD_ID = 258 VOCAB_SIZE = 259 @dataclass class ByteLMConfig: # ------------------------------ dataset -------------------------------- hf_dataset: str = "Omarrran/KS-PRET-5M_5_million_kashmiri_Pretrainning_LLM_dataset_12M_tokens_2026" hf_revision: Optional[str] = None text_col: str = "auto" # "auto" => pick the longest string column local_text_file: Optional[str] = None # bypass HF: a local UTF-8 .txt path min_ks_ratio: float = 0.90 # drop records below this script purity min_chars: int = 2 max_chars: int = 100_000 # guard against pathological mega-rows keep_mixed_script: bool = True # keep the ~18% non-Nastaliq content # ------------------------- normalization policy ------------------------ zwnj_policy: str = "keep" # keep | to_space | strip digit_policy: str = "keep_native" # keep_native | to_ascii remove_diacritics: bool = False # NEVER true for the foundation model # ------------------------------ split ---------------------------------- val_frac: float = 0.05 test_frac: float = 0.05 split_seed: int = 13 # deterministic, content-hash based dedup: bool = True # exact-hash dedup before splitting # ----------------------------- data cache ------------------------------ data_dir: str = "data" # where {train,val,test}.bin + meta.json live rebuild_data: bool = False # force re-running data_prep # ------------------------------ model ---------------------------------- variant: str = "plain" # "plain" (implemented) | "spacebyte" (P2) d_model: int = 384 n_layers: int = 6 # used by "plain" n_heads: int = 6 n_kv_heads: int = 2 # GQA; must divide n_heads mlp_ratio: float = 2.6667 # SwiGLU hidden = round(mlp_ratio * d_model) ctx_len: int = 2048 # bytes per window rope_theta: float = 10_000.0 dropout: float = 0.1 # ON: 48MB overfits a transformer fast qk_norm: bool = True # LayerNorm on Q,K -> tames attention logits tie_embeddings: bool = True # SpaceByte-only (reserved; "spacebyte" variant lands in P2) n_local_in: int = 2 n_global: int = 6 n_local_out: int = 2 max_patches: int = 320 # ---------------------------- optimisation ----------------------------- epochs: float = 4.0 # Muennighoff: <=4 epochs ~ free on repeats lr: float = 4e-4 min_lr_ratio: float = 0.1 # cosine floor = lr * min_lr_ratio warmup_ratio: float = 0.03 weight_decay: float = 0.1 beta1: float = 0.9 beta2: float = 0.95 grad_clip: float = 1.0 z_loss_weight: float = 1e-4 # auxiliary logit-norm regularizer label_smoothing: float = 0.0 batch_size: int = 16 # sequences per micro-step grad_accum: int = 4 # effective batch = batch_size * grad_accum doc_attention_mask: bool = True # block cross-document attention + RoPE reset # --------------------------- system / perf ----------------------------- device: str = "auto" # auto | cuda | cpu bf16: bool = True # L4 (Ada) supports bf16 tf32: bool = True torch_compile: bool = False # optional; can be flaky -> off by default num_workers: int = 4 seed: int = 42 # ------------------------------- eval ---------------------------------- eval_interval: int = 500 # optimiser steps between evals eval_iters: int = 100 # micro-batches per eval estimate log_interval: int = 20 generate_every: int = 1000 # emit a sample generation every N steps generate_tokens: int = 160 early_stop_patience: int = 6 # evals without val-BPB improvement max_steps: Optional[int] = None # hard cap (None => derive from epochs) # --------------------------- io / logging ------------------------------ output_dir: str = "outputs" run_name: str = "ksbyte-plain-10m" resume: str = "auto" # auto | never | save_interval: int = 500 # also saves best-on-eval regardless save_total_limit: int = 3 use_wandb: bool = False wandb_project: str = "ks-byte-lm" # ------------------------------ helpers -------------------------------- @property def run_dir(self) -> str: return f"{self.output_dir.rstrip('/')}/{self.run_name}" @property def effective_batch(self) -> int: return self.batch_size * self.grad_accum @property def head_dim(self) -> int: if self.d_model % self.n_heads != 0: raise ValueError(f"d_model {self.d_model} not divisible by n_heads {self.n_heads}") return self.d_model // self.n_heads @property def ffn_hidden(self) -> int: return int(round(self.mlp_ratio * self.d_model)) def validate(self) -> "ByteLMConfig": """Fail fast on inconsistent settings (called by train/data entrypoints).""" errs: List[str] = [] if self.n_heads % self.n_kv_heads != 0: errs.append(f"n_heads {self.n_heads} not divisible by n_kv_heads {self.n_kv_heads}") if self.d_model % self.n_heads != 0: errs.append(f"d_model {self.d_model} not divisible by n_heads {self.n_heads}") if self.variant not in ("plain", "spacebyte"): errs.append(f"unknown variant {self.variant!r}") if self.zwnj_policy not in ("keep", "to_space", "strip"): errs.append(f"unknown zwnj_policy {self.zwnj_policy!r}") if self.digit_policy not in ("keep_native", "to_ascii"): errs.append(f"unknown digit_policy {self.digit_policy!r}") if not (0.0 <= self.val_frac + self.test_frac < 0.9): errs.append("val_frac + test_frac must be in [0, 0.9)") if self.ctx_len < 8: errs.append("ctx_len too small") if self.variant == "spacebyte": if self.n_local_in < 0 or self.n_global < 0 or self.n_local_out < 0: errs.append("spacebyte layer counts must be non-negative") if self.n_global < 1: errs.append("spacebyte requires at least one global block") if self.max_patches < 1: errs.append("spacebyte max_patches must be >= 1") if errs: raise ValueError("Invalid ByteLMConfig:\n - " + "\n - ".join(errs)) return self def merge(self, overrides: Optional[dict]) -> "ByteLMConfig": """Return a copy with given keys replaced; warns (not fails) on unknown keys.""" if not overrides: return replace(self) known = {k: v for k, v in overrides.items() if k in self.__dataclass_fields__} unknown = set(overrides) - set(known) if unknown: print(f"[config] WARNING ignoring unknown overrides: {sorted(unknown)}") return replace(self, **known) def to_dict(self) -> dict: return asdict(self) @classmethod def from_dict(cls, d: dict) -> "ByteLMConfig": known = {k: v for k, v in d.items() if k in cls.__dataclass_fields__} return cls(**known) # ---- Inlined ksbyte.layers.py ------------------------------------------------- """Modern Llama-style building blocks, isolated so architectures swap easily. Components: RMSNorm (pre-norm), RoPE (position ids supplied by the loader so it resets per document), GQA attention with optional QK-Norm and document-block attention masking, and a SwiGLU MLP. """ import torch import torch.nn as nn import torch.nn.functional as F class RMSNorm(nn.Module): def __init__(self, dim: int, eps: float = 1e-6): super().__init__() self.eps = eps self.weight = nn.Parameter(torch.ones(dim)) def forward(self, x: torch.Tensor) -> torch.Tensor: dtype = x.dtype x = x.float() x = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) return (x * self.weight).to(dtype) class RotaryEmbedding(nn.Module): """Builds cos/sin tables from explicit position ids ([B,T]).""" def __init__(self, head_dim: int, theta: float = 10_000.0): super().__init__() if head_dim % 2 != 0: raise ValueError("head_dim must be even for RoPE") inv_freq = 1.0 / (theta ** (torch.arange(0, head_dim, 2).float() / head_dim)) self.register_buffer("inv_freq", inv_freq, persistent=False) def forward(self, pos_ids: torch.Tensor): # pos_ids: [B,T] -> freqs [B,T,head_dim/2] freqs = pos_ids.float().unsqueeze(-1) * self.inv_freq.to(pos_ids.device) emb = torch.cat((freqs, freqs), dim=-1) # [B,T,head_dim] return emb.cos().unsqueeze(1), emb.sin().unsqueeze(1) # [B,1,T,head_dim] def _rotate_half(x: torch.Tensor) -> torch.Tensor: x1, x2 = x.chunk(2, dim=-1) return torch.cat((-x2, x1), dim=-1) def apply_rope(q: torch.Tensor, k: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor): q = (q * cos) + (_rotate_half(q) * sin) k = (k * cos) + (_rotate_half(k) * sin) return q, k class Attention(nn.Module): def __init__(self, cfg): super().__init__() self.n_heads = cfg.n_heads self.n_kv_heads = cfg.n_kv_heads self.head_dim = cfg.head_dim self.dropout = cfg.dropout self.qk_norm = cfg.qk_norm d = cfg.d_model self.wq = nn.Linear(d, self.n_heads * self.head_dim, bias=False) self.wk = nn.Linear(d, self.n_kv_heads * self.head_dim, bias=False) self.wv = nn.Linear(d, self.n_kv_heads * self.head_dim, bias=False) self.wo = nn.Linear(self.n_heads * self.head_dim, d, bias=False) if self.qk_norm: self.q_norm = RMSNorm(self.head_dim) self.k_norm = RMSNorm(self.head_dim) def forward(self, x, cos, sin, attn_mask): B, T, _ = x.shape q = self.wq(x).view(B, T, self.n_heads, self.head_dim) k = self.wk(x).view(B, T, self.n_kv_heads, self.head_dim) v = self.wv(x).view(B, T, self.n_kv_heads, self.head_dim) if self.qk_norm: q = self.q_norm(q) k = self.k_norm(k) q = q.transpose(1, 2) # [B,H,T,D] k = k.transpose(1, 2) # [B,KV,T,D] v = v.transpose(1, 2) q, k = apply_rope(q, k, cos, sin) # GQA: expand kv heads to match query heads. if self.n_kv_heads != self.n_heads: rep = self.n_heads // self.n_kv_heads k = k.repeat_interleave(rep, dim=1) v = v.repeat_interleave(rep, dim=1) dp = self.dropout if self.training else 0.0 if attn_mask is None: out = F.scaled_dot_product_attention(q, k, v, is_causal=True, dropout_p=dp) else: # boolean mask [B,1,T,T]: True => attend. out = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask, dropout_p=dp) out = out.transpose(1, 2).contiguous().view(B, T, self.n_heads * self.head_dim) return self.wo(out) class SwiGLU(nn.Module): def __init__(self, cfg): super().__init__() d, h = cfg.d_model, cfg.ffn_hidden self.w_gate = nn.Linear(d, h, bias=False) self.w_up = nn.Linear(d, h, bias=False) self.w_down = nn.Linear(h, d, bias=False) def forward(self, x): return self.w_down(F.silu(self.w_gate(x)) * self.w_up(x)) class Block(nn.Module): def __init__(self, cfg): super().__init__() self.norm1 = RMSNorm(cfg.d_model) self.attn = Attention(cfg) self.norm2 = RMSNorm(cfg.d_model) self.mlp = SwiGLU(cfg) self.drop = nn.Dropout(cfg.dropout) def forward(self, x, cos, sin, attn_mask): x = x + self.drop(self.attn(self.norm1(x), cos, sin, attn_mask)) x = x + self.drop(self.mlp(self.norm2(x))) return x def build_doc_attn_mask(seg_ids: torch.Tensor) -> torch.Tensor: """Boolean [B,1,T,T] mask: position i may attend to j iff j<=i (causal) and they share a document segment. True => attend.""" B, T = seg_ids.shape causal = torch.tril(torch.ones(T, T, dtype=torch.bool, device=seg_ids.device)) same_doc = seg_ids.unsqueeze(2) == seg_ids.unsqueeze(1) # [B,T,T] return (causal.unsqueeze(0) & same_doc).unsqueeze(1) # ---- Inlined ksbyte.model.py -------------------------------------------------- """Byte-level decoder language models. `variant="plain"` is the P1 small Llama-style causal decoder over raw UTF-8 bytes. `variant="spacebyte"` is the P2 word-boundary hierarchy: cheap local byte blocks run everywhere, expensive global blocks run only on patch-boundary positions, then the latest patch context is scattered back to every byte before optional local output blocks. forward() returns (logits, loss, parts) where `parts` breaks the loss into the cross-entropy term and the z-loss term for logging. """ import math from typing import Dict, Optional, Tuple import torch import torch.nn as nn import torch.nn.functional as F # ASCII whitespace bytes that delimit Kashmiri words in the normalized corpus. # UTF-8 non-ASCII punctuation is deliberately not treated as a patch boundary. _SPACEBYTE_WHITESPACE = (9, 10, 11, 12, 13, 32) # \t \n \v \f \r space def build_spacebyte_boundary_mask( x: torch.Tensor, seg_ids: Optional[torch.Tensor] = None, spacelike_bytes: Tuple[int, ...] = _SPACEBYTE_WHITESPACE, ) -> torch.Tensor: """Return [B,T] positions promoted to SpaceByte global patches. A text patch starts at the first byte of each document/window and at a spacelike byte that is not itself preceded by another spacelike byte. BOS/EOS ids are also promoted because packed byte streams use them as hard document boundaries. Marking only the first byte in a run of spaces prevents wasting global compute on repeated whitespace. """ if x.dim() != 2: raise ValueError(f"x must be [B,T], got shape {tuple(x.shape)}") B, T = x.shape is_space = torch.zeros_like(x, dtype=torch.bool) for byte in spacelike_bytes: is_space |= x == byte prev_space = torch.zeros_like(is_space) prev_space[:, 1:] = is_space[:, :-1] boundary = is_space & ~prev_space boundary |= (x == BOS_ID) | (x == EOS_ID) # Every sampled window needs at least one patch. When segment ids are known, # also seed each new document segment so document-aware global attention has # a patch before the first visible space. boundary[:, 0] = True if seg_ids is not None: new_doc = torch.zeros_like(boundary) new_doc[:, 1:] = seg_ids[:, 1:] != seg_ids[:, :-1] boundary |= new_doc return boundary class ByteDecoder(nn.Module): def __init__(self, cfg: ByteLMConfig): super().__init__() cfg.validate() self.cfg = cfg self.embed = nn.Embedding(VOCAB_SIZE, cfg.d_model) self.drop = nn.Dropout(cfg.dropout) self.blocks = nn.ModuleList([Block(cfg) for _ in range(cfg.n_layers)]) self.norm_f = RMSNorm(cfg.d_model) self.lm_head = nn.Linear(cfg.d_model, VOCAB_SIZE, bias=False) self.rope = RotaryEmbedding(cfg.head_dim, cfg.rope_theta) if cfg.tie_embeddings: self.lm_head.weight = self.embed.weight self.apply(self._init_weights) # Scale residual-path projections by 1/sqrt(2*n_layers) for stable depth. scale = 1.0 / math.sqrt(2 * cfg.n_layers) for name, p in self.named_parameters(): if name.endswith("wo.weight") or name.endswith("w_down.weight"): with torch.no_grad(): p.mul_(scale) @staticmethod def _init_weights(m: nn.Module): if isinstance(m, nn.Linear): nn.init.normal_(m.weight, mean=0.0, std=0.02) if m.bias is not None: nn.init.zeros_(m.bias) elif isinstance(m, nn.Embedding): nn.init.normal_(m.weight, mean=0.0, std=0.02) def num_params(self, non_embedding: bool = False) -> int: n = sum(p.numel() for p in self.parameters()) if non_embedding and not self.cfg.tie_embeddings: n -= self.embed.weight.numel() return n def forward( self, x: torch.Tensor, y: Optional[torch.Tensor] = None, seg_ids: Optional[torch.Tensor] = None, pos_ids: Optional[torch.Tensor] = None, ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Dict[str, float]]: B, T = x.shape if pos_ids is None: pos_ids = torch.arange(T, device=x.device).unsqueeze(0).expand(B, T) cos, sin = self.rope(pos_ids) attn_mask = build_doc_attn_mask(seg_ids) if seg_ids is not None else None h = self.drop(self.embed(x)) for blk in self.blocks: h = blk(h, cos, sin, attn_mask) h = self.norm_f(h) logits = self.lm_head(h) loss, parts = _loss_parts(logits, y, self.cfg) return logits, loss, parts @torch.no_grad() def generate(self, idx: torch.Tensor, max_new_tokens: int, temperature: float = 1.0, top_k: Optional[int] = None, eos_id: Optional[int] = None) -> torch.Tensor: """Autoregressive sampling (no doc-mask; single growing sequence).""" self.eval() for _ in range(max_new_tokens): idx_cond = idx[:, -self.cfg.ctx_len:] logits, _, _ = self(idx_cond) logits = logits[:, -1, :] / max(temperature, 1e-6) if top_k is not None: v, _ = torch.topk(logits, min(top_k, logits.size(-1))) logits[logits < v[:, [-1]]] = float("-inf") probs = F.softmax(logits, dim=-1) nxt = torch.multinomial(probs, num_samples=1) idx = torch.cat((idx, nxt), dim=1) if eos_id is not None and (nxt == eos_id).all(): break return idx class SpaceByteDecoder(nn.Module): """SpaceByte-style hierarchy over raw bytes. Local-in blocks process every byte cheaply. Global blocks process only patch boundary states (word/document boundaries), with causal/document masks over the patch sequence. The latest available global patch state is then scattered back to each byte and added as context before local-out blocks and the shared LM head. Setting n_local_in=n_local_out=0 and n_global=n_layers makes this reduce exactly to `ByteDecoder` when every position is a boundary. """ def __init__(self, cfg: ByteLMConfig): super().__init__() cfg.validate() self.cfg = cfg self.embed = nn.Embedding(VOCAB_SIZE, cfg.d_model) self.drop = nn.Dropout(cfg.dropout) self.local_in_blocks = nn.ModuleList([Block(cfg) for _ in range(cfg.n_local_in)]) self.global_blocks = nn.ModuleList([Block(cfg) for _ in range(cfg.n_global)]) self.local_out_blocks = nn.ModuleList([Block(cfg) for _ in range(cfg.n_local_out)]) self.norm_f = RMSNorm(cfg.d_model) self.lm_head = nn.Linear(cfg.d_model, VOCAB_SIZE, bias=False) self.rope = RotaryEmbedding(cfg.head_dim, cfg.rope_theta) if cfg.tie_embeddings: self.lm_head.weight = self.embed.weight self.apply(ByteDecoder._init_weights) total_layers = max(1, cfg.n_local_in + cfg.n_global + cfg.n_local_out) scale = 1.0 / math.sqrt(2 * total_layers) for name, p in self.named_parameters(): if name.endswith("wo.weight") or name.endswith("w_down.weight"): with torch.no_grad(): p.mul_(scale) def num_params(self, non_embedding: bool = False) -> int: n = sum(p.numel() for p in self.parameters()) if non_embedding and not self.cfg.tie_embeddings: n -= self.embed.weight.numel() return n @staticmethod def gather_patches( h: torch.Tensor, boundary_mask: torch.Tensor, pos_ids: torch.Tensor, seg_ids: Optional[torch.Tensor], max_patches: int, ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: """Gather boundary states into a padded [B,P,D] patch sequence. Returns `(patch_h, patch_valid, patch_pos, patch_seg, patch_ids)`, where `patch_ids[b,t]` is the index of the latest patch at or before byte `t`. If a sequence has more than `max_patches` boundaries, later bytes reuse the last retained patch context; this keeps memory bounded and explicit. """ if h.dim() != 3 or boundary_mask.dim() != 2: raise ValueError("h must be [B,T,D] and boundary_mask must be [B,T]") B, T, D = h.shape counts = boundary_mask.sum(dim=1).clamp_min(1) P = int(min(max_patches, counts.max().item())) device = h.device gather_idx = torch.zeros((B, P), dtype=torch.long, device=device) patch_valid = torch.zeros((B, P), dtype=torch.bool, device=device) patch_ids = torch.empty((B, T), dtype=torch.long, device=device) for b in range(B): idx = torch.nonzero(boundary_mask[b], as_tuple=False).flatten() if idx.numel() == 0: idx = torch.zeros(1, dtype=torch.long, device=device) keep = idx[:P] n = keep.numel() gather_idx[b, :n] = keep patch_valid[b, :n] = True ordinal = torch.cumsum(boundary_mask[b].to(torch.long), dim=0) - 1 patch_ids[b] = ordinal.clamp(min=0, max=max(n - 1, 0)) patch_h = h.gather(1, gather_idx.unsqueeze(-1).expand(B, P, D)) patch_pos = pos_ids.gather(1, gather_idx) if seg_ids is None: patch_seg = torch.zeros((B, P), dtype=torch.long, device=device) else: patch_seg = seg_ids.gather(1, gather_idx) return patch_h, patch_valid, patch_pos, patch_seg, patch_ids @staticmethod def scatter_patches(patch_h: torch.Tensor, patch_ids: torch.Tensor) -> torch.Tensor: """Scatter latest patch states from [B,P,D] back to each byte [B,T,D].""" B, T = patch_ids.shape D = patch_h.size(-1) return patch_h.gather(1, patch_ids.unsqueeze(-1).expand(B, T, D)) @staticmethod def _patch_attn_mask(patch_seg: torch.Tensor, patch_valid: torch.Tensor) -> torch.Tensor: """Causal/document patch mask that keeps padded query rows finite.""" mask = build_doc_attn_mask(patch_seg) valid = patch_valid.unsqueeze(1) mask = mask & valid.unsqueeze(-1) & valid.unsqueeze(-2) eye = torch.eye(patch_seg.size(1), dtype=torch.bool, device=patch_seg.device) # Padded rows are ignored later, but SDPA still needs at least one True. mask = mask | (eye.unsqueeze(0).unsqueeze(0) & ~valid.unsqueeze(-1)) return mask def forward( self, x: torch.Tensor, y: Optional[torch.Tensor] = None, seg_ids: Optional[torch.Tensor] = None, pos_ids: Optional[torch.Tensor] = None, boundary_mask: Optional[torch.Tensor] = None, ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Dict[str, float]]: B, T = x.shape if pos_ids is None: pos_ids = torch.arange(T, device=x.device).unsqueeze(0).expand(B, T) byte_cos, byte_sin = self.rope(pos_ids) byte_attn_mask = build_doc_attn_mask(seg_ids) if seg_ids is not None else None h = self.drop(self.embed(x)) for blk in self.local_in_blocks: h = blk(h, byte_cos, byte_sin, byte_attn_mask) if boundary_mask is None: boundary_mask = build_spacebyte_boundary_mask(x, seg_ids=seg_ids) patch_h, patch_valid, patch_pos, patch_seg, patch_ids = self.gather_patches( h, boundary_mask, pos_ids, seg_ids, self.cfg.max_patches ) patch_cos, patch_sin = self.rope(patch_pos) patch_attn_mask = self._patch_attn_mask(patch_seg, patch_valid) for blk in self.global_blocks: patch_h = blk(patch_h, patch_cos, patch_sin, patch_attn_mask) patch_context = self.scatter_patches(patch_h, patch_ids) if self.cfg.n_local_in == 0 and self.cfg.n_local_out == 0: # Degenerate hierarchy: the global stream is the whole residual # stream, so all-boundary inputs exactly match the plain decoder. h = patch_context else: h = h + patch_context for blk in self.local_out_blocks: h = blk(h, byte_cos, byte_sin, byte_attn_mask) h = self.norm_f(h) logits = self.lm_head(h) loss, parts = _loss_parts(logits, y, self.cfg) return logits, loss, parts @torch.no_grad() def generate(self, idx: torch.Tensor, max_new_tokens: int, temperature: float = 1.0, top_k: Optional[int] = None, eos_id: Optional[int] = None) -> torch.Tensor: """Autoregressive sampling with SpaceByte boundary detection per window.""" self.eval() for _ in range(max_new_tokens): idx_cond = idx[:, -self.cfg.ctx_len:] logits, _, _ = self(idx_cond) logits = logits[:, -1, :] / max(temperature, 1e-6) if top_k is not None: v, _ = torch.topk(logits, min(top_k, logits.size(-1))) logits[logits < v[:, [-1]]] = float("-inf") probs = F.softmax(logits, dim=-1) nxt = torch.multinomial(probs, num_samples=1) idx = torch.cat((idx, nxt), dim=1) if eos_id is not None and (nxt == eos_id).all(): break return idx def _loss_parts( logits: torch.Tensor, y: Optional[torch.Tensor], cfg: ByteLMConfig, ) -> Tuple[Optional[torch.Tensor], Dict[str, float]]: """Shared CE + z-loss computation for both decoder variants.""" if y is None: return None, {} ce = F.cross_entropy( logits.view(-1, VOCAB_SIZE), y.reshape(-1), label_smoothing=cfg.label_smoothing, ) z = logits.logsumexp(dim=-1).pow(2).mean() loss = ce + cfg.z_loss_weight * z return loss, {"ce": ce.item(), "z": z.item(), "loss": loss.item()} def build_model(cfg: ByteLMConfig) -> nn.Module: if cfg.variant == "plain": return ByteDecoder(cfg) if cfg.variant == "spacebyte": return SpaceByteDecoder(cfg) raise ValueError(f"unknown variant {cfg.variant!r}") def _to_byte_lm_config(config: KsByteConfig) -> ByteLMConfig: keys = ByteLMConfig.__dataclass_fields__.keys() return ByteLMConfig.from_dict({k: getattr(config, k) for k in keys if hasattr(config, k)}) class KsByteForCausalLM(PreTrainedModel, GenerationMixin): config_class = KsByteConfig base_model_prefix = "model" main_input_name = "input_ids" supports_gradient_checkpointing = False _supports_cache_class = False _tied_weights_keys = ["model.lm_head.weight"] # Some Transformers releases expect this to be a dict and call .keys(); # others only need membership. Keep it dict-shaped for Colab compatibility. all_tied_weights_keys = {"model.lm_head.weight": "model.embed.weight"} def __init__(self, config: KsByteConfig): super().__init__(config) self.model = build_model(_to_byte_lm_config(config)) def get_input_embeddings(self): return self.model.embed def set_input_embeddings(self, value): self.model.embed = value if getattr(self.config, "tie_embeddings", False): self.model.lm_head.weight = self.model.embed.weight def get_output_embeddings(self): return self.model.lm_head def set_output_embeddings(self, new_embeddings): self.model.lm_head = new_embeddings def prepare_inputs_for_generation(self, input_ids, **kwargs): if input_ids.shape[1] > self.config.ctx_len: input_ids = input_ids[:, -self.config.ctx_len:] return {"input_ids": input_ids} def forward( self, input_ids: torch.LongTensor, attention_mask: Optional[torch.Tensor] = None, labels: Optional[torch.LongTensor] = None, **kwargs, ): logits, _, _ = self.model(input_ids) loss = None if labels is not None: shift_logits = logits[:, :-1, :].contiguous() shift_labels = labels[:, 1:].contiguous() loss = F.cross_entropy( shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1), ignore_index=-100, ) return CausalLMOutput(loss=loss, logits=logits) @torch.no_grad() def generate_bytes(self, input_ids, max_new_tokens=200, temperature=0.8, top_k=50): return self.model.generate( input_ids, max_new_tokens=max_new_tokens, temperature=temperature, top_k=top_k, eos_id=self.config.eos_token_id, )