blt-trainer 0.1.0__tar.gz

blt_trainer-0.1.0/PKG-INFO

@@ -0,0 +1,8 @@
+Metadata-Version: 2.5
+Name: blt-trainer
+Version: 0.1.0
+Summary: Pure functional 96-layer Byte Latent Transformer trainer with Muon and weighted streaming datamix
+Requires-Python: >=3.10
+Requires-Dist: datasets>=2.16.0
+Requires-Dist: numpy>=1.24.0
+Requires-Dist: torch>=2.2.0

blt_trainer-0.1.0/pyproject.toml

@@ -0,0 +1,21 @@
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[project]
+name = "blt-trainer"
+version = "0.1.0"
+description = "Pure functional 96-layer Byte Latent Transformer trainer with Muon and weighted streaming datamix"
+readme = "README.md"
+requires-python = ">=3.10"
+dependencies = [
+    "torch>=2.2.0",
+    "datasets>=2.16.0",
+    "numpy>=1.24.0",
+]
+
+[project.scripts]
+blt-train = "blt_trainer.cli:main"
+
+[tool.hatch.build.targets.wheel]
+packages = ["src/blt_trainer"]

blt_trainer-0.1.0/src/blt_trainer/cli.py

@@ -0,0 +1,172 @@
+import argparse
+import sys
+import torch
+import torch.nn.functional as F
+
+from blt_trainer.model import (
+    blt_local_encoder_with_surprise,
+    global_96_layer_backbone_single,
+    local_decoder_single,
+    precompute_rope_freqs,
+)
+from blt_trainer.optim import (
+    MuonWithScheduleSkip,
+    ScheduleFreeAdamW,
+    cut_cross_entropy,
+    partition_params_for_optimizers,
+)
+from blt_trainer.datamix import get_mix_loader
+
+def parse_args():
+    parser = argparse.ArgumentParser(
+        description="CLI trainer for a 96-layer Byte Latent Transformer with ablation datamix."
+    )
+    # Architecture
+    parser.add_argument("--num-layers", type=int, default=96, help="Global backbone layer depth")
+    parser.add_argument("--d-global", type=int, default=288, help="Global hidden dimension")
+    parser.add_argument("--d-local", type=int, default=128, help="Local encoder/decoder dimension")
+    parser.add_argument("--num-heads", type=int, default=6, help="Backbone MHA attention heads")
+    parser.add_argument("--max-patches", type=int, default=32, help="Max static patch capacity per sequence")
+    parser.add_argument("--entropy-tau", type=float, default=3.2, help="Surprise threshold for dynamic boundaries")
+
+    # Training & Data
+    parser.add_argument("--batch-size", type=int, default=8, help="Batch size per step")
+    parser.add_argument("--block-size", type=int, default=512, help="Byte context length")
+    parser.add_argument("--lr-muon", type=float, default=0.02, help="Learning rate for Muon (2D weights)")
+    parser.add_argument("--lr-adam", type=float, default=1e-3, help="Learning rate for AdamW (1D / Embeddings)")
+    parser.add_argument("--steps", type=int, default=10000, help="Total training steps")
+    parser.add_argument("--device", type=str, default="cuda" if torch.cuda.is_available() else "cpu")
+    parser.add_argument("--compile", action="store_true", help="Enable torch.compile kernel fusion")
+
+    return parser.parse_args()
+
+def init_all_parameters(args, device):
+    head_dim = args.d_global // args.num_heads
+    d_mlp = (args.d_global * 8) // 3
+
+    params = {
+        # Entropy predictor
+        "entropy_embed": torch.randn(256, 64, device=device) * 0.02,
+        "entropy_head":  torch.randn(256, 64, device=device) * 0.02,
+
+        # Local Encoder
+        "byte_embed":    torch.randn(256, args.d_local, device=device) * 0.02,
+        "conv_w":        torch.randn(args.d_local, 1, 3, device=device) * 0.02,
+        "q_seed":        torch.randn(args.max_patches, args.d_global, device=device) * 0.02,
+        "W_k_enc":       torch.randn(args.d_global, args.d_local, device=device) * 0.02,
+        "W_v_enc":       torch.randn(args.d_global, args.d_local, device=device) * 0.02,
+        "W_out_enc":     torch.randn(args.d_global, args.d_global, device=device) * 0.02,
+
+        # 96-Layer Global Backbone
+        "final_norm_w":  torch.ones(args.d_global, device=device),
+        
+        # Local Decoder
+        "cross_norm_w":  torch.ones(args.d_local, device=device),
+        "W_q_cross":     torch.randn(args.d_local, args.d_local, device=device) * 0.02,
+        "W_k_cross":     torch.randn(args.d_local, args.d_global, device=device) * 0.02,
+        "W_v_cross":     torch.randn(args.d_local, args.d_global, device=device) * 0.02,
+        "W_o_cross":     torch.randn(args.d_local, args.d_local, device=device) * 0.02,
+        "self_norm_w":   torch.ones(args.d_local, device=device),
+        "W_q_self":      torch.randn(args.d_local, args.d_local, device=device) * 0.02,
+        "W_k_self":      torch.randn(args.d_local, args.d_local, device=device) * 0.02,
+        "W_v_self":      torch.randn(args.d_local, args.d_local, device=device) * 0.02,
+        "W_o_self":      torch.randn(args.d_local, args.d_local, device=device) * 0.02,
+        "dec_mlp_norm":  torch.ones(args.d_local, device=device),
+        "W_gate_dec":    torch.randn(args.d_local * 3, args.d_local, device=device) * 0.02,
+        "W_up_dec":      torch.randn(args.d_local * 3, args.d_local, device=device) * 0.02,
+        "W_down_dec":    torch.randn(args.d_local, args.d_local * 3, device=device) * 0.02,
+        "dec_final_norm":torch.ones(args.d_local, device=device),
+        "lm_head":       torch.randn(256, args.d_local, device=device) * 0.02,
+    }
+
+    # Populate 96 layers
+    for i in range(args.num_layers):
+        params[f"layer_{i}"] = {
+            "norm1_w":   torch.ones(args.d_global, device=device),
+            "W_q":       torch.randn(args.d_global, args.d_global, device=device) * 0.02,
+            "W_k":       torch.randn(args.d_global, args.d_global, device=device) * 0.02,
+            "W_v":       torch.randn(args.d_global, args.d_global, device=device) * 0.02,
+            "W_o":       torch.randn(args.d_global, args.d_global, device=device) * 0.02,
+            "q_norm_w":  torch.ones(head_dim, device=device),
+            "k_norm_w":  torch.ones(head_dim, device=device),
+            "norm2_w":   torch.ones(args.d_global, device=device),
+            "W_gate":    torch.randn(d_mlp, args.d_global, device=device) * 0.02,
+            "W_up":      torch.randn(d_mlp, args.d_global, device=device) * 0.02,
+            "W_down":    torch.randn(args.d_global, d_mlp, device=device) * 0.02,
+        }
+
+    # Enable autograd for all weights
+    for v in params.values():
+        if isinstance(v, dict):
+            for sub_v in v.values():
+                sub_v.requires_grad_(True)
+        else:
+            v.requires_grad_(True)
+
+    return params
+
+def main():
+    args = parse_args()
+    device = torch.device(args.device)
+    print(f">> Initializing BLT 96-layer trainer on {device} (Compile={args.compile})...")
+
+    # 1. Initialize Model Weights
+    params = init_all_parameters(args, device)
+    head_dim = args.d_global // args.num_heads
+    cos, sin = precompute_rope_freqs(head_dim, args.max_patches, device=device)
+
+    # 2. Partition into Muon and AdamW Parameter Sets
+    muon_tensors, adam_tensors = partition_params_for_optimizers(params)
+    opt_muon = MuonWithScheduleSkip(muon_tensors, lr=args.lr_muon)
+    opt_adam = ScheduleFreeAdamW(adam_tensors, lr=args.lr_adam)
+
+    # 3. Build Dataloader from streaming weighted mix
+    print(">> Initializing Hugging Face multi-stream datamix (DCLM, FLAN, WikiMCQA, Reasoning, Ubuntu)...")
+    loader = get_mix_loader(block_size=args.block_size, batch_size=args.batch_size)
+    data_iter = iter(loader)
+
+    # 4. Training Loop
+    print(f">> Training starting: {args.steps} steps.")
+    for step in range(1, args.steps + 1):
+        batch_x, batch_y = next(data_iter)
+        batch_x, batch_y = batch_x.to(device), batch_y.to(device)
+
+        opt_muon.zero_grad()
+        opt_adam.zero_grad()
+
+        # Step 1: Dynamic entropy-based patch extraction
+        # (Uniform patch fallback indexing for the decoder byte map)
+        chunk = args.block_size // args.max_patches
+        byte_to_patch = (torch.arange(args.block_size, device=device) // chunk).unsqueeze(0).expand(args.batch_size, -1)
+
+        # Batch forward pass via functional calls
+        # 1. Local Encoder
+        patches = torch.vmap(
+            blt_local_encoder_with_surprise,
+            in_dims=(None, 0, None, None)
+        )(params, batch_x, args.entropy_tau, args.max_patches)
+
+        # 2. 96-Layer Latent Backbone
+        backbone_out = torch.vmap(
+            global_96_layer_backbone_single,
+            in_dims=(0, None, None, None, None, None, None)
+        )(patches, params, cos, sin, args.num_heads, head_dim, args.num_layers)
+
+        # 3. Local Decoder Logits
+        h_decoded = torch.vmap(
+            local_decoder_single,
+            in_dims=(0, 0, 0, None, None, None)
+        )(batch_x, backbone_out, byte_to_patch, params, args.d_local, 4)
+
+        # 4. Chunked Loss calculation (Cut Cross-Entropy)
+        loss = cut_cross_entropy(h_decoded, params["lm_head"], batch_y, chunk_size=512)
+        loss.backward()
+
+        opt_muon.step()
+        opt_adam.step()
+
+        if step % 10 == 0 or step == 1:
+            print(f"Step {step:05d} | Loss: {loss.item():.4f} | Processed: {step * args.batch_size * args.block_size} bytes")
+
+if __name__ == "__main__":
+    main()

blt_trainer-0.1.0/src/blt_trainer/datamix.py

@@ -0,0 +1,88 @@
+import random
+import numpy as np
+import torch
+from torch.utils.data import IterableDataset, DataLoader
+from datasets import load_dataset
+
+# Normalized mixture weights based on the ablation spec
+MIXTURE_WEIGHTS = {
+    "dclm": 800.0 / 1175.0,
+    "flan": 200.0 / 1175.0,
+    "wiki_mcqa": 100.0 / 1175.0,
+    "reasoning": 50.0 / 1175.0,
+    "ubuntu": 25.0 / 1175.0,
+}
+
+def create_stream_iterators():
+    """Initializes streaming iterables from Hugging Face for each split."""
+    # 1. DCLM Baseline
+    ds_dclm = iter(load_dataset("mlfoundations/dclm-baseline-1.0", split="train", streaming=True))
+    
+    # 2. FLAN (Collection / CoT)
+    ds_flan = iter(load_dataset("Muennighoff/flan", split="train", streaming=True))
+    
+    # 3. Wikipedia MCQA (e.g. WikiQA or OpenBookQA/Wikipedia)
+    ds_wiki = iter(load_dataset("wiki_qa", split="train", streaming=True))
+    
+    # 4. Reasoning Traces (OpenMath / GSM8k / NumGLUE style reasoning)
+    ds_reasoning = iter(load_dataset("open-web-math/open-web-math", split="train", streaming=True))
+    
+    # 5. Ubuntu IRC Dialogue
+    ds_ubuntu = iter(load_dataset("sedthh/ubuntu_dialogue_qa", split="train", streaming=True))
+
+    return {
+        "dclm": (ds_dclm, lambda x: x.get("text", "")),
+        "flan": (ds_flan, lambda x: f"Question: {x.get('inputs', '')}\nAnswer: {x.get('targets', '')}\n"),
+        "wiki_mcqa": (ds_wiki, lambda x: f"Question: {x.get('question', '')}\nContext: {x.get('document_title', '')}\nAnswer: {x.get('answer', '')}\n"),
+        "reasoning": (ds_reasoning, lambda x: x.get("text", "")),
+        "ubuntu": (ds_ubuntu, lambda x: f"User: {x.get('question', '')}\nReply: {x.get('answer', '')}\n"),
+    }
+
+class WeightedByteMixDataset(IterableDataset):
+    def __init__(self, block_size: int = 512, seed: int = 42):
+        self.block_size = block_size
+        self.seed = seed
+
+    def __iter__(self):
+        random.seed(self.seed)
+        streams = create_stream_iterators()
+        sources = list(MIXTURE_WEIGHTS.keys())
+        weights = [MIXTURE_WEIGHTS[s] for s in sources]
+        
+        byte_buffer = bytearray()
+        needed = self.block_size + 1
+
+        while True:
+            # Sample which dataset to fetch from based on mix ratio
+            chosen_source = random.choices(sources, weights=weights, k=1)[0]
+            stream, text_extractor = streams[chosen_source]
+            
+            try:
+                sample = next(stream)
+                text = text_extractor(sample) + "\n"
+            except StopIteration:
+                # Re-seed exhaustion fallback
+                streams = create_stream_iterators()
+                continue
+
+            byte_buffer.extend(text.encode("utf-8", errors="ignore"))
+
+            # Drain full blocks from buffer
+            while len(byte_buffer) >= needed:
+                chunk = bytes(byte_buffer[:needed])
+                del byte_buffer[:self.block_size]
+                
+                arr = np.frombuffer(chunk, dtype=np.uint8).astype(np.int64)
+                x = torch.from_numpy(arr[:-1])
+                y = torch.from_numpy(arr[1:])
+                yield x, y
+
+def get_mix_loader(block_size: int, batch_size: int, num_workers: int = 2):
+    dataset = WeightedByteMixDataset(block_size=block_size)
+    return DataLoader(
+        dataset,
+        batch_size=batch_size,
+        num_workers=num_workers,
+        pin_memory=True,
+        drop_last=True
+    )

blt_trainer-0.1.0/src/blt_trainer/model.py

@@ -0,0 +1,258 @@
+import torch
+import torch.nn.functional as F
+
+# =====================================================================
+# 1. Positional Embeddings & Normalization Utilities
+# =====================================================================
+
+def rms_norm(x: torch.Tensor, weight: torch.Tensor, eps: float = 1e-6) -> torch.Tensor:
+    """Root Mean Square Layer Normalization[cite: 1]."""
+    var = x.pow(2).mean(-1, keepdim=True)
+    return x * torch.rsqrt(var + eps) * weight
+
+
+def precompute_rope_freqs(
+    head_dim: int, max_seq_len: int, theta: float = 10000.0, device: str = "cpu"
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Precomputes complex exponential cos/sin tables for RoPE[cite: 1]."""
+    freqs = 1.0 / (theta ** (torch.arange(0, head_dim, 2, device=device).float() / head_dim))
+    t = torch.arange(max_seq_len, device=device).float()
+    freqs = torch.outer(t, freqs)
+    cos = torch.cos(freqs).repeat_interleave(2, dim=-1).unsqueeze(0)
+    sin = torch.sin(freqs).repeat_interleave(2, dim=-1).unsqueeze(0)
+    return cos, sin
+
+
+def apply_rope(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
+    """Applies Rotary Position Embeddings to Q or K tensors[cite: 1]."""
+    # x: [num_heads, seq_len, head_dim]
+    half = x.shape[-1] // 2
+    x1, x2 = x[..., :half], x[..., half:]
+    rotated = torch.cat((-x2, x1), dim=-1)
+    return (x * cos) + (rotated * sin)
+
+
+# =====================================================================
+# 2. Local Byte Encoder with Dynamic Surprise Factor
+# =====================================================================
+
+def compute_entropy_and_patches(
+    byte_ids: torch.Tensor,       # [L]
+    byte_embed: torch.Tensor,     # [256, d_small]
+    W_lm: torch.Tensor,           # [256, d_small]
+    entropy_threshold: float,
+    max_patches: int,
+    max_patch_len: int = 16,
+) -> torch.Tensor:
+    """Calculates byte-level cross-entropy surprise and returns a boolean patch mask[cite: 1]."""
+    L = byte_ids.shape[0]
+
+    # 1. Predict next byte logits with a shallow byte model[cite: 1]
+    x_small = F.embedding(byte_ids, byte_embed)  # [L, d_small][cite: 1]
+    logits = F.linear(x_small[:-1], W_lm)        # [L - 1, 256][cite: 1]
+    targets = byte_ids[1:]                       # [L - 1][cite: 1]
+
+    # 2. Surprise factor: negative log likelihood[cite: 1]
+    log_probs = F.log_softmax(logits, dim=-1) #[cite: 1]
+    surprise = -log_probs.gather(dim=-1, index=targets.unsqueeze(-1)).squeeze(-1) #[cite: 1]
+    surprise = torch.cat([torch.tensor([0.0], device=byte_ids.device), surprise])  # [L][cite: 1]
+
+    # 3. Mark boundaries where surprise spikes[cite: 1]
+    is_spike = surprise >= entropy_threshold #[cite: 1]
+    patch_mask = torch.zeros((max_patches, L), dtype=torch.bool, device=byte_ids.device) #[cite: 1]
+
+    curr_patch = 0 #[cite: 1]
+    patch_start = 0 #[cite: 1]
+    for t in range(L):
+        patch_mask[curr_patch, t] = True #[cite: 1]
+        reach_max = (t - patch_start + 1) >= max_patch_len #[cite: 1]
+        boundary = (is_spike[t] and (t > patch_start)) or reach_max #[cite: 1]
+
+        if boundary and (curr_patch < max_patches - 1): #[cite: 1]
+            curr_patch += 1 #[cite: 1]
+            patch_start = t + 1 #[cite: 1]
+
+    return patch_mask #[cite: 1]
+
+
+def blt_local_encoder_with_surprise(
+    params: dict,
+    byte_ids: torch.Tensor,       # [L]
+    entropy_threshold: float,
+    max_patches: int,
+) -> torch.Tensor:
+    """Maps raw byte tokens to [max_patches, d_global] using cross-attention pooling[cite: 1]."""
+    d_global = params["W_out_enc"].shape[0]
+    scale = d_global ** -0.5 #[cite: 1]
+
+    # Dynamically identify boundaries[cite: 1]
+    patch_mask = compute_entropy_and_patches(
+        byte_ids=byte_ids,
+        byte_embed=params["entropy_embed"],
+        W_lm=params["entropy_head"],
+        entropy_threshold=entropy_threshold,
+        max_patches=max_patches,
+    ) #[cite: 1]
+
+    # Intra-patch contextualization (1D convolution)[cite: 1]
+    x = F.embedding(byte_ids, params["byte_embed"]) #[cite: 1]
+    x_conv = x.transpose(0, 1).unsqueeze(0) #[cite: 1]
+    pad = (params["conv_w"].shape[-1] - 1) // 2 #[cite: 1]
+    x_conv = F.conv1d(x_conv, params["conv_w"], padding=pad, groups=x.shape[-1]) #[cite: 1]
+    x_ctx = F.silu(x_conv).squeeze(0).transpose(0, 1)  # [L, d_local][cite: 1]
+
+    # Projections to global space[cite: 1]
+    k = F.linear(x_ctx, params["W_k_enc"]) #[cite: 1]
+    v = F.linear(x_ctx, params["W_v_enc"]) #[cite: 1]
+    q = params["q_seed"]  # [max_patches, d_global][cite: 1]
+
+    # Cross-attention patch pooling[cite: 1]
+    scores = torch.matmul(q, k.transpose(0, 1)) * scale #[cite: 1]
+    empty_patches = ~patch_mask.any(dim=-1, keepdim=True) #[cite: 1]
+    mask_for_scores = patch_mask | empty_patches #[cite: 1]
+    scores = scores.masked_fill(~mask_for_scores, float("-inf")) #[cite: 1]
+
+    attn_weights = F.softmax(scores, dim=-1) #[cite: 1]
+    pooled = torch.matmul(attn_weights, v) #[cite: 1]
+    pooled = torch.where(empty_patches, torch.zeros_like(pooled), pooled) #[cite: 1]
+
+    return F.linear(pooled, params["W_out_enc"])
+
+
+# =====================================================================
+# 3. 96-Layer Global Backbone (MHA + QK-Norm + SwiGLU)
+# =====================================================================
+
+def narrow_deep_block(
+    x: torch.Tensor,              # [P, d][cite: 1]
+    p: dict,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+    num_heads: int,
+    head_dim: int,
+    res_scale: float,
+) -> torch.Tensor:
+    """Single transformer layer with depth-scaled residuals and QK-norm[cite: 1]."""
+    P, d = x.shape #[cite: 1]
+    scale = head_dim ** -0.5 #[cite: 1]
+
+    # --- Pre-LN MHA with QK-Norm ---[cite: 1]
+    h = rms_norm(x, p["norm1_w"]) #[cite: 1]
+
+    q = F.linear(h, p["W_q"]).view(P, num_heads, head_dim).transpose(0, 1)  # [H, P, d_k][cite: 1]
+    k = F.linear(h, p["W_k"]).view(P, num_heads, head_dim).transpose(0, 1)  # [H, P, d_k][cite: 1]
+    v = F.linear(h, p["W_v"]).view(P, num_heads, head_dim).transpose(0, 1)  # [H, P, d_k][cite: 1]
+
+    # Normalize queries and keys to prevent logit explosion across deep layers[cite: 1]
+    q = rms_norm(q, p["q_norm_w"]) #[cite: 1]
+    k = rms_norm(k, p["k_norm_w"]) #[cite: 1]
+
+    q = apply_rope(q, cos, sin) #[cite: 1]
+    k = apply_rope(k, cos, sin) #[cite: 1]
+
+    # Causal self-attention[cite: 1]
+    scores = torch.bmm(q, k.transpose(1, 2)) * scale #[cite: 1]
+    causal_mask = torch.triu(torch.full((P, P), float("-inf"), device=x.device), diagonal=1) #[cite: 1]
+    scores = scores + causal_mask.unsqueeze(0) #[cite: 1]
+    attn_weights = F.softmax(scores, dim=-1) #[cite: 1]
+
+    attn_out = torch.bmm(attn_weights, v).transpose(0, 1).contiguous().view(P, num_heads * head_dim) #[cite: 1]
+    h_attn = F.linear(attn_out, p["W_o"]) #[cite: 1]
+    x = x + (res_scale * h_attn) #[cite: 1]
+
+    # --- Pre-LN SwiGLU MLP ---[cite: 1]
+    h_mlp = rms_norm(x, p["norm2_w"]) #[cite: 1]
+    gate = F.linear(h_mlp, p["W_gate"]) #[cite: 1]
+    up = F.linear(h_mlp, p["W_up"]) #[cite: 1]
+    swiglu = F.silu(gate) * up #[cite: 1]
+    x_mlp = F.linear(swiglu, p["W_down"]) #[cite: 1]
+
+    return x + (res_scale * x_mlp) #[cite: 1]
+
+
+def global_96_layer_backbone_single(
+    patches: torch.Tensor,        # [P, d][cite: 1]
+    params: dict,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+    num_heads: int = 6,
+    head_dim: int = 48,
+    num_layers: int = 96,
+) -> torch.Tensor:
+    """Iterates through all 96 blocks with scaled residual updates[cite: 1]."""
+    h = patches #[cite: 1]
+    res_scale = (2.0 * num_layers) ** -0.5 #[cite: 1]
+
+    for i in range(num_layers): #[cite: 1]
+        block_p = params[f"layer_{i}"] #[cite: 1]
+        h = narrow_deep_block( #[cite: 1]
+            h, block_p, cos, sin,
+            num_heads=num_heads,
+            head_dim=head_dim,
+            res_scale=res_scale,
+        )
+
+    return rms_norm(h, params["final_norm_w"]) #[cite: 1]
+
+
+# =====================================================================
+# 4. Local Byte Decoder
+# =====================================================================
+
+def local_decoder_single(
+    byte_tokens: torch.Tensor,        # [L][cite: 1]
+    latent_patches: torch.Tensor,     # [P, d_global][cite: 1]
+    byte_to_patch_idx: torch.Tensor,  # [L][cite: 1]
+    params: dict,
+    d_local: int = 128,
+    num_heads: int = 4,
+) -> torch.Tensor:
+    """Cross-attends from bytes to global patches and runs local byte self-attention[cite: 1]."""
+    L = byte_tokens.shape[0] #[cite: 1]
+    P, d_global = latent_patches.shape #[cite: 1]
+    head_dim = d_local // num_heads #[cite: 1]
+    scale = head_dim ** -0.5 #[cite: 1]
+
+    # Byte Embedding Lookup[cite: 1]
+    x_bytes = F.embedding(byte_tokens, params["byte_embed"]) #[cite: 1]
+
+    # --- Patch Cross-Attention ---[cite: 1]
+    h_bytes = rms_norm(x_bytes, params["cross_norm_w"]) #[cite: 1]
+    q_b = F.linear(h_bytes, params["W_q_cross"]).view(L, num_heads, head_dim).transpose(0, 1) #[cite: 1]
+    k_p = F.linear(latent_patches, params["W_k_cross"]).view(P, num_heads, head_dim).transpose(0, 1) #[cite: 1]
+    v_p = F.linear(latent_patches, params["W_v_cross"]).view(P, num_heads, head_dim).transpose(0, 1) #[cite: 1]
+
+    cross_scores = torch.bmm(q_b, k_p.transpose(1, 2)) * scale #[cite: 1]
+
+    # Byte t cannot attend to future patches (p' > p)[cite: 1]
+    p_indices = torch.arange(P, device=byte_tokens.device).unsqueeze(0) #[cite: 1]
+    b_indices = byte_to_patch_idx.unsqueeze(1) #[cite: 1]
+    causal_patch_mask = (p_indices > b_indices).unsqueeze(0) #[cite: 1]
+    cross_scores = cross_scores.masked_fill(causal_patch_mask, float("-inf")) #[cite: 1]
+    cross_weights = F.softmax(cross_scores, dim=-1) #[cite: 1]
+
+    cross_out = torch.bmm(cross_weights, v_p).transpose(0, 1).contiguous().view(L, d_local) #[cite: 1]
+    x_bytes = x_bytes + F.linear(cross_out, params["W_o_cross"]) #[cite: 1]
+
+    # --- Local Intra-Patch Causal Self-Attention ---[cite: 1]
+    h_self = rms_norm(x_bytes, params["self_norm_w"]) #[cite: 1]
+    q_s = F.linear(h_self, params["W_q_self"]).view(L, num_heads, head_dim).transpose(0, 1) #[cite: 1]
+    k_s = F.linear(h_self, params["W_k_self"]).view(L, num_heads, head_dim).transpose(0, 1) #[cite: 1]
+    v_s = F.linear(h_self, params["W_v_self"]).view(L, num_heads, head_dim).transpose(0, 1) #[cite: 1]
+
+    self_scores = torch.bmm(q_s, k_s.transpose(1, 2)) * scale #[cite: 1]
+    causal_byte_mask = torch.triu(torch.full((L, L), float("-inf"), device=byte_tokens.device), diagonal=1) #[cite: 1]
+    self_scores = self_scores + causal_byte_mask.unsqueeze(0) #[cite: 1]
+    self_weights = F.softmax(self_scores, dim=-1) #[cite: 1]
+
+    self_out = torch.bmm(self_weights, v_s).transpose(0, 1).contiguous().view(L, d_local) #[cite: 1]
+    x_bytes = x_bytes + F.linear(self_out, params["W_o_self"]) #[cite: 1]
+
+    # --- Local SwiGLU MLP ---
+    h_mlp = rms_norm(x_bytes, params["dec_mlp_norm"])
+    gate = F.linear(h_mlp, params["W_gate_dec"])
+    up = F.linear(h_mlp, params["W_up_dec"])
+    x_bytes = x_bytes + F.linear(F.silu(gate) * up, params["W_down_dec"])
+
+    # --- Final Output Projection ---
+    return rms_norm(x_bytes, params["dec_final_norm"])

blt_trainer-0.1.0/src/blt_trainer/optim.py

@@ -0,0 +1,188 @@
+import torch
+import torch.nn.functional as F
+from torch.optim.optimizer import Optimizer
+
+# =====================================================================
+# 1. Cut Cross-Entropy (Chunked Loss)
+# =====================================================================
+
+def cut_cross_entropy(
+    h_final: torch.Tensor,       # [B, L, d_local]
+    lm_head: torch.Tensor,       # [256, d_local]
+    targets: torch.Tensor,       # [B, L]
+    chunk_size: int = 1024,
+) -> torch.Tensor:
+    """
+    Computes cross-entropy loss along the sequence dimension in slices,
+    preventing high memory consumption from materializing the entire logit tensor.
+    """
+    B, L, D = h_final.shape
+    h_flat = h_final.view(-1, D)
+    targets_flat = targets.view(-1)
+
+    total_tokens = h_flat.size(0)
+    total_loss = torch.tensor(0.0, device=h_final.device)
+
+    for start in range(0, total_tokens, chunk_size):
+        end = min(start + chunk_size, total_tokens)
+        h_chunk = h_flat[start:end]
+        t_chunk = targets_flat[start:end]
+
+        logits_chunk = F.linear(h_chunk, lm_head)  # [chunk, 256]
+        loss_chunk = F.cross_entropy(logits_chunk, t_chunk, reduction="sum")
+        total_loss = total_loss + loss_chunk
+
+    return total_loss / total_tokens
+
+
+# =====================================================================
+# 2. Muon Optimizer with ScheduleSkip (Newton-Schulz Orthogonalization)
+# =====================================================================
+
+def zeropower_via_newtonschulz5(G: torch.Tensor, steps: int = 5, eps: float = 1e-7) -> torch.Tensor:
+    """
+    Newton-Schulz iteration (degree 5) to compute the matrix sign/orthogonalization
+    of gradient tensors for Muon updates.
+    """
+    assert G.ndim >= 2
+    a, b, c = (3.4445, -4.7750, 2.0315)
+    X = G.bfloat16() if G.is_cuda else G.float()
+    X /= (X.norm() + eps)
+
+    transposed = False
+    if X.size(0) > X.size(1):
+        X = X.T
+        transposed = True
+
+    for _ in range(steps):
+        A = X @ X.T
+        B = b * A + c * (A @ A)
+        X = a * X + B @ X
+
+    if transposed:
+        X = X.T
+    return X.type_as(G)
+
+
+class MuonWithScheduleSkip(Optimizer):
+    """
+    Muon optimizer applied to hidden 2D projection matrices.
+    Implements schedule-free state tracking ('z') so learning rates do not
+    require fixed decay step counts.
+    """
+    def __init__(self, params, lr: float = 0.02, momentum: float = 0.95, weight_decay: float = 0.01):
+        defaults = dict(lr=lr, momentum=momentum, weight_decay=weight_decay)
+        super().__init__(params, defaults)
+
+    @torch.no_grad()
+    def step(self):
+        for group in self.param_groups:
+            lr = group["lr"]
+            mom = group["momentum"]
+            wd = group["weight_decay"]
+
+            for p in group["params"]:
+                if p.grad is None:
+                    continue
+                g = p.grad
+                state = self.state[p]
+
+                if "momentum_buffer" not in state:
+                    state["momentum_buffer"] = torch.zeros_like(g)
+                    state["z"] = p.clone()
+
+                buf = state["momentum_buffer"]
+                buf.mul_(mom).add_(g)
+
+                # Matrix orthogonalization
+                u = zeropower_via_newtonschulz5(buf)
+
+                # Weight decay applied directly to auxiliary sequence
+                state["z"].mul_(1.0 - lr * wd)
+                state["z"].add_(u, alpha=-lr)
+
+                p.copy_(state["z"])
+
+
+# =====================================================================
+# 3. Schedule-Free AdamW (for 1D Norms and Embeddings)
+# =====================================================================
+
+class ScheduleFreeAdamW(Optimizer):
+    """
+    AdamW variant that maintains momentum buffers without requiring 
+    a scheduled cosine learning rate multiplier.
+    """
+    def __init__(
+        self,
+        params,
+        lr: float = 1e-3,
+        betas: tuple[float, float] = (0.9, 0.999),
+        eps: float = 1e-8,
+        weight_decay: float = 0.0,
+    ):
+        defaults = dict(lr=lr, betas=betas, eps=eps, weight_decay=weight_decay)
+        super().__init__(params, defaults)
+
+    @torch.no_grad()
+    def step(self):
+        for group in self.param_groups:
+            lr = group["lr"]
+            beta1, beta2 = group["betas"]
+            eps = group["eps"]
+            wd = group["weight_decay"]
+
+            for p in group["params"]:
+                if p.grad is None:
+                    continue
+                g = p.grad
+                state = self.state[p]
+
+                if "step" not in state:
+                    state["step"] = 0
+                    state["exp_avg"] = torch.zeros_like(g)
+                    state["exp_avg_sq"] = torch.zeros_like(g)
+                    state["z"] = p.clone()
+
+                state["step"] += 1
+                exp_avg, exp_avg_sq = state["exp_avg"], state["exp_avg_sq"]
+
+                exp_avg.mul_(beta1).add_(g, alpha=1.0 - beta1)
+                exp_avg_sq.mul_(beta2).addcmul_(g, g, value=1.0 - beta2)
+
+                denom = exp_avg_sq.sqrt().add_(eps)
+                step_size = lr / (1.0 - beta1 ** state["step"])
+
+                state["z"].mul_(1.0 - lr * wd)
+                state["z"].addcdiv_(exp_avg, denom, value=-step_size)
+
+                p.copy_(state["z"])
+
+
+# =====================================================================
+# 4. Optimizer Parameter Routing
+# =====================================================================
+
+def partition_params_for_optimizers(params: dict) -> tuple[list[torch.Tensor], list[torch.Tensor]]:
+    """
+    Traverses nested parameter dictionaries and separates parameters:
+    - 2D internal weight matrices -> Muon
+    - 1D normalization gains, biases, and embedding tables -> AdamW
+    """
+    muon_tensors = []
+    adam_tensors = []
+
+    def _collect(subdict, prefix=""):
+        for k, v in subdict.items():
+            full_name = f"{prefix}.{k}" if prefix else k
+            if isinstance(v, dict):
+                _collect(v, full_name)
+            elif isinstance(v, torch.Tensor):
+                # Only 2D hidden matrices route to Muon
+                if v.ndim == 2 and not any(tag in full_name for tag in ["embed", "lm_head", "entropy"]):
+                    muon_tensors.append(v)
+                else:
+                    adam_tensors.append(v)
+
+    _collect(params)
+    return muon_tensors, adam_tensors