ark65m.py

#!/usr/bin/env python3
"""
ARK-65M — ModotAI
Usage:
  python ark65m.py train \
    --data-config data_config.json \
    --bin-dir pretokenized \
    --tokenizer ThingAI/msqark-tokenizer \
    --output-dir checkpoints-ark65m \
    --batch-size 4 --grad-accum 16 --lr 3e-4 \
    --max-steps 49591 --save-every 500 --compile
"""

import os, sys, json, time, argparse
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import Dataset, DataLoader
from torch.utils.checkpoint import checkpoint
from transformers import AutoTokenizer
from dataclasses import dataclass


@dataclass
class ARKConfig:
    vocab_size: int           = 32768
    d_model: int              = 576
    n_heads: int              = 8
    n_kv_heads: int           = 2
    n_layers: int             = 14
    d_ff: int                 = 1536
    max_seq_len: int          = 2048
    rope_theta: float         = 500000.0
    rms_eps: float            = 1e-5
    dropout: float            = 0.0
    router_start_layer: int   = 7
    gradient_checkpointing: bool = True
    @property
    def head_dim(self): return self.d_model // self.n_heads


class RMSNorm(nn.Module):
    def __init__(self, dim, eps=1e-5):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(dim))
        self.eps = eps
    def forward(self, x):
        norm = x.float().pow(2).mean(-1, keepdim=True).add(self.eps).rsqrt()
        return (x.float() * norm).type_as(x) * self.weight

class RotaryEmbedding(nn.Module):
    def __init__(self, dim, theta=500000.0, max_seq_len=2048):
        super().__init__()
        inv_freq = 1.0 / (theta ** (torch.arange(0, dim, 2).float() / dim))
        self.register_buffer("inv_freq", inv_freq, persistent=False)
        t = torch.arange(max_seq_len, dtype=self.inv_freq.dtype, device=self.inv_freq.device)
        freqs = torch.outer(t, self.inv_freq)
        emb = torch.cat([freqs, freqs], dim=-1)
        self.register_buffer("cos_cache", emb.cos(), persistent=False)
        self.register_buffer("sin_cache", emb.sin(), persistent=False)
    def forward(self, positions):
        positions = positions.clamp(0, self.cos_cache.shape[0] - 1)
        return self.cos_cache[positions], self.sin_cache[positions]

def rotate_half(x):
    x1, x2 = x.chunk(2, dim=-1)
    return torch.cat((-x2, x1), dim=-1)

def apply_rotary_pos_emb(q, k, cos, sin):
    cos, sin = cos.unsqueeze(1), sin.unsqueeze(1)
    return (q * cos) + (rotate_half(q) * sin), (k * cos) + (rotate_half(k) * sin)

class SwiGLU(nn.Module):
    def __init__(self, d_model, d_ff):
        super().__init__()
        self.gate_proj = nn.Linear(d_model, d_ff, bias=False)
        self.up_proj   = nn.Linear(d_model, d_ff, bias=False)
        self.down_proj = nn.Linear(d_ff, d_model, bias=False)
    def forward(self, x):
        return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))

