Upload model_arch.py with huggingface_hub

model_arch.py

@@ -1,343 +1,152 @@
-#!/usr/bin/env python3
-"""
-Multilingual 3B GPT — SFT Training
-
-Fine-tunes the base model on instruction data (Aya + Bactrian-X + FLORES translations).
-Uses the same architecture as pretraining with LoRA-free full fine-tuning
-(model is 3B params, fits in 24GB A10G in bf16).
-
-Usage:
-    python train_sft_3b.py --checkpoint /path/to/best_model.pt \
-        --tokenizer /path/to/multilingual_32k.model \
-        --data-dir /path/to/sft_data/ \
-        --output /path/to/sft_model.pt
-"""
-
-import os, sys, json, math, time, random, argparse
-sys.stdout.reconfigure(line_buffering=True)
-import gc
-import numpy as np
+"""Shared model architecture for multilingual 3B GPT — must match training exactly."""
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-import sentencepiece as spm
+import math
 
-# ============ MODEL (must match training) ============
 VOCAB_SIZE = 32000
 DIM = 3072
 DEPTH = 26
 N_HEADS = 24
+HEAD_DIM = DIM // N_HEADS  # 128
 MAX_SEQ_LEN = 2048
-ROPE_THETA = 10000
+ROPE_THETA = 10000.0
+HIDDEN_DIM = ((int(2 * DIM * 4 / 3) + 63) // 64) * 64  # SwiGLU hidden
+
 
 class RMSNorm(nn.Module):
     def __init__(self, dim, eps=1e-6):
         super().__init__()
-        self.weight = nn.Parameter(torch.ones(dim))
         self.eps = eps
-    def forward(self, x):
-        return x * torch.rsqrt(x.float().pow(2).mean(-1, keepdim=True) + self.eps).type_as(x) * self.weight
+        self.weight = nn.Parameter(torch.ones(dim))
 
-class SwiGLU(nn.Module):
-    def __init__(self, dim, hidden_dim):
-        super().__init__()
-        self.gate = nn.Linear(dim, hidden_dim, bias=False)
-        self.up = nn.Linear(dim, hidden_dim, bias=False)
-        self.down = nn.Linear(hidden_dim, dim, bias=False)
     def forward(self, x):
-        return self.down(F.silu(self.gate(x)) * self.up(x))
+        norm = x.float().pow(2).mean(-1, keepdim=True).add(self.eps).rsqrt()
+        return (x.float() * norm).type_as(x) * self.weight
+
+
+def precompute_freqs_cis(dim, max_seq_len, theta=ROPE_THETA):
+    freqs = 1.0 / (theta ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim))
+    t = torch.arange(max_seq_len, dtype=torch.float32)
+    freqs = torch.outer(t, freqs)
+    return torch.polar(torch.ones_like(freqs), freqs)
+
 
