Upload folder using huggingface_hub

README.md

@@ -1,13 +1,48 @@
 ---
 title: HobbyLM Playground
-emoji: 🏆
-colorFrom: red
+emoji: 🪶
+colorFrom: indigo
 colorTo: pink
 sdk: gradio
-sdk_version: 6.19.0
-python_version: '3.13'
+sdk_version: 4.44.1
 app_file: app.py
 pinned: false
+license: apache-2.0
+short_description: Chat, see & generate with the 500M HobbyLM MoE family
+models:
+  - rootxhacker/HobbyLM-Base
+  - rootxhacker/HobbyLM-Chat
+  - rootxhacker/HobbyLM-Computer-Use
+  - rootxhacker/HobbyLM-Omni
+  - rootxhacker/HobbyLM-Diffusion
+  - rootxhacker/HobbyLM-Image
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# 🪶 HobbyLM Playground
+
+An interactive demo of **HobbyLM** — a from-scratch **500M sparse Mixture-of-Experts** language-model family
+(plus a 333M text-to-image DiT), all trained on a hobby budget. One Space, three things to try:
+
+- **💬 Chat** — talk to any variant: Base, Chat, Computer-Use, the multimodal Omni core, or the
+  masked-diffusion model (which decodes by iterative denoising, not left-to-right).
+- **🖼️ Ask about an image** — upload a picture and question the multimodal **Omni** model (SigLIP2 vision
+  encoder → MoE LLM).
+- **🎨 Generate an image** — text-to-image at 1024px with **HobbyLM-Image** (a flow-matching DiT in the
+  DC-AE latent space, conditioned on CLIP-L).
+
+The models use a custom `hobbylm` architecture, so this Space vendors the small reference implementation
+(`hobbylm/`, `hobby_image/`) rather than going through `transformers.AutoModel`.
+
+## Hardware
+
+This Space is written for **ZeroGPU** (the heavy functions are decorated with `@spaces.GPU`). Enable
+*ZeroGPU* in the Space's hardware settings for fast chat, image understanding, and 1024px generation. It
+also runs on CPU (chat is slow; image generation is impractical there).
+
+## Links
+
+- Models: <https://huggingface.co/rootxhacker>
+- Code + the from-scratch Rust CPU engine: <https://github.com/harishsg993010/HobbyLM>
+
+These are tiny research models — genuinely fluent and fun, but with the capability ceiling of a 500M model
+(hallucination, weak strict-format following, soft hands / multi-person in image generation). Apache-2.0.

app.py

@@ -0,0 +1,300 @@
+"""HobbyLM Playground — a Gradio Space to chat with the HobbyLM models, ask questions about an
+image (the multimodal Omni model), and generate images (the 1024px DiT + DC-AE pipeline).
+
+All models are the from-scratch 500M sparse-MoE family (+ a 333M image DiT) published at
+https://huggingface.co/rootxhacker . They use a custom architecture, so the Space vendors the
+reference implementation (`hobbylm/`, `hobby_image/`) instead of going through transformers' AutoModel.
+
+Runs on ZeroGPU (the heavy functions are @spaces.GPU); falls back to CPU when run locally.
+"""
+import json
+import threading
+
+import gradio as gr
+import torch
+from huggingface_hub import hf_hub_download
+from safetensors.torch import load_file
+
+# ZeroGPU decorator — with a no-op fallback so the app also runs on plain CPU / locally.
+try:
+    import spaces
+except Exception:  # not on a ZeroGPU Space
+    class _Spaces:
+        @staticmethod
+        def GPU(*a, **k):
+            if a and callable(a[0]):
+                return a[0]
+            def deco(f):
+                return f
+            return deco
+    spaces = _Spaces()
+
+HF_USER = "rootxhacker"
+VISION_ID = "google/siglip2-so400m-patch16-512"   # the encoder HobbyLM-Omni was trained with
+DCAE_ID = "mit-han-lab/dc-ae-f32c32-sana-1.1-diffusers"
+CLIP_ID = "openai/clip-vit-large-patch14"
+NEG_DEFAULT = "blurry, low quality, watermark, signature, text, jpeg artifacts, deformed, distorted"
+
+# chat dropdown -> (repo suffix, decoding kind)
+CHAT_MODELS = {
+    "HobbyLM-Chat — instruction / conversation": ("HobbyLM-Chat", "chat"),
+    "HobbyLM-Base — raw text completion": ("HobbyLM-Base", "base"),
+    "HobbyLM-Computer-Use — tools / GUI agent": ("HobbyLM-Computer-Use", "chat"),
+    "HobbyLM-Omni — multimodal core (text)": ("HobbyLM-Omni", "chat"),
+    "HobbyLM-Diffusion — masked-diffusion LM": ("HobbyLM-Diffusion", "diffusion"),
+}
+
+_cache = {}
+_lock = threading.Lock()
+
+
+def _device():
+    return "cuda" if torch.cuda.is_available() else "cpu"
+
+
+def _enc():
+    import tiktoken
+    return tiktoken.get_encoding("gpt2")
+
+
+# --------------------------------------------------------------------------- loaders (cached)
+
+def _load_llm(repo, device):
+    key = ("llm", repo)
+    with _lock:
+        if key not in _cache:
+            from hobbylm.config import ModelConfig
+            from hobbylm.model import MoETransformer
+            cfg_d = {k: v for k, v in json.load(
+                open(hf_hub_download(f"{HF_USER}/{repo}", "config.json"))).items() if k != "preset"}
+            cfg = ModelConfig(**cfg_d)
+            cfg.expert_backend = "bmm"                      # universal MoE backend (CPU + GPU)
+            model = MoETransformer(cfg).eval()
+            model.load_state_dict(load_file(hf_hub_download(f"{HF_USER}/{repo}", "model.safetensors")))
+            _cache[key] = (model, cfg)
+    model, cfg = _cache[key]
+    model.to(device)
+    return model, cfg
+
+
+def _load_vlm(device):
+    key = ("vlm",)
+    with _lock:
+        if key not in _cache:
+            from hobbylm.vision import SiglipVision
+            from hobbylm.multimodal import MoEVLM
+            llm, _ = _load_llm("HobbyLM-Omni", device)
+            enc = SiglipVision(model_id=VISION_ID, device=device, dtype=torch.float32)
+            vlm = MoEVLM(llm, vision_dim=enc.hidden)
+            vlm.mm_projector.load_state_dict(
+                load_file(hf_hub_download(f"{HF_USER}/HobbyLM-Omni", "vision_projector.safetensors")))
+            vlm.eval()
+            _cache[key] = (vlm, enc)
+    vlm, enc = _cache[key]
+    vlm.to(device)
+    return vlm, enc
+
+
+def _load_image_models(device):
+    with _lock:
+        if ("dit",) not in _cache:
+            from hobby_image.dit import HobbyImageDiT, DiTConfig
+            cfg = json.load(open(hf_hub_download(f"{HF_USER}/HobbyLM-Image", "config.json")))
+            dit = HobbyImageDiT(DiTConfig(**cfg["dit_config"])).eval()
+            dit.load_state_dict(load_file(hf_hub_download(f"{HF_USER}/HobbyLM-Image", "model.safetensors")))
+            _cache[("dit",)] = (dit, cfg["dit_config"]["latent_h"], float(cfg["lat_std"]), float(cfg["scaling_factor"]))
+        if ("dcae",) not in _cache:
+            from diffusers import AutoencoderDC
+            _cache[("dcae",)] = AutoencoderDC.from_pretrained(DCAE_ID, torch_dtype=torch.float16).eval()
+        if ("clip",) not in _cache:
+            from transformers import CLIPTextModel, CLIPTokenizer
+            _cache[("clip",)] = (CLIPTokenizer.from_pretrained(CLIP_ID),
+                                 CLIPTextModel.from_pretrained(CLIP_ID, torch_dtype=torch.float16).eval())
+    dit, lat, lat_std, sf = _cache[("dit",)]
+    ae = _cache[("dcae",)]
+    tok, clip = _cache[("clip",)]
+    dit.to(device)
+    ae.to(device)
+    clip.to(device)
+    return dit, lat, lat_std, sf, ae, tok, clip
+
+
+# --------------------------------------------------------------------------- chat
+
+def _build_prompt(repo, message, history):
+    if repo == "HobbyLM-Base":
+        return message                                     # base = pure completion
+    s = ""
+    for turn in history or []:
+        if isinstance(turn, (list, tuple)) and len(turn) == 2:
+            u, a = turn
+            if u:
+                s += f"USER: {u}\n"
+            if a:
+                s += f"ASSISTANT: {a}\n"
+    return s + f"USER: {message}\nASSISTANT:"
+
+
+@spaces.GPU(duration=120)
+def chat_fn(message, history, model_name, max_new, temperature):
+    from hobbylm.generate import generate as ar_generate
+    repo, kind = CHAT_MODELS[model_name]
+    dev = _device()
+    enc = _enc()
+    prompt = _build_prompt(repo, message, history)
+
+    if kind == "diffusion":
+        from hobbylm.diffusion import generate as dgen
+        model, _ = _load_llm(repo, dev)
+        ids = torch.tensor([enc.encode_ordinary(prompt)], device=dev)
+        gen_len = int(max_new)
+        out = dgen(model, ids, gen_len=gen_len, steps=max(32, 2 * gen_len),
+                   temperature=max(0.0, float(temperature) - 0.4), rep_penalty=1.5, remask_steps=2)
+        return enc.decode(out[0].tolist()).strip()
+
+    model, cfg = _load_llm(repo, dev)
+    ids = torch.tensor([enc.encode_ordinary(prompt)], device=dev)
+    ctx_len = min(getattr(cfg, "context_length", 1024), 2048)
+    out = ar_generate(model, ids, int(max_new), float(temperature), 0, torch.device(dev),
+                      top_p=0.95, repetition_penalty=1.3, no_repeat_ngram_size=3, ctx_len=ctx_len)
+    return enc.decode(out[0, ids.shape[1]:].tolist()).strip()
+
+
+# --------------------------------------------------------------------------- image understanding (Omni)
+
+@spaces.GPU(duration=120)
+def understand_fn(image, question, max_new):
+    if image is None:
+        return "Please upload an image first."
+    from hobbylm.multimodal import IMAGE_TOKEN
+    from hobbylm.generate import GPT2_VALID, EOT
+    dev = _device()
+    enc = _enc()
+    vlm, venc = _load_vlm(dev)
+
+    with torch.no_grad():
+        feats = venc.encode([image.convert("RGB")]).float()
+        q = (question or "Describe this image in detail.").strip()
+        pre = enc.encode_ordinary(f"USER: {q}\nASSISTANT:")
+        ids = torch.tensor([[IMAGE_TOKEN] + pre], device=dev)
+        cur, _ = vlm.build_inputs_embeds(ids, image_features=feats)
+        outs = []
+        for _ in range(int(max_new)):
+            logits, _ = vlm.llm(inputs_embeds=cur)
+            lg = logits[:, -1, :].float()
+            lg[:, GPT2_VALID:] = -float("inf")
+            if outs:                                        # repetition penalty
+                u = torch.tensor(sorted(set(outs)), device=dev)
+                v = lg[0, u]
+                lg[0, u] = torch.where(v > 0, v / 1.3, v * 1.3)
+            t = int(lg.argmax(-1).item())
+            if t == EOT:
+                break
+            outs.append(t)
+            e = vlm.llm.embed(torch.tensor([[t]], device=dev)).to(cur.dtype)
+            cur = torch.cat([cur, e], dim=1)
+    return enc.decode(outs).strip() or "(no answer)"
+
+
+# --------------------------------------------------------------------------- image generation
+
+@spaces.GPU(duration=180)
+def generate_image_fn(prompt, negative, steps, guidance, seed, progress=gr.Progress()):
+    if not prompt or not prompt.strip():
+        raise gr.Error("Enter a prompt.")
+    from PIL import Image
+    import numpy as np
+    dev = _device()
+    dit, lat, lat_std, sf, ae, tok, clip = _load_image_models(dev)
+    steps = int(steps)
+    neg = (negative or "").strip()
+
+    def clip_encode(texts):
+        ids = tok(texts, padding="max_length", max_length=64, truncation=True,
+                  return_tensors="pt").input_ids.to(dev)
+        with torch.no_grad():
+            return clip(ids).last_hidden_state.float()
+
+    g = torch.Generator(device=dev).manual_seed(int(seed))
+    ctx = clip_encode([prompt])
+    uncond = clip_encode([neg]) if neg else torch.zeros_like(ctx)
+    task = torch.zeros(1, dtype=torch.long, device=dev)
+    z = torch.randn(1, 32, lat, lat, generator=g, device=dev)
+    zs = torch.zeros(1, 32, lat, lat, device=dev)
+    em = torch.zeros(1, 1, lat, 2 * lat, device=dev)
+    amp = torch.autocast("cuda", dtype=torch.float16) if dev == "cuda" else torch.autocast("cpu", enabled=False)
+    with torch.no_grad():
+        for i in progress.tqdm(range(steps), desc="denoising"):
+            tt = torch.full((1,), i / steps, device=dev)
+            inp = torch.cat([torch.cat([z, zs], dim=-1), em, em], dim=1)
+            with amp:
+                vc = dit(inp, tt, ctx, task)[..., :lat].float()
+                vu = dit(inp, tt, uncond, task)[..., :lat].float()
+            z = z + (vu + float(guidance) * (vc - vu)) / steps
+        with amp:
+            img = ae.decode((z * lat_std / sf).half() if dev == "cuda" else (z * lat_std / sf)).sample
+    img = img.float().clamp(-1, 1)[0]
+    arr = ((img.permute(1, 2, 0).cpu().numpy() + 1) * 127.5).clip(0, 255).astype(np.uint8)
+    return Image.fromarray(arr)
+
+
+# --------------------------------------------------------------------------- UI
+
+INTRO = """# 🪶 HobbyLM Playground
+
+A from-scratch **500M sparse Mixture-of-Experts** model family (+ a 333M image DiT), trained on a hobby
+budget. Chat with any variant, ask questions about an image with the multimodal **Omni** model, or
+generate a 1024px image. Models: [rootxhacker on Hugging Face](https://huggingface.co/rootxhacker) ·
+code: [GitHub](https://github.com/harishsg993010/HobbyLM).
+
+*These are tiny research models — fluent and fun, with the capability ceiling of a 500M model.*
+"""
+
+with gr.Blocks(title="HobbyLM Playground", theme=gr.themes.Soft()) as demo:
+    gr.Markdown(INTRO)
+
+    with gr.Tab("💬 Chat"):
+        model_dd = gr.Dropdown(list(CHAT_MODELS), value=list(CHAT_MODELS)[0], label="Model")
+        with gr.Row():
+            max_new = gr.Slider(16, 512, value=200, step=8, label="Max new tokens")
+            temp = gr.Slider(0.0, 1.5, value=0.7, step=0.05, label="Temperature (0 = greedy)")
+        gr.ChatInterface(
+            fn=chat_fn,
+            additional_inputs=[model_dd, max_new, temp],
+            examples=[["Give me three tips for better sleep."],
+                      ["Explain a mixture-of-experts model in one sentence."],
+                      ["Write a short poem about the ocean."]],
+        )
+
+    with gr.Tab("🖼️ Ask about an image"):
+        gr.Markdown("Upload an image and ask the **HobbyLM-Omni** multimodal model about it.")
+        with gr.Row():
+            with gr.Column():
+                u_img = gr.Image(type="pil", label="Image")
+                u_q = gr.Textbox(label="Question", value="Describe this image in detail.")
+                u_max = gr.Slider(16, 128, value=48, step=8, label="Max new tokens")
+                u_btn = gr.Button("Ask", variant="primary")
+            u_out = gr.Textbox(label="Answer", lines=6)
+        u_btn.click(understand_fn, [u_img, u_q, u_max], u_out)
+
+    with gr.Tab("🎨 Generate an image"):
+        gr.Markdown("Text-to-image with **HobbyLM-Image** (1024px DiT in DC-AE latent space). "
+                    "Strongest on single objects and cinematic scenes.")
+        with gr.Row():
+            with gr.Column():
+                g_prompt = gr.Textbox(label="Prompt", value="a red convertible car on a coastal road, golden hour")
+                g_neg = gr.Textbox(label="Negative prompt", value=NEG_DEFAULT)
+                with gr.Row():
+                    g_steps = gr.Slider(20, 120, value=60, step=5, label="Steps")
+                    g_cfg = gr.Slider(1.0, 10.0, value=5.0, step=0.5, label="Guidance (CFG)")
+                    g_seed = gr.Number(value=1234, label="Seed", precision=0)
+                g_btn = gr.Button("Generate", variant="primary")
+            g_out = gr.Image(label="Result", height=512)
+        g_btn.click(generate_image_fn, [g_prompt, g_neg, g_steps, g_cfg, g_seed], g_out)
+        gr.Examples([["a photograph of a single red apple on a plain white background", NEG_DEFAULT, 60, 5.0, 1234],
+                     ["a cozy library with tall wooden bookshelves, warm light", NEG_DEFAULT, 80, 5.0, 7],
+                     ["a bowl of fresh strawberries, studio food photography", NEG_DEFAULT, 60, 5.0, 42]],
+                    [g_prompt, g_neg, g_steps, g_cfg, g_seed])
+
+if __name__ == "__main__":
+    demo.queue(max_size=20).launch()