class GQAttention(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.n_heads    = config.n_heads
        self.n_kv_heads = config.n_kv_heads
        self.head_dim   = config.head_dim
        self.n_rep      = self.n_heads // self.n_kv_heads
        self.q_proj = nn.Linear(config.d_model, config.n_heads * self.head_dim, bias=False)
        self.k_proj = nn.Linear(config.d_model, config.n_kv_heads * self.head_dim, bias=False)
        self.v_proj = nn.Linear(config.d_model, config.n_kv_heads * self.head_dim, bias=False)
        self.o_proj = nn.Linear(config.n_heads * self.head_dim, config.d_model, bias=False)
        self.q_norm = RMSNorm(self.head_dim, eps=config.rms_eps)
        self.k_norm = RMSNorm(self.head_dim, eps=config.rms_eps)
    def forward(self, x, cos, sin, mask=None):
        B, L, _ = x.shape
        q = self.q_proj(x).view(B, L, self.n_heads, self.head_dim).transpose(1, 2)
        k = self.k_proj(x).view(B, L, self.n_kv_heads, self.head_dim).transpose(1, 2)
        v = self.v_proj(x).view(B, L, self.n_kv_heads, self.head_dim).transpose(1, 2)
        q, k = self.q_norm(q), self.k_norm(k)
        q, k = apply_rotary_pos_emb(q, k, cos, sin)
        k = k.repeat_interleave(self.n_rep, dim=1)
        v = v.repeat_interleave(self.n_rep, dim=1)
        out = F.scaled_dot_product_attention(q, k, v, attn_mask=mask, is_causal=(mask is None))
        return self.o_proj(out.transpose(1, 2).contiguous().view(B, L, -1))

class MSAAttention(nn.Module):
    def __init__(self, config, layer_idx):
        super().__init__()
        self.n_heads    = config.n_heads
        self.n_kv_heads = config.n_kv_heads
        self.head_dim   = config.head_dim
        self.n_rep      = self.n_heads // self.n_kv_heads
        self.window_size = 512
        self.n_global    = 64
        self.q_proj = nn.Linear(config.d_model, config.n_heads * self.head_dim, bias=False)
        self.k_proj = nn.Linear(config.d_model, config.n_kv_heads * self.head_dim, bias=False)
        self.v_proj = nn.Linear(config.d_model, config.n_kv_heads * self.head_dim, bias=False)
        self.o_proj = nn.Linear(config.n_heads * self.head_dim, config.d_model, bias=False)
        self.q_norm = RMSNorm(self.head_dim, eps=config.rms_eps)
        self.k_norm = RMSNorm(self.head_dim, eps=config.rms_eps)
        self.router_q = nn.Linear(config.d_model, self.head_dim, bias=False)
        self.router_k = nn.Linear(config.d_model, self.head_dim, bias=False)
    def forward(self, x, cos, sin, mask=None):
        B, L, _ = x.shape
        q = self.q_proj(x).view(B, L, self.n_heads, self.head_dim).transpose(1, 2)
        k = self.k_proj(x).view(B, L, self.n_kv_heads, self.head_dim).transpose(1, 2)
        v = self.v_proj(x).view(B, L, self.n_kv_heads, self.head_dim).transpose(1, 2)
        q, k = self.q_norm(q), self.k_norm(k)
        q, k = apply_rotary_pos_emb(q, k, cos, sin)
        k = k.repeat_interleave(self.n_rep, dim=1)
        v = v.repeat_interleave(self.n_rep, dim=1)
        if L <= self.window_size + self.n_global:
            out = F.scaled_dot_product_attention(q, k, v, is_causal=True)
        else:
            out = self._sparse_attention(q, k, v, L)
        return self.o_proj(out.transpose(1, 2).contiguous().view(B, L, -1))
    def _sparse_attention(self, q, k, v, L):
        output = torch.zeros_like(q)
        for i in range(0, L, self.window_size):
            end         = min(i + self.window_size, L)
            q_chunk     = q[:, :, i:end]
            local_start = max(0, i - self.window_size)
            k_local     = k[:, :, local_start:end]
            v_local     = v[:, :, local_start:end]
            q_pos = torch.arange(i, end, device=q.device)
            if i > self.n_global:
                k_combined = torch.cat([k[:, :, :self.n_global], k_local], dim=2)
                v_combined = torch.cat([v[:, :, :self.n_global], v_local], dim=2)
                k_pos = torch.cat([
                    torch.arange(self.n_global, device=q.device),
                    torch.arange(local_start, end, device=q.device),
                ])
            else:
                k_combined, v_combined = k_local, v_local
                k_pos = torch.arange(local_start, end, device=q.device)
            causal    = q_pos.unsqueeze(1) >= k_pos.unsqueeze(0)
            attn_mask = torch.where(causal, 0.0, float('-inf')).unsqueeze(0).unsqueeze(0)
            output[:, :, i:end] = F.scaled_dot_product_attention(
                q_chunk, k_combined, v_combined, attn_mask=attn_mask
            )
        return output

class TransformerBlock(nn.Module):
    def __init__(self, config, layer_idx):
        super().__init__()
        self.attn  = MSAAttention(config, layer_idx) if layer_idx >= config.router_start_layer else GQAttention(config)
        self.ffn   = SwiGLU(config.d_model, config.d_ff)
        self.norm1 = RMSNorm(config.d_model, eps=config.rms_eps)
        self.norm2 = RMSNorm(config.d_model, eps=config.rms_eps)
    def forward(self, x, cos, sin, mask=None):
        x = x + self.attn(self.norm1(x), cos, sin, mask)
        x = x + self.ffn(self.norm2(x))
        return x

class ARK65M(nn.Module):
    def __init__(self, config: ARKConfig):
        super().__init__()
        self.config  = config
        self.embed   = nn.Embedding(config.vocab_size, config.d_model)
        self.rope    = RotaryEmbedding(config.head_dim, config.rope_theta, config.max_seq_len)
        self.layers  = nn.ModuleList([TransformerBlock(config, i) for i in range(config.n_layers)])
        self.norm    = RMSNorm(config.d_model, eps=config.rms_eps)
        self.lm_head = nn.Linear(config.d_model, config.vocab_size, bias=False)
        self.lm_head.weight = self.embed.weight
        self.apply(self._init_weights)
    def _init_weights(self, module):
        if isinstance(module, nn.Linear):
            nn.init.normal_(module.weight, mean=0.0, std=0.02)
        elif isinstance(module, nn.Embedding):
            nn.init.normal_(module.weight, mean=0.0, std=0.02)
    def forward(self, input_ids, labels=None):
        B, L = input_ids.shape
        x = self.embed(input_ids)
        positions = torch.arange(L, device=input_ids.device).unsqueeze(0).expand(B, -1)
        cos, sin = self.rope(positions)
        for layer in self.layers:
            if self.config.gradient_checkpointing and self.training:
                x = checkpoint(layer, x, cos, sin, None, use_reentrant=False)
            else:
                x = layer(x, cos, sin)
        logits = self.lm_head(self.norm(x))
        loss = None
        if labels is not None:
            loss = F.cross_entropy(
                logits[:, :-1].contiguous().view(-1, self.config.vocab_size),
                labels[:, 1:].contiguous().view(-1),
                ignore_index=-100,
            )
        return {"logits": logits, "loss": loss}
    def count_parameters(self):
        total = sum(p.numel() for p in self.parameters())
        embed = self.embed.weight.numel()
        return {
            "total": total, "embedding": embed,
            "embed_pct": f"{embed/total*100:.1f}%",
            "transformer": total - embed,
        }

def _unwrap(model):
    return model._orig_mod if hasattr(model, '_orig_mod') else model


def _safe_filename(name):
    return name.replace(' ', '_').replace('/', '_').replace('\\', '_')

class PreTokenizedDataset(Dataset):
    def __init__(self, data_config, bin_dir, max_len=2048):
        super().__init__()
        self.max_len = max_len
        with open(data_config, 'r', encoding='utf-8') as f:
            sources_cfg = json.load(f)
        self.sources = []
        for src in sources_cfg:
            safe_name = _safe_filename(src["name"])
            bin_path  = os.path.join(bin_dir, f"{safe_name}.bin")
            if not os.path.exists(bin_path):
                print(f"Warning: {bin_path} not found, skipping {src['name']}")
                continue
            file_bytes = os.path.getsize(bin_path)
            dtype      = np.uint16 if file_bytes % 2 == 0 else np.uint32
            memmap     = np.memmap(bin_path, dtype=dtype, mode='r')
            if len(memmap) <= max_len:
                continue
            self.sources.append({
                "name": src["name"], "weight": float(src["weight"]),
                "memmap": memmap, "length": len(memmap),
            })
        if not self.sources:
            raise ValueError("No valid sources!")
        weights = np.array([s["weight"] for s in self.sources], dtype=np.float64)
        weights /= weights.sum()
        self.cum_weights = np.cumsum(weights)
        self._len = int(sum(s["length"] for s in self.sources) // max_len * 2)
    def __len__(self): return self._len
    def __getitem__(self, idx):
        rng    = np.random.RandomState(idx)
        source = self.sources[int(np.searchsorted(self.cum_weights, rng.rand()))]
        start  = rng.randint(0, source["length"] - self.max_len)
        tokens = source["memmap"][start: start + self.max_len]
        input_ids = torch.from_numpy(tokens.astype(np.int64)).long()
        return {"input_ids": input_ids, "labels": input_ids.clone()}


def train():
    parser = argparse.ArgumentParser()
    parser.add_argument("--data-config",   required=True)
    parser.add_argument("--bin-dir",       default="pretokenized")
    parser.add_argument("--tokenizer",     default="ThingAI/msqark-tokenizer")
    parser.add_argument("--output-dir",    default="checkpoints-ark65m")
    parser.add_argument("--batch-size",    type=int,   default=4)
    parser.add_argument("--grad-accum",    type=int,   default=16)
    parser.add_argument("--lr",            type=float, default=3e-4)
    parser.add_argument("--max-len",       type=int,   default=2048)
    parser.add_argument("--warmup-steps",  type=int,   default=500)
    parser.add_argument("--max-steps",     type=int,   default=0)
    parser.add_argument("--save-every",    type=int,   default=500)
    parser.add_argument("--num-workers",   type=int,   default=4)
    parser.add_argument("--resume",        type=str,   default=None)
    parser.add_argument("--compile",       action="store_true")
    parser.add_argument("--no-checkpoint", action="store_true")
    args = parser.parse_args()

    torch.backends.cuda.matmul.allow_tf32 = True
    torch.backends.cudnn.allow_tf32 = True
    torch.set_float32_matmul_precision('high')
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    os.makedirs(args.output_dir, exist_ok=True)

    tokenizer = AutoTokenizer.from_pretrained(args.tokenizer, trust_remote_code=True)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token

    config = ARKConfig(
        vocab_size=max(tokenizer.vocab_size, 32768),
        max_seq_len=args.max_len,
        gradient_checkpointing=not args.no_checkpoint,
    )
    model = ARK65M(config).to(device)

    if args.compile:
        try:
            print("Compiling with torch.compile...")
            model = torch.compile(model, mode="default")
            print("✅ Compiled.")
        except Exception as e:
            print(f"⚠️ torch.compile failed: {e}")

    params = _unwrap(model).count_parameters()
    print(f"\n{'═'*55}")
    print(f"  ARK-65M — ModotAI")
    print(f"{'═'*55}")
    print(f"  Parameters:     {params['total']:,}")
    print(f"  Embedding:      {params['embed_pct']}")
    print(f"  Transformer:    {params['transformer']:,}")
    print(f"  Context:        {config.max_seq_len:,} tokens")
    print(f"  Layers:         {config.n_layers} ({config.router_start_layer} GQA + {config.n_layers - config.router_start_layer} MSA)")
    print(f"  d_model:        {config.d_model}")
    print(f"  Batch eff.:     {args.batch_size} x {args.grad_accum} = {args.batch_size * args.grad_accum}")
    print(f"  LR:             {args.lr}")
    print(f"  Grad ckpt:      {config.gradient_checkpointing}")
    print(f"  torch.compile:  {args.compile}")
    print(f"  Device:         {device}")
    if torch.cuda.is_available():
        print(f"  GPU:            {torch.cuda.get_device_name()}")
        print(f"  VRAM:           {torch.cuda.get_device_properties(0).total_memory/1e9:.1f} GB")
    print(f"{'═'*55}\n")

    dataset = PreTokenizedDataset(args.data_config, args.bin_dir, max_len=args.max_len)
    loader  = DataLoader(
        dataset, batch_size=args.batch_size, shuffle=True,
        num_workers=args.num_workers,
        persistent_workers=args.num_workers > 0,
        pin_memory=True, drop_last=True,
    )

    print(f"Dataset: {len(dataset):,} samples")
    print(f"Sources loaded: {len(dataset.sources)}")
    for s in dataset.sources:
        print(f"  • {s['name']:.<35} {s['length']:>12,} tokens")
    print()

    optimizer = torch.optim.AdamW(
        model.parameters(), lr=args.lr,
        betas=(0.9, 0.95), weight_decay=0.1,
    )
    def lr_lambda(step):
        if step < args.warmup_steps:
            return step / max(args.warmup_steps, 1)
        return 1.0
    scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda)
    scaler    = torch.amp.GradScaler("cuda")

    step        = 0
    tokens_seen = 0
    if args.resume:
        print(f"🔄 Resuming from {args.resume}")
        ckpt = torch.load(args.resume, map_location=device, weights_only=False)
        _unwrap(model).load_state_dict(ckpt["model_state_dict"])
        optimizer.load_state_dict(ckpt["optimizer_state_dict"])
        scheduler.load_state_dict(ckpt["scheduler_state_dict"])
        scaler.load_state_dict(ckpt["scaler_state_dict"])
        step        = ckpt.get("step", 0)
        tokens_seen = ckpt.get("tokens_seen", 0)
        print(f"   ✅ Resumed at step {step:,} | {tokens_seen/1e9:.3f}B tokens | LR {scheduler.get_last_lr()[0]:.2e}\n")

    model.train()
    loss_accum   = 0.0
    log_steps    = 0
    start_time   = time.time()
    last_log     = start_time
    tokens_inst  = 0
    done         = False

    print("Starting training...\n")

    while not done:
        for batch_idx, batch in enumerate(loader):
            input_ids = batch["input_ids"].to(device, non_blocking=True)
            labels    = batch["labels"].to(device, non_blocking=True)

            with torch.amp.autocast("cuda", dtype=torch.bfloat16):
                out  = model(input_ids, labels=labels)
                loss = out["loss"] / args.grad_accum

            scaler.scale(loss).backward()
            loss_accum += loss.item()

            batch_tokens = input_ids.numel()
            tokens_seen += batch_tokens
            tokens_inst += batch_tokens

            if (batch_idx + 1) % args.grad_accum == 0:
                scaler.unscale_(optimizer)
                torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
                scaler.step(optimizer)
                scaler.update()
                optimizer.zero_grad()
                scheduler.step()
                step     += 1
                log_steps += 1

                if step % 10 == 0:
                    now           = time.time()
                    elapsed_total = now - start_time
                    elapsed_inst  = now - last_log
                    avg_loss      = loss_accum / log_steps
                    tok_avg       = tokens_seen / elapsed_total if elapsed_total > 0 else 0
                    tok_inst      = tokens_inst / elapsed_inst  if elapsed_inst  > 0 else 0
                    lr_now        = scheduler.get_last_lr()[0]
                    vram          = torch.cuda.memory_allocated() / 1e9
                    print(
                        f"step {step:>6} │ loss {avg_loss:.4f} │ lr {lr_now:.2e} │ "
                        f"avg {tok_avg:>7,.0f} tok/s │ inst {tok_inst:>7,.0f} tok/s │ "
                        f"{tokens_seen/1e9:.3f}B tok │ VRAM {vram:.1f}GB │ {elapsed_total:.0f}s"
                    )
                    loss_accum  = 0.0
                    log_steps   = 0
                    tokens_inst = 0
                    last_log    = now

                if step % args.save_every == 0:
                    path = os.path.join(args.output_dir, f"step-{step}.pt")
                    torch.save({
                        "model_state_dict":     _unwrap(model).state_dict(),
                        "optimizer_state_dict": optimizer.state_dict(),
                        "scheduler_state_dict": scheduler.state_dict(),
                        "scaler_state_dict":    scaler.state_dict(),
                        "config":               config.__dict__,
                        "step":                 step,
                        "tokens_seen":          tokens_seen,
                    }, path)
                    print(f"  💾 Saved: {path}")

                if args.max_steps > 0 and step >= args.max_steps:
                    done = True
                    break

        if done:
            break

    final = os.path.join(args.output_dir, "final.pt")
    torch.save({
        "model_state_dict": _unwrap(model).state_dict(),
        "config":           config.__dict__,
        "step":             step,
        "tokens_seen":      tokens_seen,
    }, final)
    elapsed = time.time() - start_time
    print(f"\n{'═'*55}")
    print(f"  Done!")
    print(f"  Steps:   {step:,}")
    print(f"  Tokens:  {tokens_seen/1e9:.3f}B")
    print(f"  Time:    {elapsed/3600:.1f}h")
    print(f"  Saved:   {final}")
    print(f"{'═'*55}")


if __name__ == "__main__":
    if len(sys.argv) > 1 and sys.argv[1] == "train":
        sys.argv.pop(1)
        train()
    else:
        config = ARKConfig(gradient_checkpointing=False)
        model  = ARK65M(config)
        params = model.count_parameters()
        print(f"{'═'*45}")
        print(f"  ARK-65M — ModotAI")
        print(f"{'═'*45}")
        for k, v in params.items():
            if isinstance(v, int):
                print(f"  {k:.<25} {v:>12,}")
            else:
                print(f"  {k:.<25} {v:>12}")
        x   = torch.randint(0, config.vocab_size, (2, 128))
        out = model(x, labels=x)
        print(f"\n  Test forward OK")
        print(f"  Loss: {out['loss'].item():.4f}")
        print(f"{'═'*45}")