-def apply_rope(x, cos, sin):
-    x1, x2 = x[..., ::2], x[..., 1::2]
-    return torch.stack((x1*cos - x2*sin, x1*sin + x2*cos), dim=-1).flatten(-2)
+def apply_rotary_emb(x, freqs_cis):
+    # x: (B, n_heads, S, head_dim)
+    B, H, S, D = x.shape
+    x_complex = torch.view_as_complex(x.float().reshape(B, H, S, D // 2, 2))
+    freqs = freqs_cis[:S].unsqueeze(0).unsqueeze(1)  # (1, 1, S, D//2)
+    x_rot = torch.view_as_real(x_complex * freqs).reshape(B, H, S, D)
+    return x_rot.type_as(x)
 
-class Attention(nn.Module):
+
+class FusedAttention(nn.Module):
     def __init__(self, dim, n_heads):
         super().__init__()
         self.n_heads = n_heads
         self.head_dim = dim // n_heads
-        self.qkv = nn.Linear(dim, 3*dim, bias=False)
-        self.proj = nn.Linear(dim, dim, bias=False)
-    def forward(self, x, cos, sin):
-        B, T, C = x.shape
-        qkv = self.qkv(x).reshape(B, T, 3, self.n_heads, self.head_dim).permute(2, 0, 3, 1, 4)
-        q, k, v = qkv[0], qkv[1], qkv[2]
-        q, k = apply_rope(q, cos, sin), apply_rope(k, cos, sin)
-        y = F.scaled_dot_product_attention(q, k, v, is_causal=True)
-        return self.proj(y.transpose(1, 2).contiguous().view(B, T, C))
+        self.qkv = nn.Linear(dim, 3 * dim, bias=False)
+        self.out_proj = nn.Linear(dim, dim, bias=False)
+
+    def forward(self, x, freqs_cis, mask=None):
+        B, S, D = x.shape
+        qkv = self.qkv(x).reshape(B, S, 3, self.n_heads, self.head_dim)
+        q, k, v = qkv[:, :, 0], qkv[:, :, 1], qkv[:, :, 2]
+        q = q.transpose(1, 2)  # (B, H, S, D)
+        k = k.transpose(1, 2)
+        v = v.transpose(1, 2)
+        q = apply_rotary_emb(q, freqs_cis)
+        k = apply_rotary_emb(k, freqs_cis)
+        # Scaled dot-product attention
+        scale = math.sqrt(self.head_dim)
+        attn = (q @ k.transpose(-2, -1)) / scale
+        if mask is not None:
+            attn = attn + mask
+        attn = F.softmax(attn, dim=-1)
+        out = (attn @ v).transpose(1, 2).reshape(B, S, D)
+        return self.out_proj(out)
+
+
+class SwiGLUFFN(nn.Module):
+    def __init__(self, dim, hidden_dim):
+        super().__init__()
+        self.w1 = nn.Linear(dim, hidden_dim, bias=False)
+        self.w2 = nn.Linear(hidden_dim, dim, bias=False)
+        self.w3 = nn.Linear(dim, hidden_dim, bias=False)
 
-class Block(nn.Module):
-    def __init__(self, dim, n_heads, mlp_dim):
+    def forward(self, x):
+        return self.w2(F.silu(self.w1(x)) * self.w3(x))
+
+
+class TransformerBlock(nn.Module):
+    def __init__(self, dim, n_heads, hidden_dim):
         super().__init__()
-        self.ln1 = RMSNorm(dim)
-        self.attn = Attention(dim, n_heads)
-        self.ln2 = RMSNorm(dim)
-        self.mlp = SwiGLU(dim, mlp_dim)
-    def forward(self, x, cos, sin):
-        x = x + self.attn(self.ln1(x), cos, sin)
-        x = x + self.mlp(self.ln2(x))
+        self.attn_norm = RMSNorm(dim)
+        self.attn = FusedAttention(dim, n_heads)
+        self.ffn_norm = RMSNorm(dim)
+        self.ffn = SwiGLUFFN(dim, hidden_dim)
+
+    def forward(self, x, freqs_cis, mask=None):
+        x = x + self.attn(self.attn_norm(x), freqs_cis, mask)
+        x = x + self.ffn(self.ffn_norm(x))
         return x
 
-class GPT(nn.Module):
+
+class MultilingualGPT(nn.Module):
     def __init__(self):
         super().__init__()
         self.tok_emb = nn.Embedding(VOCAB_SIZE, DIM)
-        mlp_dim = ((int(2 * DIM * 4 / 3) + 63) // 64) * 64
-        self.blocks = nn.ModuleList([Block(DIM, N_HEADS, mlp_dim) for _ in range(DEPTH)])
-        self.ln_f = RMSNorm(DIM)
+        self.layers = nn.ModuleList([
+            TransformerBlock(DIM, N_HEADS, HIDDEN_DIM) for _ in range(DEPTH)
+        ])
+        self.norm = RMSNorm(DIM)
         self.head = nn.Linear(DIM, VOCAB_SIZE, bias=False)
+        # Tied embeddings
         self.head.weight = self.tok_emb.weight
-        hd = DIM // N_HEADS
-        freqs = 1.0 / (ROPE_THETA ** (torch.arange(0, hd, 2).float() / hd))
-        angles = torch.outer(torch.arange(MAX_SEQ_LEN).float(), freqs)
-        self.register_buffer('rope_cos', angles.cos())
-        self.register_buffer('rope_sin', angles.sin())
-
-    def forward(self, idx):
-        B, T = idx.shape
-        x = self.tok_emb(idx)
-        cos = self.rope_cos[:T][None, None]
-        sin = self.rope_sin[:T][None, None]
-        for block in self.blocks:
-            x = block(x, cos, sin)
-        return self.head(self.ln_f(x))
-
-    @torch.no_grad()
-    def generate(self, idx, max_new=200, temp=0.7, top_k=40, rep_penalty=1.2):
-        for _ in range(max_new):
-            idx_c = idx[:, -MAX_SEQ_LEN:]
-            logits = self(idx_c)[:, -1, :]
-            if rep_penalty > 1.0:
-                for token_id in set(idx[0].tolist()[-50:]):
-                    logits[0, token_id] /= rep_penalty
-            logits = logits / temp
-            if top_k > 0:
-                v, _ = torch.topk(logits, top_k)
-                logits[logits < v[:, [-1]]] = float('-inf')
-            probs = F.softmax(logits, dim=-1)
-            nx = torch.multinomial(probs, 1)
-            idx = torch.cat([idx, nx], dim=1)
-            if nx.item() == 2:
-                break
-        return idx
-
-
-# ============ DATA LOADING ============
-USER_PREFIX = "### User:\n"
-ASSISTANT_PREFIX = "### Assistant:\n"
-TURN_END = "\n\n"
-
-def load_sft_data(data_dir, split='train'):
-    """Load tokenized SFT data."""
-    filepath = os.path.join(data_dir, f'{split}_sft.bin')
-    data = np.fromfile(filepath, dtype=np.uint16)
-    return torch.from_numpy(data.astype(np.int64))
-
-def get_batch(data, batch_size, seq_len, device):
-    """Get a random batch of sequences."""
-    ix = torch.randint(len(data) - seq_len - 1, (batch_size,))
-    x = torch.stack([data[i:i+seq_len] for i in ix]).to(device)
-    y = torch.stack([data[i+1:i+seq_len+1] for i in ix]).to(device)
-    return x, y
-
-
-# ============ TRAINING ============
-def train(args):
-    device = args.device
-    print(f"Device: {device}")
-    
-    # Load tokenizer
-    print(f"Loading tokenizer: {args.tokenizer}")
-    sp = spm.SentencePieceProcessor(args.tokenizer)
-    
-    # Load model
-    print(f"Loading base model: {args.checkpoint}")
-    model = GPT()
-    ckpt = torch.load(args.checkpoint, map_location='cpu', weights_only=False)
-    state_dict = ckpt.get('model_state_dict', ckpt.get('model', ckpt))
-    clean_sd = {}
-    for k, v in state_dict.items():
-        k = k.replace('_orig_mod.', '').replace('module.', '')
-        clean_sd[k] = v
-    model.load_state_dict(clean_sd, strict=False)
-    del ckpt, state_dict, clean_sd
-    gc.collect()
-    
-    model = model.to(device).train()
-    # Use bf16 for memory efficiency
-    model = model.to(torch.bfloat16)
-    
-    param_count = sum(p.numel() for p in model.parameters())
-    print(f"Model loaded: {param_count/1e9:.2f}B parameters")
-    
-    # Load data
-    print(f"Loading SFT data from: {args.data_dir}")
-    train_data = load_sft_data(args.data_dir, 'train')
-    val_data = load_sft_data(args.data_dir, 'val')
-    print(f"Train: {len(train_data)} tokens, Val: {len(val_data)} tokens")
-    
-    # Optimizer — 8-bit Adam for memory efficiency (halves optimizer states)
-    try:
-        import bitsandbytes as bnb
-        optimizer = bnb.optim.AdamW8bit(
-            model.parameters(),
-            lr=args.lr,
-            betas=(0.9, 0.95),
-            weight_decay=0.01,
-        )
-        print("Using 8-bit AdamW (bitsandbytes)")
-    except ImportError:
-        optimizer = torch.optim.AdamW(
-            model.parameters(),
-            lr=args.lr,
-            betas=(0.9, 0.95),
-            weight_decay=0.01,
-        )
-        print("Using standard AdamW")
-    
-    # Cosine schedule with warmup
-    def get_lr(step):
-        if step < args.warmup_steps:
-            return args.lr * step / args.warmup_steps
-        decay_ratio = (step - args.warmup_steps) / (args.max_steps - args.warmup_steps)
-        return args.lr * 0.1 + 0.9 * args.lr * 0.5 * (1 + math.cos(math.pi * decay_ratio))
-    
-    # Enable gradient checkpointing to save VRAM
-    for block in model.blocks:
-        block._gradient_checkpointing = True
-    original_block_forward = Block.forward
-    def checkpointed_forward(self, x, cos, sin):
-        if self.training and hasattr(self, '_gradient_checkpointing') and self._gradient_checkpointing:
-            return torch.utils.checkpoint.checkpoint(original_block_forward, self, x, cos, sin, use_reentrant=False)
-        return original_block_forward(self, x, cos, sin)
-    Block.forward = checkpointed_forward
-    
-    # Training loop
-    best_val_loss = float('inf')
-    grad_accum = args.grad_accum
-    print(f"\nStarting SFT training for {args.max_steps} steps...")
-    print(f"Batch size: {args.batch_size} x {grad_accum} accum = {args.batch_size * grad_accum} effective, Seq len: {MAX_SEQ_LEN}, LR: {args.lr}")
-    
-    t0 = time.time()
-    for step in range(1, args.max_steps + 1):
-        # LR schedule
-        lr = get_lr(step)
-        for pg in optimizer.param_groups:
-            pg['lr'] = lr
-        
-        # Gradient accumulation
-        optimizer.zero_grad(set_to_none=True)
-        accum_loss = 0.0
-        for micro in range(grad_accum):
-            x, y = get_batch(train_data, args.batch_size, MAX_SEQ_LEN, device)
-            with torch.autocast(device_type='cuda', dtype=torch.bfloat16):
-                logits = model(x)
-                loss = F.cross_entropy(logits.view(-1, VOCAB_SIZE), y.view(-1)) / grad_accum
-            loss.backward()
-            accum_loss += loss.item()
-        
-        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
-        optimizer.step()
-        loss = type('obj', (object,), {'item': lambda self: accum_loss})()  # For logging
-        
-        # Logging
-        if step % 10 == 0:
-            elapsed = time.time() - t0
-            tps = step * args.batch_size * grad_accum * MAX_SEQ_LEN / elapsed
-            print(f"Step {step}/{args.max_steps} | Loss: {accum_loss:.4f} | LR: {lr:.6f} | TPS: {tps:.0f} | {elapsed:.0f}s")
-        
-        # Eval
-        if step % args.eval_every == 0 or step == args.max_steps:
-            model.eval()
-            val_losses = []
-            for _ in range(20):
-                x, y = get_batch(val_data, args.batch_size, MAX_SEQ_LEN, device)
-                with torch.no_grad(), torch.autocast(device_type='cuda', dtype=torch.bfloat16):
-                    logits = model(x)
-                    val_loss = F.cross_entropy(logits.view(-1, VOCAB_SIZE), y.view(-1))
-                val_losses.append(val_loss.item())
-            avg_val = sum(val_losses) / len(val_losses)
-            print(f"  📊 Val loss: {avg_val:.4f} {'(NEW BEST!)' if avg_val < best_val_loss else ''}")
-            
-            if avg_val < best_val_loss:
-                best_val_loss = avg_val
-                torch.save({
-                    'model_state_dict': model.state_dict(),
-                    'step': step,
-                    'val_loss': avg_val,
-                    'config': {
-                        'vocab_size': VOCAB_SIZE, 'dim': DIM, 'depth': DEPTH,
-                        'n_heads': N_HEADS, 'max_seq_len': MAX_SEQ_LEN,
-                    }
-                }, args.output)
-                print(f"  💾 Best model saved to {args.output}")
-            
-            model.train()
-        
-        # Generate samples periodically
-        if step % args.sample_every == 0 or step == args.max_steps:
-            model.eval()
-            prompts = [
-                ("EN", "### User:\nWhat is the capital of France?\n\n### Assistant:\n"),
-                ("HE", "### User:\nמה בירת צרפת?\n\n### Assistant:\n"),
-                ("AR", "### User:\nما هي عاصمة فرنسا؟\n\n### Assistant:\n"),
-                ("FA", "### User:\nپایتخت فرانسه کجاست؟\n\n### Assistant:\n"),
-                ("TRANSLATE", "### User:\nTranslate the following Hebrew text to English:\nשלום עולם, איך אתה היום?\n\n### Assistant:\n"),
-            ]
-            print(f"\n  🔤 Generation samples (step {step}):")
-            for label, prompt in prompts:
-                ids = torch.tensor([sp.encode(prompt)], device=device, dtype=torch.long)
-                with torch.no_grad(), torch.autocast(device_type='cuda', dtype=torch.bfloat16):
-                    out = model.generate(ids, max_new=100, temp=0.7, top_k=40)
-                text = sp.decode(out[0].tolist())
-                # Just show the assistant response
-                if "### Assistant:" in text:
-                    response = text.split("### Assistant:")[-1].strip()[:200]
-                else:
-                    response = text[len(prompt):].strip()[:200]
-                print(f"    [{label}] {response}")
-            print()
-            model.train()
-    
-    # Final save
-    elapsed = time.time() - t0
-    print(f"\n{'='*60}")
-    print(f"SFT TRAINING COMPLETE")
-    print(f"Steps: {args.max_steps}, Time: {elapsed:.0f}s ({elapsed/60:.1f}min)")
-    print(f"Best val loss: {best_val_loss:.4f}")
-    print(f"Model saved to: {args.output}")
-    print(f"{'='*60}")
-    
-    # Upload to S3
-    print("Uploading to S3...")
-    os.system(f"aws s3 cp {args.output} s3://autoresearch-dashboard-196766918360/multilingual-7b/checkpoints/3b-v1-fsdp/sft_model.pt --quiet")
-    os.system(f"aws s3 cp /tmp/sft/sft.log s3://autoresearch-dashboard-196766918360/multilingual-7b/eval/sft_3b.log --quiet 2>/dev/null")
-    print("Done!")
-
-
-def main():
-    parser = argparse.ArgumentParser()
-    parser.add_argument('--checkpoint', required=True)
-    parser.add_argument('--tokenizer', required=True)
-    parser.add_argument('--data-dir', required=True)
-    parser.add_argument('--output', default='/tmp/sft/sft_model.pt')
-    parser.add_argument('--device', default='cuda')
-    parser.add_argument('--batch-size', type=int, default=1)  # 1 for 24GB GPU
-    parser.add_argument('--grad-accum', type=int, default=4)  # Effective batch = 4
-    parser.add_argument('--lr', type=float, default=2e-5)
-    parser.add_argument('--max-steps', type=int, default=2000)
-    parser.add_argument('--warmup-steps', type=int, default=100)
-    parser.add_argument('--eval-every', type=int, default=200)
-    parser.add_argument('--sample-every', type=int, default=500)
-    parser.add_argument('--seed', type=int, default=42)
-    args = parser.parse_args()
-    
-    random.seed(args.seed)
-    torch.manual_seed(args.seed)
-    os.makedirs(os.path.dirname(args.output), exist_ok=True)
-    
-    train(args)
-
-
-if __name__ == '__main__':
-    main()
+        # Precompute RoPE
+        self.register_buffer('freqs_cis', precompute_freqs_cis(HEAD_DIM, MAX_SEQ_LEN))
+
+    def forward(self, tokens, targets=None):
+        B, S = tokens.shape
+        x = self.tok_emb(tokens)
+        mask = torch.triu(torch.full((S, S), float('-inf'), device=tokens.device), diagonal=1)
+        mask = mask.unsqueeze(0).unsqueeze(0)  # (1, 1, S, S)
+        for layer in self.layers:
+            x = layer(x, self.freqs_cis, mask)
+        x = self.norm(x)
+        logits = self.head(x)
+        loss = None
+        if targets is not None:
+            loss = F.cross_entropy(logits.view(-1, VOCAB_SIZE), targets.view(-1))
+        return logits, loss
+
+
+def load_model(path, device='cuda'):
+    """Load model from checkpoint, stripping prefixes."""
+    model = MultilingualGPT()
+    ckpt = torch.load(path, map_location='cpu', weights_only=False)
+    state = ckpt.get('model_state_dict', ckpt)
+    # Strip prefixes
+    cleaned = {}
+    for k, v in state.items():
+        new_k = k
+        for prefix in ['_orig_mod.', 'module.']:
+            if new_k.startswith(prefix):
+                new_k = new_k[len(prefix):]
+        cleaned[new_k] = v
+    # Handle tied weights - remove head.weight if present (will be tied)
+    if 'head.weight' in cleaned and 'tok_emb.weight' in cleaned:
+        if torch.equal(cleaned['head.weight'], cleaned['tok_emb.weight']):
+            del cleaned['head.weight']
+    model.load_state_dict(cleaned, strict=False)
+    model = model.to(device).eval()
+    return model
+
+
+def load_tokenizer(path):
+    """Load SentencePiece tokenizer."""
+    import sentencepiece as spm
+    sp = spm.SentencePieceProcessor()
+    sp.Load(path)
+    return sp