hobby_image/__init__.py

@@ -0,0 +1 @@
+"""hobby_image — the HobbyLM-Image text-to-image DiT (vendored for the Space)."""

hobby_image/dit.py

@@ -0,0 +1,173 @@
+"""DreamLite in-context DiT on a DC-AE f32c32 latent — the $300 path.
+
+DC-AE compresses 32x spatially (vs VAE f8's 8x), so a 512px image -> 16x16x32 latent. The two-panel
+canvas (target|source, width-concat) is 16x32x32 = ONLY 512 tokens (vs the conv U-Net's 64x128=8,192
+positions). A DiT on 512 tokens is ~order-of-magnitude cheaper to train than the conv U-Net, which is
+how the budget drops from ~$1.66k toward ~$300. Trade-off: f32 reconstruction is slightly weaker than
+f8c16 on fine text (acceptable for the V0 "global + simple-local edits" scope).
+
+Standard DiT (AdaLN-Zero on self-attn + MLP) + cross-attention to the frozen-VLM refiner tokens, flow-
+matching velocity. Two-panel handled by a learned left/right panel embedding; loss uses the left half.
+"""
+from __future__ import annotations
+
+import math
+from dataclasses import dataclass
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch import Tensor
+
+
+@dataclass
+class DiTConfig:
+    in_channels: int = 34      # 32 DC-AE latent ch + 2 optional mask ch
+    out_channels: int = 32
+    latent_h: int = 16         # 512px / 32 = 16 (256px -> 8)
+    panel_w: int = 16          # per-panel width; canvas width = 2*panel_w
+    patch: int = 1             # DC-AE latent is already compressed -> patch 1
+    d_model: int = 1024
+    depth: int = 16
+    heads: int = 16
+    mlp_ratio: float = 3.0
+    ctx_dim: int = 1024
+    n_tasks: int = 3
+
+
+def sinusoidal(t: Tensor, dim: int) -> Tensor:
+    half = dim // 2
+    f = torch.exp(-math.log(10000.0) * torch.arange(half, device=t.device) / half)
+    a = t.float()[:, None] * f[None, :]
+    return torch.cat([a.cos(), a.sin()], dim=-1)
+
+
+def modulate(x: Tensor, shift: Tensor, scale: Tensor) -> Tensor:
+    return x * (1 + scale[:, None]) + shift[:, None]
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, dim: int, eps: float = 1e-6):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(dim))
+        self.eps = eps
+
+    def forward(self, x: Tensor) -> Tensor:
+        return F.rms_norm(x, (x.size(-1),), self.weight, self.eps)
+
+
+class Attention(nn.Module):
+    def __init__(self, dim: int, heads: int, ctx_dim: int | None = None, qk_norm: bool = True):
+        super().__init__()
+        self.heads = heads
+        self.hd = dim // heads
+        self.cross = ctx_dim is not None
+        self.q = nn.Linear(dim, dim, bias=False)
+        self.kv = nn.Linear(ctx_dim or dim, 2 * dim, bias=False)
+        self.o = nn.Linear(dim, dim, bias=False)
+        # QK-norm: RMSNorm on Q,K (over head_dim) BEFORE the dot product — the key training-stability
+        # fix for DiTs (Inf-DiT/Lumina-Next/SD3); bounds the attention logits, prevents divergence.
+        self.qn = RMSNorm(self.hd) if qk_norm else nn.Identity()
+        self.kn = RMSNorm(self.hd) if qk_norm else nn.Identity()
+
+    def forward(self, x: Tensor, ctx: Tensor | None = None) -> Tensor:
+        B, N, _ = x.shape
+        src = ctx if self.cross else x
+        q = self.qn(self.q(x).view(B, N, self.heads, self.hd)).transpose(1, 2)
+        kv = self.kv(src).view(B, src.shape[1], 2, self.heads, self.hd)
+        k = self.kn(kv[:, :, 0]).transpose(1, 2)
+        v = kv[:, :, 1].transpose(1, 2)
+        o = F.scaled_dot_product_attention(q, k, v)
+        return self.o(o.transpose(1, 2).reshape(B, N, -1))
+
+
+class DiTBlock(nn.Module):
+    def __init__(self, cfg: DiTConfig):
+        super().__init__()
+        d = cfg.d_model
+        self.ln1 = nn.LayerNorm(d, elementwise_affine=False, eps=1e-6)
+        self.sa = Attention(d, cfg.heads)
+        self.lnc = nn.LayerNorm(d, eps=1e-6)
+        self.ca = Attention(d, cfg.heads, ctx_dim=cfg.ctx_dim)
+        self.ln2 = nn.LayerNorm(d, elementwise_affine=False, eps=1e-6)
+        h = int(d * cfg.mlp_ratio)
+        self.mlp = nn.Sequential(nn.Linear(d, h), nn.GELU(approximate="tanh"), nn.Linear(h, d))
+        self.adaln = nn.Sequential(nn.SiLU(), nn.Linear(d, 6 * d))
+        nn.init.zeros_(self.adaln[-1].weight); nn.init.zeros_(self.adaln[-1].bias)
+
+    def forward(self, x: Tensor, cond: Tensor, ctx: Tensor) -> Tensor:
+        sa_s, sa_sc, sa_g, ml_s, ml_sc, ml_g = self.adaln(cond).chunk(6, dim=-1)
+        x = x + sa_g[:, None] * self.sa(modulate(self.ln1(x), sa_s, sa_sc))
+        x = x + self.ca(self.lnc(x), ctx)                       # cross-attn to VLM tokens
+        x = x + ml_g[:, None] * self.mlp(modulate(self.ln2(x), ml_s, ml_sc))
+        return x
+
+
+class HobbyImageDiT(nn.Module):
+    def __init__(self, cfg: DiTConfig):
+        super().__init__()
+        self.cfg = cfg
+        d = cfg.d_model
+        self.canvas_w = 2 * cfg.panel_w
+        self.n_tokens = (cfg.latent_h // cfg.patch) * (self.canvas_w // cfg.patch)
+        self.patch_embed = nn.Conv2d(cfg.in_channels, d, cfg.patch, stride=cfg.patch)
+        self.pos = nn.Parameter(torch.zeros(1, self.n_tokens, d))
+        self.panel = nn.Parameter(torch.zeros(2, d))            # left/right
+        self.t_mlp = nn.Sequential(nn.Linear(d, d), nn.SiLU(), nn.Linear(d, d))
+        self.task_emb = nn.Embedding(cfg.n_tasks, d)
+        self.blocks = nn.ModuleList([DiTBlock(cfg) for _ in range(cfg.depth)])
+        self.ln_f = nn.LayerNorm(d, elementwise_affine=False, eps=1e-6)
+        self.adaln_f = nn.Sequential(nn.SiLU(), nn.Linear(d, 2 * d))
+        nn.init.zeros_(self.adaln_f[-1].weight); nn.init.zeros_(self.adaln_f[-1].bias)
+        self.head = nn.Linear(d, cfg.patch * cfg.patch * cfg.out_channels)
+        nn.init.zeros_(self.head.weight); nn.init.zeros_(self.head.bias)
+        nn.init.trunc_normal_(self.pos, std=0.02)
+
+    def forward(self, x: Tensor, t: Tensor, ctx: Tensor, task: Tensor | None = None) -> Tensor:
+        """x: (B, in_ch, H, 2*panel_w) two-panel DC-AE latent. t: (B,). ctx: (B, M, ctx_dim).
+        Returns velocity (B, out_ch, H, 2*panel_w)."""
+        cfg = self.cfg
+        B, _, H, W = x.shape
+        h = self.patch_embed(x).flatten(2).transpose(1, 2)      # (B, N, d)
+        h = h + self.pos
+        # panel embedding: tokens in the left half vs right half (by column)
+        gw = W // cfg.patch
+        col = torch.arange(self.n_tokens, device=x.device) % gw
+        left = (col < gw // 2)
+        h = h + torch.where(left[None, :, None], self.panel[0], self.panel[1])
+        cond = self.t_mlp(sinusoidal(t * 1000.0, cfg.d_model))
+        if task is not None:
+            cond = cond + self.task_emb(task)
+        for blk in self.blocks:
+            h = blk(h, cond, ctx)
+        s, sc = self.adaln_f(cond).chunk(2, dim=-1)
+        h = modulate(self.ln_f(h), s, sc)
+        h = self.head(h)                                        # (B, N, p*p*out)
+        # unpatchify
+        gh = H // cfg.patch
+        h = h.view(B, gh, gw, cfg.patch, cfg.patch, cfg.out_channels)
+        h = h.permute(0, 5, 1, 3, 2, 4).reshape(B, cfg.out_channels, H, W)
+        return h
+
+
+def count_params(m: nn.Module) -> int:
+    return sum(p.numel() for p in m.parameters())
+
+
+# 512px: DC-AE f32 -> 16x16x32 per panel; two-panel canvas 16x32; 512 tokens.
+V0_DCAE_512 = DiTConfig(in_channels=34, out_channels=32, latent_h=16, panel_w=16, patch=1,
+                        d_model=1024, depth=16, heads=16, mlp_ratio=3.0, ctx_dim=1024)
+# 256px pilot: 8x8x32 per panel; canvas 8x16; 128 tokens.
+V0_DCAE_256 = DiTConfig(in_channels=34, out_channels=32, latent_h=8, panel_w=8, patch=1,
+                        d_model=1024, depth=16, heads=16, mlp_ratio=3.0, ctx_dim=1024)
+
+
+if __name__ == "__main__":
+    for name, cfg in [("512", V0_DCAE_512), ("256", V0_DCAE_256)]:
+        m = HobbyImageDiT(cfg)
+        x = torch.randn(2, cfg.in_channels, cfg.latent_h, 2 * cfg.panel_w)
+        t = torch.rand(2); ctx = torch.randn(2, 256, cfg.ctx_dim); task = torch.zeros(2, dtype=torch.long)
+        with torch.no_grad():
+            y = m(x, t, ctx, task)
+        print(f"DiT {name}: {count_params(m)/1e6:.1f}M params | {m.n_tokens} tokens | "
+              f"in {tuple(x.shape)} -> out {tuple(y.shape)}")

hobbylm/__init__.py

@@ -0,0 +1,6 @@
+"""hobbylm — core library package for the MoE LLM project.
+
+Contains the model/MoE definitions, configs, optimizers, data pipelines,
+multimodal encoders, generation/decoding, and tool/eval utilities that are
+imported by the training, export, agents, eval, and Modal pipeline scripts.
+"""

hobbylm/config.py

@@ -0,0 +1,142 @@
+"""Model + training configuration for the MoE lab.
+
+A single ModelConfig drives the architecture; ablation knobs are explicit fields so an
+ablation = one config override. PRESETS hold the 130M / 500M / 1B starting points
+(exact dims are validated against targets by count_params.py).
+"""
+from __future__ import annotations
+
+from dataclasses import dataclass, field, asdict
+from typing import Literal
+
+
+@dataclass
+class ModelConfig:
+    # ---- shape ----
+    vocab_size: int = 50304            # GPT-2 (50257) padded to mult of 128
+    d_model: int = 512
+    n_layers: int = 12
+    n_dense_layers: int = 1            # first-k layers use a dense FFN (DeepSeekMoE-16B style)
+    # ---- attention (GQA + QK-norm) ----
+    n_q_heads: int = 8
+    n_kv_heads: int = 2                # GQA group = n_q_heads / n_kv_heads
+    head_dim: int = 64                 # may be decoupled from d_model (Qwen3 trick)
+    qk_norm: bool = True               # RMSNorm per-head on Q,K before RoPE
+    rope_theta: float = 10000.0
+    attn_softcap: float = 0.0          # 0 disables; logit softcap fallback if not using qk_norm
+    # ---- FFN / experts (SwiGLU) ----
+    dense_ffn: int = 1536              # intermediate dim of the dense FFN layers
+    expert_ffn: int = 256              # intermediate dim of ONE routed/shared expert (fine-grained)
+    n_experts: int = 32                # routed experts per MoE layer
+    top_k: int = 4                     # active routed experts per token
+    n_shared: int = 0                  # always-on shared experts (0 or 1 at this scale)
+    # ---- routing / balancing (ablation forks) ----
+    gating: Literal["sigmoid", "softmax"] = "sigmoid"
+    norm_topk_prob: bool = False       # ablation winner: NOT renormalizing top-k gates (OLMoE-style) beat renorm by -0.015
+    balancing: Literal["aux_free", "aux_loss"] = "aux_free"
+    aux_loss_coef: float = 1e-3        # global-batch balance loss (safety net w/ aux_free; 1e-2 if aux_loss only)
+    z_loss_coef: float = 1e-3          # router z-loss
+    bias_update_rate: float = 1e-3     # aux-loss-free per-expert bias step (u); annealed to 0 near end
+    router_init_std: float = 0.02      # router gets a small (0.1x-ish) init; see model init
+    # ---- embeddings / output ----
+    tie_embeddings: bool = True
+    scale_embeddings: bool = False     # multiply token embeds by sqrt(d_model) (Gemma)
+    final_z_loss_coef: float = 1e-4    # z-loss on final logits
+    logit_softcap: float = 0.0         # 0 disables; e.g. 15-30 (Gemma-2)
+    # ---- diffusion conversion (LLaDA/MDLM: bidirectional attn + masked-token objective) ----
+    diffusion: bool = False            # True => full bidirectional attention + masked-diffusion loss (no AR)
+    mask_token_id: int = 50257         # free sentinel in the 50304-padded vocab (>= GPT2_VALID, never an AR token)
+    mask_eps: float = 1e-3             # min mask ratio: t ~ U(eps, 1) per sequence
+    # ---- multi-token prediction (optional) ----
+    n_mtp: int = 0                     # 0 disables; 1 = predict t+2 with one extra head
+    mtp_weight: float = 0.1
+    # ---- init ----
+    init_std: float = 0.02
+    # ---- impl ----
+    expert_backend: Literal["grouped", "bmm", "loop"] = "grouped"  # grouped=GPU fast; bmm/loop=CPU-testable ref
+    # ---- throughput opts (nanogpt-inspired; see docs/ARCHITECTURE_RESEARCH.md §8) ----
+    fused_ce: bool = False             # chunked cross-entropy: never materialize the full (T,vocab) fp32 logits
+    ce_chunk: int = 4096               # rows per CE chunk (fused_ce only)
+    fp8_head: bool = False             # EXPERIMENTAL/BROKEN: FP8 lm_head (untied). No speedup at 1B + zero-grad
+                                       # backward (loss frozen at init in the 130M ablation). Do not use.
+    fp8_x_scale: float = 1.0           # fp8 activation/weight/grad scales (pre-head x is ~unit RMS, so 1.0 is safe)
+    fp8_w_scale: float = 1.0
+    fp8_grad_scale: float = 1.0
+
+    def __post_init__(self):
+        assert self.n_q_heads % self.n_kv_heads == 0, "n_q_heads must be divisible by n_kv_heads"
+        assert self.n_dense_layers <= self.n_layers
+        assert self.top_k <= self.n_experts
+
+    @property
+    def n_moe_layers(self) -> int:
+        return self.n_layers - self.n_dense_layers
+
+    def to_dict(self) -> dict:
+        return asdict(self)
+
+
+@dataclass
+class TrainConfig:
+    # data
+    data_dir: str = "data/fineweb10B"
+    train_pattern: str = "fineweb_train_*.bin"
+    val_pattern: str = "fineweb_val_*.bin"
+    seq_len: int = 1024
+    batch_tokens: int = 256 * 1024     # tokens per optimizer step (global)
+    micro_batch_seqs: int = 16         # sequences per micro-batch per GPU
+    # schedule
+    max_steps: int = 4000
+    warmup_steps: int = 100
+    cooldown_frac: float = 0.4         # last fraction of steps linearly decays lr -> lr*final_frac
+    final_lr_frac: float = 0.1
+    # optimizer
+    muon_lr: float = 0.02
+    muon_momentum: float = 0.95
+    muon_wd: float = 0.1
+    muon_ns_steps: int = 5
+    orthogonalizer: Literal["ns5", "polar"] = "ns5"  # ns5=Newton-Schulz (baseline); polar=Polar Express schedule
+    adam_lr: float = 3e-4
+    adam_betas: tuple = (0.9, 0.95)
+    adam_wd: float = 0.1
+    grad_clip: float = 1.0
+    # bias anneal
+    bias_anneal_frac: float = 0.95     # disable aux-free bias updates after this fraction of steps
+    # eval / logging
+    val_every: int = 250
+    val_tokens: int = 10 * 1024 * 1024
+    log_every: int = 10
+    # run
+    seed: int = 1337
+    compile: bool = True
+    bf16: bool = True
+    out_dir: str = "runs"
+    run_name: str = "default"
+
+
+# ---- preset architectures (starting points; tune with count_params.py) ----
+PRESETS: dict[str, ModelConfig] = {
+    # dims tuned so TOTAL params hit targets (see count_params.py); G = dense_ffn/expert_ffn
+    "130M": ModelConfig(   # ~140M total / ~62M active, G=8; top_k bumped 4->8 (ablation: -0.025 val loss)
+        d_model=512, n_layers=12, n_dense_layers=1,
+        n_q_heads=8, n_kv_heads=2, head_dim=64,
+        dense_ffn=1536, expert_ffn=192, n_experts=32, top_k=8, n_shared=0,
+    ),
+    "500M": ModelConfig(   # ~500M total, G=7.2
+        d_model=768, n_layers=16, n_dense_layers=1,
+        n_q_heads=12, n_kv_heads=3, head_dim=128,
+        dense_ffn=2304, expert_ffn=320, n_experts=36, top_k=6, n_shared=1,
+    ),
+    "1B": ModelConfig(     # ~1.02B total, G=12.6
+        d_model=1024, n_layers=20, n_dense_layers=1,
+        n_q_heads=16, n_kv_heads=8, head_dim=128,
+        dense_ffn=2816, expert_ffn=224, n_experts=64, top_k=8, n_shared=1,
+    ),
+}
+
+
+def get_config(preset: str) -> ModelConfig:
+    if preset not in PRESETS:
+        raise KeyError(f"unknown preset {preset!r}; choose from {list(PRESETS)}")
+    # return a copy so callers can mutate for ablations
+    return ModelConfig(**PRESETS[preset].to_dict())

hobbylm/diffusion.py

@@ -0,0 +1,159 @@
+"""Pure masked-diffusion (LLaDA / MDLM) conversion utilities for the MoE LLM.
+
+The model itself only flips one thing for diffusion: attention becomes bidirectional
+(model.py, gated on cfg.diffusion). Everything else lives here:
+
+  forward_mask : the forward (noising) process used at train time.
+  generate     : iterative-denoising sampler with semi-autoregressive blocks.
+
+The TRAIN loss is model.diffusion_cross_entropy (fused/chunked for big batches). The
+unfused `diffusion_loss` below is for tests / sanity checks only.
+"""
+from __future__ import annotations
+
+import torch
+import torch.nn.functional as F
+from torch import Tensor
+
+
+def forward_mask(input_ids: Tensor, mask_id: int, eps: float = 1e-3,
+                 prompt_lens: Tensor | None = None, generator: torch.Generator | None = None):
+    """LLaDA forward process.
+
+    One mask ratio t ~ U(eps, 1) per sequence; mask each token iid with prob t.
+    Returns (noisy, labels, p_mask):
+      noisy  : input with masked positions replaced by mask_id
+      labels : original token at masked positions, -1 (ignore_index) elsewhere
+      p_mask : per-token mask probability t (broadcast), used for the 1/p reweighting
+    `prompt_lens` (B,) optionally protects a prompt prefix from being masked/scored (SFT).
+    """
+    b, l = input_ids.shape
+    dev = input_ids.device
+    t = torch.rand(b, device=dev, generator=generator) * (1 - eps) + eps   # (b,)
+    p_mask = t[:, None].expand(b, l).contiguous()                          # (b, l)
+    mask = torch.rand(b, l, device=dev, generator=generator) < p_mask
+    if prompt_lens is not None:
+        pos = torch.arange(l, device=dev)[None, :]
+        mask &= pos >= prompt_lens[:, None]
+    # guarantee >=1 masked token per sequence so no micro-batch contributes a zero loss
+    none_masked = ~mask.any(dim=1)
+    if none_masked.any():
+        rows = none_masked.nonzero(as_tuple=True)[0]
+        lo = 0 if prompt_lens is None else int(prompt_lens.min().item())
+        j = torch.randint(lo, l, (rows.numel(),), device=dev, generator=generator)
+        mask[rows, j] = True
+    noisy = torch.where(mask, torch.full_like(input_ids, mask_id), input_ids)
+    labels = torch.where(mask, input_ids, torch.full_like(input_ids, -1))
+    return noisy, labels, p_mask
+
+
+def diffusion_loss(logits: Tensor, labels: Tensor, p_mask: Tensor) -> Tensor:
+    """Unfused LLaDA loss (tests): sum_{masked} CE / p_mask, normalized by B*L."""
+    b, l, _ = logits.shape
+    m = labels != -1
+    if int(m.sum()) == 0:
+        return logits.sum() * 0.0
+    ce = F.cross_entropy(logits[m].float(), labels[m], reduction="none")
+    return (ce / p_mask[m]).sum() / (b * l)
+
+
+def get_num_transfer_tokens(n: int, steps: int) -> list[int]:
+    """Spread n unmask events as evenly as possible across `steps` (sums to n)."""
+    base = n // steps
+    out = [base] * steps
+    for i in range(n - base * steps):
+        out[i] += 1
+    return out
+
+
+def add_gumbel_noise(logits: Tensor, temperature: float,
+                     generator: torch.Generator | None = None) -> Tensor:
+    """Gumbel-max categorical sampling (LLaDA). argmax of this == a sample at `temperature`.
+    temperature<=0 -> identity (argmax == greedy)."""
+    if temperature <= 0:
+        return logits
+    logits = logits.to(torch.float64)
+    noise = torch.rand(logits.shape, dtype=torch.float64, device=logits.device, generator=generator)
+    gumbel = (-torch.log(noise + 1e-12)) ** temperature
+    return logits.exp() / gumbel
+
+
+def _rep_penalty(blk: Tensor, present_ids: Tensor, penalty: float) -> Tensor:
+    """CTRL-style penalty across the canvas: damp logits of tokens already present (prompt +
+    committed) so the denoiser stops filling many slots with the same token. In-place on blk."""
+    if penalty == 1.0 or present_ids.numel() == 0:
+        return blk
+    col = blk[:, present_ids]
+    blk[:, present_ids] = torch.where(col > 0, col / penalty, col * penalty)
+    return blk
+
+
+@torch.no_grad()
+def generate(model, prompt_ids: Tensor, gen_len: int = 256, block: int = 32, steps: int = 64,
+             mask_id: int = 50257, temperature: float = 0.0, rep_penalty: float = 1.0,
+             remask_steps: int = 0, remask_frac: float = 0.3, valid_vocab: int = 50257,
+             eos_id: int | None = None, generator: torch.Generator | None = None) -> Tensor:
+    """Semi-autoregressive iterative denoising. prompt_ids: (1, P). Returns generated ids (1, <=gen_len).
+
+    Each block of `block` masked slots is filled over ~`steps*block/gen_len` steps, committing the
+    highest-confidence still-masked positions each step (low-confidence-remasking selection). Then
+    `remask_steps` refinement passes re-mask the lowest-confidence committed tokens and re-predict
+    them with full bidirectional context — this is what lets the model fix repetition/mistakes.
+    Sentinels (>= valid_vocab, incl. mask_id) are banned from being emitted. Blocks are causal
+    w.r.t. each other (a block attends to the committed prefix + itself), bidirectional within.
+    """
+    was_training = model.training
+    model.eval()
+    dev = prompt_ids.device
+    x = torch.cat([prompt_ids, torch.full((1, gen_len), mask_id, device=dev, dtype=prompt_ids.dtype)], dim=1)
+    P = prompt_ids.shape[1]
+
+    def block_logits(b1: int, b0: int) -> Tensor:
+        """Forward prefix+block; return (blk_len, V) logits with sentinels banned + rep-penalty."""
+        logits, _ = model(x[:, :b1])
+        blk = logits[0, b0:b1].float()
+        blk[:, valid_vocab:] = -float("inf")                  # never emit mask/sentinel ids
+        present = torch.unique(x[0, :b1])
+        present = present[(present < valid_vocab) & (present != mask_id)]
+        return _rep_penalty(blk, present, rep_penalty)
+
+    def predict(blk: Tensor):
+        prob = blk.softmax(-1)
+        pred = add_gumbel_noise(blk, temperature, generator).argmax(-1) if temperature > 0 else blk.argmax(-1)
+        return pred, prob
+
+    for b0 in range(P, P + gen_len, block):
+        b1 = min(b0 + block, P + gen_len)
+        blk_len = b1 - b0
+        sb = max(1, round(steps * blk_len / gen_len))
+        sched = get_num_transfer_tokens(blk_len, sb)
+        # --- fill: commit the most-confident still-masked positions over sb steps ---
+        for s in range(sb):
+            pred, prob = predict(block_logits(b1, b0))
+            conf = prob.gather(-1, pred.unsqueeze(-1)).squeeze(-1)
+            still = x[0, b0:b1] == mask_id
+            conf = torch.where(still, conf, torch.full_like(conf, -1.0))
+            k = min(sched[s], int(still.sum()))
+            if k <= 0:
+                continue
+            idx = conf.topk(k).indices
+            x[0, b0 + idx] = pred[idx].to(x.dtype)
+        # --- refine: re-mask the least-confident committed tokens and re-predict them ---
+        for _ in range(remask_steps):
+            blk = block_logits(b1, b0)
+            prob = blk.softmax(-1)
+            cur = x[0, b0:b1]
+            cur_conf = prob.gather(-1, cur.unsqueeze(-1)).squeeze(-1)   # confidence in current tokens
+            r = max(1, int(blk_len * remask_frac))
+            x[0, b0 + cur_conf.topk(r, largest=False).indices] = mask_id
+            pred, _ = predict(block_logits(b1, b0))
+            still = (x[0, b0:b1] == mask_id).nonzero(as_tuple=True)[0]
+            x[0, b0 + still] = pred[still].to(x.dtype)
+        if eos_id is not None and bool((x[0, b0:b1] == eos_id).any()):
+            rel = int((x[0, b0:b1] == eos_id).nonzero(as_tuple=True)[0][0].item())
+            if was_training:
+                model.train()
+            return x[:, P:b0 + rel + 1]
+    if was_training:
+        model.train()
+    return x[:, P:]

hobbylm/generate.py

@@ -0,0 +1,132 @@
+"""Inference / text generation from a trained moe-lab checkpoint.
+
+Loads a saved checkpoint (model.pt / ckpt_*.pt), rebuilds the model from its embedded config,
+and autoregressively samples text with the GPT-2 (tiktoken) tokenizer.
+
+  python generate.py --ckpt runs/130M_10B/model.pt --prompt "The meaning of life is"
+"""
+from __future__ import annotations
+
+import argparse
+from contextlib import nullcontext
+
+import tiktoken
+import torch
+import torch.nn.functional as F
+
+from .config import ModelConfig
+from .model import MoETransformer
+
+GPT2_VALID = 50257          # real GPT-2 tokens (rest of vocab is padding)
+EOT = 50256                 # <|endoftext|>
+
+
+def load_model(ckpt_path: str, device):
+    ck = torch.load(ckpt_path, map_location=device, weights_only=False)
+    cfg_d = dict(ck["config"])
+    cfg_d.pop("preset", None)
+    cfg = ModelConfig(**cfg_d)
+    cfg.expert_backend = "grouped" if device.type == "cuda" else "bmm"
+    model = MoETransformer(cfg).to(device)
+    model.load_state_dict(ck["model"])
+    model.eval()
+    return model, cfg, ck.get("val_loss"), ck.get("step")
+
+
+def _banned_ngram_tokens(prev: list[int], n: int) -> list[int]:
+    """Tokens that would complete an already-seen n-gram (no-repeat-ngram blocking)."""
+    if n <= 0 or len(prev) < n:
+        return []
+    seen: dict[tuple, list[int]] = {}
+    for i in range(len(prev) - n + 1):
+        ng = tuple(prev[i:i + n])
+        seen.setdefault(ng[:-1], []).append(ng[-1])
+    return seen.get(tuple(prev[-(n - 1):]), [])
+
+
+@torch.no_grad()
+def generate(model, idx, max_new_tokens, temperature, top_k, device,
+             top_p=0.95, repetition_penalty=1.3, no_repeat_ngram_size=3, ctx_len=1024):
+    amp = torch.autocast("cuda", dtype=torch.bfloat16) if device.type == "cuda" else nullcontext()
+    for _ in range(max_new_tokens):
+        idx_cond = idx[:, -ctx_len:]
+        with amp:
+            logits, _ = model(idx_cond)
+        logits = logits[:, -1, :].float()
+        logits[:, GPT2_VALID:] = -float("inf")          # never emit padding tokens
+
+        seq = idx[0].tolist()
+        # repetition penalty (CTRL-style): damp logits of already-generated tokens
+        if repetition_penalty and repetition_penalty != 1.0:
+            uniq = torch.tensor(sorted(set(seq)), device=logits.device)
+            lg = logits[0, uniq]
+            logits[0, uniq] = torch.where(lg > 0, lg / repetition_penalty, lg * repetition_penalty)
+        # no-repeat n-gram blocking
+        for t in _banned_ngram_tokens(seq, no_repeat_ngram_size):
+            logits[0, t] = -float("inf")
+
+        if temperature > 0:
+            logits = logits / temperature
+            if top_k:
+                v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
+                logits[logits < v[:, [-1]]] = -float("inf")
+            if top_p and top_p < 1.0:                    # nucleus filtering
+                s_logits, s_idx = torch.sort(logits, descending=True)
+                cum = torch.cumsum(F.softmax(s_logits, dim=-1), dim=-1)
+                rm = cum > top_p
+                rm[..., 1:] = rm[..., :-1].clone()
+                rm[..., 0] = False
+                logits[0, s_idx[0, rm[0]]] = -float("inf")
+            nxt = torch.multinomial(F.softmax(logits, dim=-1), 1)
+        else:
+            nxt = logits.argmax(-1, keepdim=True)
+        idx = torch.cat([idx, nxt], dim=1)
+        if nxt.item() == EOT:
+            break
+    return idx
+
+
+def run(ckpt_path, prompts, max_new_tokens=120, temperature=0.9, top_k=0, device=None,
+        top_p=0.95, repetition_penalty=1.3, no_repeat_ngram_size=3):
+    device = device or torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    model, cfg, val_loss, step = load_model(ckpt_path, device)
+    enc = tiktoken.get_encoding("gpt2")
+    print(f"loaded {ckpt_path} | step={step} val_loss={val_loss} | "
+          f"d_model={cfg.d_model} layers={cfg.n_layers} experts={cfg.n_experts}/top{cfg.top_k}", flush=True)
+    print(f"sampling: temp={temperature} top_k={top_k} top_p={top_p} "
+          f"rep_penalty={repetition_penalty} no_repeat_ngram={no_repeat_ngram_size}\n", flush=True)
+    for p in prompts:
+        ids = torch.tensor([enc.encode_ordinary(p)], dtype=torch.long, device=device)
+        out = generate(model, ids, max_new_tokens, temperature, top_k, device,
+                       top_p=top_p, repetition_penalty=repetition_penalty,
+                       no_repeat_ngram_size=no_repeat_ngram_size)
+        text = enc.decode(out[0].tolist())
+        print("=" * 70)
+        print(f"PROMPT: {p!r}")
+        print(text)
+        print(flush=True)
+
+
+if __name__ == "__main__":
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--ckpt", default="runs/130M_10B/model.pt")
+    ap.add_argument("--prompt", default=None)
+    ap.add_argument("--max_new_tokens", type=int, default=120)
+    ap.add_argument("--temperature", type=float, default=0.9)
+    ap.add_argument("--top_k", type=int, default=0)
+    ap.add_argument("--top_p", type=float, default=0.95)
+    ap.add_argument("--repetition_penalty", type=float, default=1.3)
+    ap.add_argument("--no_repeat_ngram_size", type=int, default=3)
+    args = ap.parse_args()
+    default_prompts = [
+        "The meaning of life is",
+        "Once upon a time, there was a",
+        "The capital of France is",
+        "In 2023, scientists discovered that",
+        "To make a good cup of coffee, you",
+        "The most important thing about climate change is",
+    ]
+    prompts = [args.prompt] if args.prompt else default_prompts
+    run(args.ckpt, prompts, args.max_new_tokens, args.temperature, args.top_k,
+        top_p=args.top_p, repetition_penalty=args.repetition_penalty,
+        no_repeat_ngram_size=args.no_repeat_ngram_size)

hobbylm/model.py

@@ -0,0 +1,410 @@
+"""The MoE transformer: GQA + QK-norm + RoPE attention, dense/MoE SwiGLU FFNs, tied head.
+
+Pre-norm RMSNorm blocks. First cfg.n_dense_layers use a dense SwiGLU FFN; the rest use MoE.
+Loss = cross-entropy + final-logit z-loss + sum of per-layer MoE aux/z losses.
+"""
+from __future__ import annotations
+
+import math
+import torch
+import torch.distributed as dist
+import torch.nn as nn
+import torch.nn.functional as F
+from torch import Tensor
+
+from .config import ModelConfig
+from .moe import MoE, SwiGLUWeights
+
+
+# -----------------------------------------------------------------------------
+# FP8 matmul for the lm_head (the single largest GEMM). Adapted from modded-nanogpt
+# (@YouJiacheng): weight stored transposed (in, out) so the gradient w.r.t. the weight
+# lands in the natural layout. Forward in e4m3, backward grad in e5m2. bf16 outputs.
+# Wrapped as a custom op with an explicit autograd rule (torch._scaled_mm is not
+# differentiable on its own).
+
+@torch.library.custom_op("moelab::mm_t", mutates_args=())
+def _mm_t(x: Tensor, w: Tensor, x_s: float, w_s: float, grad_s: float) -> tuple[Tensor, Tensor, Tensor]:
+    """y = x @ w with x:(M,in), w:(in,out). Returns (y_bf16, x_f8, w_f8) for backward reuse."""
+    @torch.compile
+    def impl(x: Tensor, w: Tensor):
+        x_f8 = x.div(x_s).to(torch.float8_e4m3fn)
+        w_f8 = w.div(w_s).to(torch.float8_e4m3fn)
+        w_col = w_f8.T.contiguous().T  # _scaled_mm needs column-major B
+        out = torch._scaled_mm(x_f8, w_col, out_dtype=torch.bfloat16,
+                               scale_a=x.new_tensor(x_s, dtype=torch.float32),
+                               scale_b=x.new_tensor(w_s, dtype=torch.float32),
+                               use_fast_accum=True)
+        return out, x_f8, w_f8
+    return impl(x, w)
+
+
+@_mm_t.register_fake
+def _(x: Tensor, w: Tensor, *_):
+    return x @ w, x.to(torch.float8_e4m3fn), w.to(torch.float8_e4m3fn)
+
+
+@torch.library.custom_op("moelab::mm_t_backward", mutates_args=())
+def _mm_t_backward(g: Tensor, x_f8: Tensor, w_f8: Tensor, x_s: float, w_s: float, grad_s: float) -> tuple[Tensor, Tensor]:
+    @torch.compile
+    def impl(g: Tensor, x_f8: Tensor, w_f8: Tensor):
+        x_scale = g.new_tensor(x_s, dtype=torch.float32)
+        w_scale = g.new_tensor(w_s, dtype=torch.float32)
+        g_scale = g.new_tensor(grad_s, dtype=torch.float32)
+        g_f8 = g.div(grad_s).to(torch.float8_e5m2)
+        grad_x = torch._scaled_mm(g_f8, w_f8.T, out_dtype=torch.bfloat16,
+                                  scale_a=g_scale, scale_b=w_scale, use_fast_accum=False)
+        grad_w = torch._scaled_mm(x_f8.T.contiguous(), g_f8.T.contiguous().T, out_dtype=torch.float32,
+                                  scale_a=x_scale, scale_b=g_scale, use_fast_accum=False)
+        return grad_x, grad_w
+    return impl(g, x_f8, w_f8)
+
+
+@_mm_t_backward.register_fake
+def _(g: Tensor, x_f8: Tensor, w_f8: Tensor, *_):
+    return x_f8.to(torch.bfloat16), w_f8.to(torch.float32)
+
+
+def _mm_t_setup(ctx, inputs, output):
+    *_, x_s, w_s, grad_s = inputs
+    _, x_f8, w_f8 = output
+    ctx.save_for_backward(x_f8, w_f8)
+    ctx.scales = (x_s, w_s, grad_s)
+    ctx.set_materialize_grads(False)
+
+
+def _mm_t_bwd(ctx, grad_out: Tensor, *_):
+    x_f8, w_f8 = ctx.saved_tensors
+    x_s, w_s, grad_s = ctx.scales
+    gx, gw = torch.ops.moelab.mm_t_backward(grad_out, x_f8, w_f8, x_s, w_s, grad_s)
+    return gx, gw, None, None, None
+
+
+_mm_t.register_autograd(_mm_t_bwd, setup_context=_mm_t_setup)
+
+
+class FP8Linear(nn.Module):
+    """Bias-free linear with FP8 matmul in training (CUDA) and a bf16 fallback elsewhere.
+    Weight is stored transposed (in_features, out_features)."""
+    def __init__(self, in_features: int, out_features: int, x_s=1.0, w_s=1.0, grad_s=1.0):
+        super().__init__()
+        self.in_features, self.out_features = in_features, out_features
+        self.x_s, self.w_s, self.grad_s = x_s, w_s, grad_s
+        self.weight = nn.Parameter(torch.empty(in_features, out_features))
+
+    def forward(self, x: Tensor) -> Tensor:
+        if self.training and x.is_cuda:
+            flat = x.flatten(0, -2).bfloat16().contiguous()
+            out = torch.ops.moelab.mm_t(flat, self.weight.bfloat16().contiguous(),
+                                        self.x_s, self.w_s, self.grad_s)[0]
+            return out.reshape(*x.shape[:-1], self.out_features)
+        return x @ self.weight.type_as(x)
+
+
+# -----------------------------------------------------------------------------
+# Fused cross-entropy: process the vocab projection in row-chunks under activation
+# checkpointing so the full (T, vocab) fp32 logit tensor is never materialized or
+# saved for backward. Numerically identical to a plain CE + final z-loss.
+
+def _ce_chunk(x_c: Tensor, weight: Tensor, tgt_c: Tensor, softcap: float, tied: bool):
+    logits = (x_c @ weight.T) if tied else (x_c @ weight)   # tied: weight (V,d); fp8 head: weight (d,V)
+    if softcap > 0:
+        logits = softcap * torch.tanh(logits / softcap)
+    logits = logits.float()
+    lse = torch.logsumexp(logits, dim=-1)                   # (c,)
+    z_sum = (lse * lse).sum()
+    valid = (tgt_c != -1)
+    tgt_logit = logits.gather(-1, tgt_c.clamp_min(0).unsqueeze(-1)).squeeze(-1)
+    ce_sum = ((lse - tgt_logit) * valid).sum()
+    return ce_sum, z_sum
+
+
+def fused_cross_entropy(x: Tensor, weight: Tensor, targets: Tensor, *,
+                        z_coef: float, softcap: float, chunk: int, tied: bool):
+    """Returns (loss = mean_ce + z_coef * mean(lse^2), z_mean.detach()). Memory-light."""
+    from torch.utils.checkpoint import checkpoint
+    T = x.shape[0]
+    n_valid = (targets != -1).sum().clamp_min(1)
+    ce_sum = x.new_zeros((), dtype=torch.float32)
+    z_sum = x.new_zeros((), dtype=torch.float32)
+    for i in range(0, T, chunk):
+        c_ce, c_z = checkpoint(_ce_chunk, x[i:i + chunk], weight, targets[i:i + chunk],
+                               softcap, tied, use_reentrant=False)
+        ce_sum = ce_sum + c_ce
+        z_sum = z_sum + c_z
+    ce = ce_sum / n_valid
+    z_mean = z_sum / T
+    return ce + z_coef * z_mean, z_mean.detach()
+
+
+# Masked-diffusion (LLaDA/MDLM) loss: cross-entropy on the masked positions only, reweighted
+# by 1/p_mask, summed and normalized by the total token count (B*L). Chunked + activation-
+# checkpointed like fused_cross_entropy so the (T, vocab) logits never fully materialize.
+
+def _dce_chunk(x_c: Tensor, weight: Tensor, tgt_c: Tensor, pm_c: Tensor, softcap: float, tied: bool):
+    logits = (x_c @ weight.T) if tied else (x_c @ weight)
+    if softcap > 0:
+        logits = softcap * torch.tanh(logits / softcap)
+    logits = logits.float()
+    lse = torch.logsumexp(logits, dim=-1)
+    valid = (tgt_c != -1)                                    # unmasked positions carry target -1
+    tgt_logit = logits.gather(-1, tgt_c.clamp_min(0).unsqueeze(-1)).squeeze(-1)
+    ce = (lse - tgt_logit) * valid / pm_c                    # 1/p_mask reweight; unmasked -> 0
+    return ce.sum()
+
+
+def diffusion_cross_entropy(x: Tensor, weight: Tensor, targets: Tensor, p_mask: Tensor, *,
+                            softcap: float, chunk: int, tied: bool):
+    """LLaDA loss = sum_{masked} CE / p_mask, normalized by B*L. Memory-light."""
+    from torch.utils.checkpoint import checkpoint
+    T = x.shape[0]
+    ce_sum = x.new_zeros((), dtype=torch.float32)
+    for i in range(0, T, chunk):
+        ce_sum = ce_sum + checkpoint(_dce_chunk, x[i:i + chunk], weight, targets[i:i + chunk],
+                                     p_mask[i:i + chunk], softcap, tied, use_reentrant=False)
+    return ce_sum / T
+
+
+def rms_norm(x: Tensor, weight: Tensor | None = None, eps: float = 1e-6) -> Tensor:
+    out = F.rms_norm(x, (x.size(-1),), eps=eps)
+    return out * weight if weight is not None else out
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, dim: int, eps: float = 1e-6):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(dim))
+        self.eps = eps
+
+    def forward(self, x: Tensor) -> Tensor:
+        return rms_norm(x, self.weight, self.eps)
+
+
+def precompute_rope(head_dim: int, max_seq: int, theta: float, device) -> tuple[Tensor, Tensor]:
+    inv_freq = 1.0 / (theta ** (torch.arange(0, head_dim, 2, device=device).float() / head_dim))
+    t = torch.arange(max_seq, device=device).float()
+    freqs = torch.outer(t, inv_freq)            # (S, head_dim/2)
+    return freqs.cos(), freqs.sin()             # each (S, head_dim/2)
+
+
+def apply_rope(x: Tensor, cos: Tensor, sin: Tensor) -> Tensor:
+    # x: (B, H, S, D). Rotate-half formulation.
+    S, D = x.shape[-2], x.shape[-1]
+    cos = cos[:S].view(1, 1, S, D // 2)
+    sin = sin[:S].view(1, 1, S, D // 2)
+    x1, x2 = x[..., : D // 2], x[..., D // 2:]
+    return torch.cat([x1 * cos - x2 * sin, x2 * cos + x1 * sin], dim=-1)
+
+
+class Attention(nn.Module):
+    def __init__(self, cfg: ModelConfig):
+        super().__init__()
+        self.cfg = cfg
+        self.nq, self.nkv, self.hd = cfg.n_q_heads, cfg.n_kv_heads, cfg.head_dim
+        self.rep = self.nq // self.nkv
+        qkv_out = (self.nq + 2 * self.nkv) * self.hd
+        self.qkv = nn.Linear(cfg.d_model, qkv_out, bias=False)
+        self.proj = nn.Linear(self.nq * self.hd, cfg.d_model, bias=False)
+        if cfg.qk_norm:
+            self.q_norm = RMSNorm(self.hd)
+            self.k_norm = RMSNorm(self.hd)
+
+    def forward(self, x: Tensor, cos: Tensor, sin: Tensor) -> Tensor:
+        B, S, _ = x.shape
+        qkv = self.qkv(x)
+        q, k, v = qkv.split([self.nq * self.hd, self.nkv * self.hd, self.nkv * self.hd], dim=-1)
+        q = q.view(B, S, self.nq, self.hd).transpose(1, 2)    # (B, nq, S, hd)
+        k = k.view(B, S, self.nkv, self.hd).transpose(1, 2)
+        v = v.view(B, S, self.nkv, self.hd).transpose(1, 2)
+        if self.cfg.qk_norm:
+            q, k = self.q_norm(q), self.k_norm(k)             # per-head RMSNorm before RoPE
+        q, k = apply_rope(q, cos, sin), apply_rope(k, cos, sin)
+        # GQA: expand kv heads to match q heads
+        k = k.repeat_interleave(self.rep, dim=1)
+        v = v.repeat_interleave(self.rep, dim=1)
+        # diffusion (LLaDA) models see the whole noised canvas -> bidirectional; AR stays causal.
+        o = F.scaled_dot_product_attention(q, k, v, is_causal=not self.cfg.diffusion)
+        o = o.transpose(1, 2).reshape(B, S, self.nq * self.hd)
+        return self.proj(o)
+
+
+class DenseFFN(nn.Module):
+    def __init__(self, cfg: ModelConfig):
+        super().__init__()
+        self.w13 = nn.Linear(cfg.d_model, 2 * cfg.dense_ffn, bias=False)
+        self.w2 = nn.Linear(cfg.dense_ffn, cfg.d_model, bias=False)
+
+    def forward(self, x: Tensor) -> Tensor:
+        gate, up = self.w13(x).chunk(2, dim=-1)
+        return self.w2(F.silu(gate) * up)
+
+
+class Block(nn.Module):
+    def __init__(self, cfg: ModelConfig, layer_idx: int):
+        super().__init__()
+        self.attn_norm = RMSNorm(cfg.d_model)
+        self.attn = Attention(cfg)
+        self.ffn_norm = RMSNorm(cfg.d_model)
+        self.is_moe = layer_idx >= cfg.n_dense_layers
+        self.ffn = MoE(cfg) if self.is_moe else DenseFFN(cfg)
+
+    def forward(self, x: Tensor, cos: Tensor, sin: Tensor):
+        x = x + self.attn(self.attn_norm(x), cos, sin)
+        if self.is_moe:
+            out, aux = self.ffn(self.ffn_norm(x))
+            return x + out, aux
+        return x + self.ffn(self.ffn_norm(x)), x.new_zeros(())
+
+
+class MoETransformer(nn.Module):
+    def __init__(self, cfg: ModelConfig):
+        super().__init__()
+        self.cfg = cfg
+        self.embed = nn.Embedding(cfg.vocab_size, cfg.d_model)
+        self.blocks = nn.ModuleList([Block(cfg, i) for i in range(cfg.n_layers)])
+        self.final_norm = RMSNorm(cfg.d_model)
+        # FP8 head must be untied (it stores the weight transposed, (d, vocab), for fp8 grad layout).
+        self.fp8_head = cfg.fp8_head
+        if cfg.fp8_head:
+            import warnings
+            warnings.warn("fp8_head is EXPERIMENTAL/BROKEN: zero gradient flows through the FP8 "
+                          "backward (model won't train) and it gave no speedup. Use fused_ce instead.")
+            self.lm_head = FP8Linear(cfg.d_model, cfg.vocab_size,
+                                     x_s=cfg.fp8_x_scale, w_s=cfg.fp8_w_scale, grad_s=cfg.fp8_grad_scale)
+        else:
+            self.lm_head = nn.Linear(cfg.d_model, cfg.vocab_size, bias=False)
+            if cfg.tie_embeddings:
+                self.lm_head.weight = self.embed.weight
+        self._rope_cache: dict = {}
+        self.apply(self._init)
+        if cfg.fp8_head:
+            # start the untied head from the (tied) embedding weights so the ablation isolates fp8,
+            # not a different head initialization.
+            with torch.no_grad():
+                self.lm_head.weight.copy_(self.embed.weight.t())
+        self._scale_residual_init()
+
+    # ---- init ----
+    def _init(self, m: nn.Module):
+        std = self.cfg.init_std
+        if isinstance(m, nn.Linear):
+            nn.init.normal_(m.weight, mean=0.0, std=std)
+            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=std)
+        elif isinstance(m, SwiGLUWeights):
+            nn.init.normal_(m.w13, mean=0.0, std=std)
+            nn.init.normal_(m.w2, mean=0.0, std=std)
+
+    def _scale_residual_init(self):
+        # scale residual-projection weights by 1/sqrt(2*n_layers) (GPT-2/Megatron), critical deep-thin
+        scale = (2 * self.cfg.n_layers) ** -0.5
+        for blk in self.blocks:
+            with torch.no_grad():
+                blk.attn.proj.weight.mul_(scale)
+                if isinstance(blk.ffn, DenseFFN):
+                    blk.ffn.w2.weight.mul_(scale)
+                else:
+                    blk.ffn.experts.w2.mul_(scale)
+                    if self.cfg.n_shared > 0:
+                        blk.ffn.shared.w2.mul_(scale)
+                    # small router init (~0.1x)
+                    blk.ffn.gate.weight.mul_(0.1)
+
+    def rope(self, S: int, device, dtype):
+        key = (S, device, dtype)
+        if key not in self._rope_cache:
+            cos, sin = precompute_rope(self.cfg.head_dim, S, self.cfg.rope_theta, device)
+            self._rope_cache[key] = (cos.to(dtype), sin.to(dtype))
+        return self._rope_cache[key]
+
+    def forward(self, idx: Tensor | None = None, targets: Tensor | None = None,
+                inputs_embeds: Tensor | None = None, p_mask: Tensor | None = None):
+        # accept either token ids OR precomputed embeddings (inputs_embeds), for multimodal splicing.
+        if inputs_embeds is None:
+            x = self.embed(idx)
+            if self.cfg.scale_embeddings:
+                x = x * (self.cfg.d_model ** 0.5)
+            device = idx.device
+        else:
+            x = inputs_embeds
+            device = inputs_embeds.device
+        B, S = x.shape[0], x.shape[1]
+        cos, sin = self.rope(S, device, x.dtype)
+        aux_sum = x.new_zeros(())
+        for blk in self.blocks:
+            x, aux = blk(x, cos, sin)
+            aux_sum = aux_sum + aux
+        x = self.final_norm(x)
+        cfg = self.cfg
+        sc = cfg.logit_softcap
+
+        # ---- inference: return logits ----
+        if targets is None:
+            logits = self.lm_head(x)
+            if sc > 0:
+                logits = sc * torch.tanh(logits / sc)
+            return logits, aux_sum
+
+        # ---- training: compute loss ----
+        if cfg.diffusion:
+            # LLaDA masked-diffusion loss on the noised positions (input already masked upstream;
+            # targets hold the original token at masked positions, -1 elsewhere; p_mask = per-token t).
+            assert p_mask is not None, "diffusion forward needs p_mask (use diffusion.forward_mask)"
+            loss_ce = diffusion_cross_entropy(
+                x.reshape(-1, x.size(-1)), self.lm_head.weight, targets.reshape(-1),
+                p_mask.reshape(-1), softcap=sc, chunk=cfg.ce_chunk, tied=not cfg.fp8_head)
+            loss = loss_ce + aux_sum
+            return loss, {"ce": loss_ce.detach(), "aux": aux_sum.detach(), "z": x.new_zeros(())}
+
+        if cfg.fused_ce and not cfg.fp8_head:
+            # chunked CE on the tied weight (V, d); never materializes the full fp32 logits.
+            loss_cez, z = fused_cross_entropy(
+                x.reshape(-1, x.size(-1)), self.lm_head.weight, targets.reshape(-1),
+                z_coef=cfg.final_z_loss_coef, softcap=sc, chunk=cfg.ce_chunk, tied=True)
+            loss = loss_cez + aux_sum
+            ce = (loss_cez - cfg.final_z_loss_coef * z).detach()
+            return loss, {"ce": ce, "aux": aux_sum.detach(), "z": z}
+
+        logits = self.lm_head(x)                       # fp8 head -> bf16, else nn.Linear
+        if sc > 0:
+            logits = sc * torch.tanh(logits / sc)
+        logits = logits.float()
+        ce = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1), ignore_index=-1)
+        z = (torch.logsumexp(logits, dim=-1) ** 2).mean()
+        loss = ce + cfg.final_z_loss_coef * z + aux_sum
+        return loss, {"ce": ce.detach(), "aux": aux_sum.detach(), "z": z.detach()}
+
+    @torch.no_grad()
+    def set_bias_update_rate(self, rate: float):
+        for blk in self.blocks:
+            if isinstance(blk.ffn, MoE):
+                blk.ffn.bias_update_rate = rate
+
+    @torch.no_grad()
+    def sync_expert_bias(self):
+        """Average aux-free bias buffers across DDP ranks so they stay identical
+        (each rank updates from local token counts; DDP doesn't sync buffers)."""
+        if not (dist.is_available() and dist.is_initialized()):
+            return
+        world = dist.get_world_size()
+        for blk in self.blocks:
+            if isinstance(blk.ffn, MoE):
+                dist.all_reduce(blk.ffn.expert_bias, op=dist.ReduceOp.SUM)
+                blk.ffn.expert_bias.div_(world)
+
+
+def count_params(model: MoETransformer) -> dict:
+    cfg = model.cfg
+    total = sum(p.numel() for p in model.parameters())
+    # subtract tied head double-count is already avoided (shared weight counted once).
+    # fp8_head forces an untied head, so it costs a second vocab x d_model matrix.
+    tied = cfg.tie_embeddings and not cfg.fp8_head
+    embed = cfg.vocab_size * cfg.d_model * (1 if tied else 2)
+    # active = total - inactive routed experts. Per MoE layer, only top_k of n_experts run.
+    per_expert = cfg.d_model * 2 * cfg.expert_ffn + cfg.expert_ffn * cfg.d_model
+    inactive_per_moe = (cfg.n_experts - cfg.top_k) * per_expert
+    active = total - cfg.n_moe_layers * inactive_per_moe
+    return {"total": total, "active": active, "embed": embed,
+            "active_pct": 100 * active / total, "sparsity": total / active}

hobbylm/moe.py

@@ -0,0 +1,154 @@
+"""Mixture-of-Experts layer: fp32 router, dropless expert compute, load balancing.
+
+Backends for the expert compute (selected by cfg.expert_backend):
+  - "grouped": sort tokens by expert -> torch grouped_mm -> scatter. Fast on H100/A100, bf16.
+  - "bmm":     loop-free reference using per-expert masked matmuls. CPU-testable.
+  - "loop":    explicit python loop over experts. Slowest, clearest reference.
+
+Balancing:
+  - "aux_free":  DeepSeek-V3 gradient-free per-expert bias added to the SELECTION scores only
+                 (not the gate weights), updated by sign(load error). Plus a tiny aux loss safety net.
+  - "aux_loss":  classic Switch/OLMoE load-balance loss (coef ~1e-2) only.
+  Both add a router z-loss.
+"""
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch import Tensor
+
+from .config import ModelConfig
+
+
+def _grouped_mm(a: Tensor, b: Tensor, offs: Tensor) -> Tensor:
+    """a: (T, d_in) expert-sorted; b: (E, d_in, d_out); offs: (E,) int32 cumulative row counts.
+    grouped_mm requires bf16 inputs; cast here (autocast does not cover this custom op)."""
+    fn = getattr(F, "grouped_mm", None) or getattr(torch, "_grouped_mm", None)
+    if fn is None:
+        raise RuntimeError("grouped_mm not available in this torch build; use expert_backend='bmm'")
+    return fn(a.bfloat16(), b.bfloat16(), offs=offs)
+
+
+class SwiGLUWeights(nn.Module):
+    """A stack of E SwiGLU experts. w13: (E, d, 2f) gate+up fused; w2: (E, f, d) down."""
+    def __init__(self, n_experts: int, d_model: int, ffn: int):
+        super().__init__()
+        self.n_experts = n_experts
+        self.d_model = d_model
+        self.ffn = ffn
+        self.w13 = nn.Parameter(torch.empty(n_experts, d_model, 2 * ffn))
+        self.w2 = nn.Parameter(torch.empty(n_experts, ffn, d_model))
+
+    def expert_glu(self, x: Tensor, e: int) -> Tensor:
+        h = x @ self.w13[e]
+        gate, up = h.chunk(2, dim=-1)
+        return (F.silu(gate) * up) @ self.w2[e]
+
+
+class MoE(nn.Module):
+    def __init__(self, cfg: ModelConfig):
+        super().__init__()
+        self.cfg = cfg
+        self.n_experts = cfg.n_experts
+        self.top_k = cfg.top_k
+        self.n_shared = cfg.n_shared
+        self.backend = cfg.expert_backend
+
+        # fp32 router (gate). Initialized small in model init.
+        self.gate = nn.Linear(cfg.d_model, cfg.n_experts, bias=False)
+        self.experts = SwiGLUWeights(cfg.n_experts, cfg.d_model, cfg.expert_ffn)
+        if cfg.n_shared > 0:
+            self.shared = SwiGLUWeights(cfg.n_shared, cfg.d_model, cfg.expert_ffn)
+
+        # aux-loss-free per-expert routing bias (no grad), updated by sign(load error)
+        self.register_buffer("expert_bias", torch.zeros(cfg.n_experts))
+        self.bias_update_rate = cfg.bias_update_rate  # set to 0 to freeze (anneal at end of training)
+
+    # ---- routing (always fp32, even under bf16 autocast) ----
+    def _route(self, x: Tensor):
+        """x: (T, d). Returns topi (T,k) long, topv (T,k) fp32 gate weights, aux_loss scalar."""
+        cfg = self.cfg
+        with torch.autocast(device_type=x.device.type, enabled=False):
+            logits = F.linear(x.float(), self.gate.weight.float())   # fp32 router
+            scores = torch.sigmoid(logits) if cfg.gating == "sigmoid" else torch.softmax(logits, dim=-1)
+
+            # selection scores: add gradient-free bias (aux_free) for balanced routing
+            sel = scores + self.expert_bias.float() if cfg.balancing == "aux_free" else scores
+            topi = torch.topk(sel, self.top_k, dim=-1).indices       # (T, k)
+            topv = torch.gather(scores, -1, topi)                    # gate weights from ORIGINAL scores
+            if cfg.norm_topk_prob:
+                topv = topv / (topv.sum(-1, keepdim=True) + 1e-9)
+
+            # ---- balancing losses ----
+            T = x.shape[0]
+            counts = torch.bincount(topi.reshape(-1), minlength=self.n_experts).float()
+            f_i = counts / (T * self.top_k)
+            P_i = scores.mean(dim=0)
+            aux = self.n_experts * (f_i.detach() * P_i).sum()
+            z_loss = (torch.logsumexp(logits, dim=-1) ** 2).mean()
+            aux_total = cfg.aux_loss_coef * aux + cfg.z_loss_coef * z_loss
+
+            # update aux-free bias (no grad): under-loaded experts up, over-loaded down
+            if cfg.balancing == "aux_free" and self.training and self.bias_update_rate > 0:
+                with torch.no_grad():
+                    ideal = T * self.top_k / self.n_experts
+                    self.expert_bias.add_(self.bias_update_rate * torch.sign(ideal - counts))
+
+        return topi, topv, aux_total
+
+    # ---- expert compute backends ----
+    def _experts_grouped(self, x: Tensor, topi: Tensor, topv: Tensor) -> Tensor:
+        T, d = x.shape
+        k = self.top_k
+        flat_e = topi.reshape(-1)
+        flat_tok = torch.arange(T, device=x.device).repeat_interleave(k)
+        order = torch.argsort(flat_e)
+        sort_e = flat_e[order]
+        sort_tok = flat_tok[order]
+        xs = x[sort_tok]                                            # (T*k, d)
+        counts = torch.bincount(sort_e, minlength=self.n_experts)
+        offs = torch.cumsum(counts, 0).to(torch.int32)
+        h = _grouped_mm(xs, self.experts.w13, offs)                # (T*k, 2f) bf16
+        gate, up = h.chunk(2, dim=-1)
+        h = F.silu(gate) * up
+        y = _grouped_mm(h, self.experts.w2, offs).to(x.dtype)      # (T*k, d) -> residual dtype
+        y = y * topv.reshape(-1)[order].unsqueeze(-1).to(x.dtype)
+        out = torch.zeros_like(x)
+        out.index_add_(0, sort_tok, y)
+        return out
+
+    def _experts_bmm(self, x: Tensor, topi: Tensor, topv: Tensor) -> Tensor:
+        # reference path: per-expert masked matmul (works on CPU)
+        T, d = x.shape
+        flat_e = topi.reshape(-1)
+        flat_tok = torch.arange(T, device=x.device).repeat_interleave(self.top_k)
+        flat_v = topv.reshape(-1)
+        out = torch.zeros_like(x)
+        for e in range(self.n_experts):
+            sel = (flat_e == e).nonzero(as_tuple=True)[0]
+            if sel.numel() == 0:
+                continue
+            toks = flat_tok[sel]
+            ye = self.experts.expert_glu(x[toks], e).to(x.dtype) * flat_v[sel].unsqueeze(-1).to(x.dtype)
+            out.index_add_(0, toks, ye)
+        return out
+
+    def _shared(self, x: Tensor) -> Tensor:
+        out = x.new_zeros(x.shape)
+        for e in range(self.n_shared):
+            out = out + self.shared.expert_glu(x, e).to(x.dtype)
+        return out
+
+    def forward(self, x: Tensor):
+        """x: (B, S, d) -> (out (B,S,d), aux_loss scalar)."""
+        B, S, d = x.shape
+        xf = x.reshape(-1, d)
+        topi, topv, aux = self._route(xf)
+        if self.backend == "grouped":
+            out = self._experts_grouped(xf, topi, topv)
+        else:
+            out = self._experts_bmm(xf, topi, topv)
+        if self.n_shared > 0:
+            out = out + self._shared(xf)
+        return out.reshape(B, S, d), aux

hobbylm/multimodal.py

@@ -0,0 +1,151 @@
+"""Multimodal (image + audio) wrapper around the MoE LLM — TinyLLaVA-style.
+
+A frozen vision/audio encoder produces patch features; a small MLP projector maps them into the
+LLM's embedding space; the LLM consumes them as ordinary token embeddings spliced at `<image>` /
+`<audio>` sentinel positions. The LLM is unchanged except `forward(inputs_embeds=...)`.
+
+Sentinels live in the padded-but-unused GPT-2 vocab slots (50257-50303), so they never collide with
+real tokens and the lm_head already masks them at decode (see generate.py).
+
+v1 scope: encoders frozen (features precomputed/cached); projector(s) + (optionally) LLM are trained.
+Batching assumes uniform feature count per modality (fixed-resolution encoders); samples may carry an
+image, audio, both, or neither (text-only) — the merge right-pads, which is safe under causal attention.
+"""
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+from torch import Tensor
+
+from .config import ModelConfig
+from .model import MoETransformer
+
+# ---- special tokens in the free vocab slots (vocab 50257 real -> padded 50304) ----
+IMAGE_TOKEN = 50257   # one sentinel per image; replaced by N projected patch features
+AUDIO_TOKEN = 50258   # one sentinel per audio clip
+IM_START = 50259
+IM_END = 50260
+VIDEO_TOKEN = 50261   # video = sampled frames through the SAME vision encoder + mm_projector (no new encoder)
+SPEECH_TOKEN = 50262  # spoken language via a Whisper encoder + speech_projector (distinct from CLAP <audio>)
+IGNORE_INDEX = -1     # target value for non-text positions; matches model CE ignore_index
+
+
+class Projector(nn.Module):
+    """TinyLLaVA `mlp2x_gelu` connector: Linear -> GELU -> Linear, encoder_dim -> d_model."""
+    def __init__(self, in_dim: int, out_dim: int):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.Linear(in_dim, out_dim),
+            nn.GELU(),
+            nn.Linear(out_dim, out_dim),
+        )
+
+    def forward(self, x: Tensor) -> Tensor:
+        return self.net(x)
+
+
+class MoEVLM(nn.Module):
+    """Wraps a MoETransformer with image (and optional audio) projectors.
+
+    forward inputs:
+      input_ids:       (B, L) token ids containing IMAGE_TOKEN / AUDIO_TOKEN sentinels.
+      image_features:  (B, Ni, vision_dim) RAW frozen-encoder features, or None.
+      audio_features:  (B, Na, audio_dim) RAW frozen-encoder features, or None.
+      targets:         (B, L) next-token targets aligned to input_ids (sentinels get IGNORE_INDEX), or None.
+    """
+    def __init__(self, llm: MoETransformer, vision_dim: int = 1152, audio_dim: int | None = None,
+                 speech_dim: int | None = None):
+        super().__init__()
+        self.llm = llm
+        self.d_model = llm.cfg.d_model
+        self.mm_projector = Projector(vision_dim, self.d_model)
+        self.audio_projector = Projector(audio_dim, self.d_model) if audio_dim else None
+        self.speech_projector = Projector(speech_dim, self.d_model) if speech_dim else None
+
+    # ---- build the merged (B, L', d) embedding sequence by splicing modality features ----
+    def build_inputs_embeds(self, input_ids: Tensor, image_features: Tensor | None = None,
+                            audio_features: Tensor | None = None, targets: Tensor | None = None,
+                            video_features: Tensor | None = None, speech_features: Tensor | None = None):
+        B, L = input_ids.shape
+        dev = input_ids.device
+        img_proj = self.mm_projector(image_features) if image_features is not None else None  # (B,Ni,d)
+        vid_proj = self.mm_projector(video_features) if video_features is not None else None   # video reuses mm_projector
+        aud_proj = (self.audio_projector(audio_features)
+                    if (audio_features is not None and self.audio_projector is not None) else None)
+        spk_proj = (self.speech_projector(speech_features)
+                    if (speech_features is not None and self.speech_projector is not None) else None)
+
+        seqs_e, seqs_t = [], []
+        for b in range(B):
+            ids = input_ids[b]
+            text_emb = self.llm.embed(ids)                              # (L, d); sentinel rows get sliced out
+            # ordered list of (position, feature_block) for every sentinel in this sample
+            spots = []
+            if img_proj is not None:
+                for p in (ids == IMAGE_TOKEN).nonzero(as_tuple=True)[0]:
+                    spots.append((int(p), img_proj[b]))
+            if vid_proj is not None:
+                for p in (ids == VIDEO_TOKEN).nonzero(as_tuple=True)[0]:
+                    spots.append((int(p), vid_proj[b]))
+            if aud_proj is not None:
+                for p in (ids == AUDIO_TOKEN).nonzero(as_tuple=True)[0]:
+                    spots.append((int(p), aud_proj[b]))
+            if spk_proj is not None:
+                for p in (ids == SPEECH_TOKEN).nonzero(as_tuple=True)[0]:
+                    spots.append((int(p), spk_proj[b]))
+            spots.sort(key=lambda s: s[0])
+
+            e_parts, t_parts, prev = [], [], 0
+            for pos, feat in spots:
+                e_parts.append(text_emb[prev:pos])
+                e_parts.append(feat.to(text_emb.dtype))   # encoder/projector may be bf16 under autocast
+                if targets is not None:
+                    t_parts.append(targets[b][prev:pos])
+                    # next-token: the LAST feature predicts the token following the sentinel (no internal
+                    # label shift in our model), so carry targets[pos] onto it; the rest are ignored.
+                    ft = torch.full((feat.shape[0],), IGNORE_INDEX, dtype=targets.dtype, device=dev)
+                    ft[-1] = targets[b][pos]
+                    t_parts.append(ft)
+                prev = pos + 1
+            e_parts.append(text_emb[prev:])
+            if targets is not None:
+                t_parts.append(targets[b][prev:])
+            seqs_e.append(torch.cat(e_parts, dim=0))
+            if targets is not None:
+                seqs_t.append(torch.cat(t_parts, dim=0))
+
+        # right-pad to the longest merged sequence (causal attention -> pads don't affect real tokens)
+        Lmax = max(e.shape[0] for e in seqs_e)
+        inputs_embeds = seqs_e[0].new_zeros(B, Lmax, self.d_model)
+        new_targets = None
+        if targets is not None:
+            new_targets = torch.full((B, Lmax), IGNORE_INDEX, dtype=targets.dtype, device=dev)
+        for b, e in enumerate(seqs_e):
+            inputs_embeds[b, :e.shape[0]] = e
+            if targets is not None:
+                new_targets[b, :seqs_t[b].shape[0]] = seqs_t[b]
+        return inputs_embeds, new_targets
+
+    def forward(self, input_ids: Tensor, image_features: Tensor | None = None,
+                audio_features: Tensor | None = None, targets: Tensor | None = None,
+                video_features: Tensor | None = None, speech_features: Tensor | None = None):
+        inputs_embeds, new_targets = self.build_inputs_embeds(
+            input_ids, image_features, audio_features, targets, video_features, speech_features)
+        return self.llm(inputs_embeds=inputs_embeds, targets=new_targets)
+
+    # ---- param groups: freeze encoders (external), optionally freeze the LLM (stage 1) ----
+    def set_llm_trainable(self, trainable: bool):
+        for p in self.llm.parameters():
+            p.requires_grad = trainable
+
+    def projector_parameters(self):
+        ps = list(self.mm_projector.parameters())
+        if self.audio_projector is not None:
+            ps += list(self.audio_projector.parameters())
+        if self.speech_projector is not None:
+            ps += list(self.speech_projector.parameters())
+        return ps
+
+
+def build_vlm(cfg: ModelConfig, vision_dim: int = 1152, audio_dim: int | None = None) -> MoEVLM:
+    return MoEVLM(MoETransformer(cfg), vision_dim=vision_dim, audio_dim=audio_dim)

hobbylm/vision.py

@@ -0,0 +1,43 @@
+"""Frozen SigLIP2 vision encoder wrapper for the MoE-VLM.
+
+Loads `google/siglip2-so400m-patch14-384` (or any SigLIP/SigLIP2), runs images through its vision
+tower under no_grad, and returns patch features (B, N, hidden) to be projected + spliced by MoEVLM.
+The encoder is frozen in every training stage, so we run it on the fly (precomputing features for
+558K images would be ~900 GB). Lazy transformers import so the module is CPU-importable.
+"""
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+
+SIGLIP2_ID = "google/siglip2-so400m-patch14-384"
+
+
+class SiglipVision(nn.Module):
+    def __init__(self, model_id: str = SIGLIP2_ID, device="cuda", dtype=torch.bfloat16):
+        super().__init__()
+        from transformers import AutoModel, AutoProcessor
+        self.processor = AutoProcessor.from_pretrained(model_id)
+        full = AutoModel.from_pretrained(model_id, torch_dtype=dtype)
+        self.vision = full.vision_model.to(device).eval()
+        for p in self.vision.parameters():
+            p.requires_grad = False
+        self.device = device
+        self.dtype = dtype
+        self.hidden = self.vision.config.hidden_size
+
+    @torch.no_grad()
+    def preprocess(self, images) -> torch.Tensor:
+        """images: list of PIL.Image -> pixel_values (B, 3, H, W) on device."""
+        px = self.processor(images=images, return_tensors="pt").pixel_values
+        return px.to(self.device, self.dtype)
+
+    @torch.no_grad()
+    def encode_pixels(self, pixel_values: torch.Tensor) -> torch.Tensor:
+        """pixel_values (B,3,H,W) -> patch features (B, N, hidden)."""
+        return self.vision(pixel_values=pixel_values).last_hidden_state
+
+    @torch.no_grad()
+    def encode(self, images) -> torch.Tensor:
+        """list of PIL.Image -> (B, N, hidden) patch features."""
+        return self.encode_pixels(self.preprocess(images))

requirements.txt

@@ -0,0 +1,11 @@
+torch
+gradio>=4.44
+transformers>=4.49
+diffusers==0.32.2
+accelerate
+safetensors
+tiktoken
+huggingface_hub>=0.25
+pillow
+numpy
+sentencepiece