Upload train_1b_v3.py with huggingface_hub

train_1b_v3.py

@@ -0,0 +1,427 @@
+#!/usr/bin/env python3
+"""
+CogNet-1B Training Script — Full 1B on 16GB VRAM
+===================================================
+Strategy: CPU optimizer offloading + gradient checkpointing
+- Model FP16 on GPU: ~2GB
+- Gradients on GPU: ~2GB  
+- Activations (checkpointed): ~2-3GB
+- Optimizer states on CPU: ~8GB RAM
+- Peak GPU: ~7-8GB → fits easily in 16GB
+"""
+
+import os, sys, time, math, json, torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.utils.data import Dataset, DataLoader
+
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+from cognet_1b import CogNet1B
+
+# ─── Full 1B Config ─────────────────────────────────────────────
+VOCAB_SIZE = 136
+HIDDEN_DIM = 2048
+NUM_BLOCKS = 13
+NUM_CHANNELS = 8
+CHANNEL_DIM = 256
+FF_DIM = 4096
+MAX_SEQ_LEN = 256
+WORKING_SLOTS = 64
+EPISODIC_SLOTS = 128
+SEMANTIC_SLOTS = 256
+
+BATCH_SIZE = 2
+GRAD_ACCUM = 32          # Effective BS = 64
+MAX_STEPS = 100000
+LR = 2e-4
+WEIGHT_DECAY = 0.1
+WARMUP_STEPS = 2000
+MIN_LR = 1e-5
+SEQ_LEN = MAX_SEQ_LEN
+
+EVAL_EVERY = 2000
+SAVE_EVERY = 5000
+LOG_EVERY = 10
+
+# ─── Char Tokenizer ─────────────────────────────────────────────
+class CharTokenizer:
+    def __init__(self):
+        self.chars = sorted(set(
+            [chr(i) for i in range(32, 127)]
+            + list('\u00e0\u00e2\u00e4\u00e9\u00e8\u00ea\u00eb\u00ef\u00ee\u00f4\u00f9\u00fb\u00fc\u00ff\u00e7\u0153\u00e6'
+                   '\u00c0\u00c2\u00c4\u00c9\u00c8\u00ca\u00cb\u00cf\u00ce\u00d4\u00d9\u00db\u00dc\u0178\u00c7\u0152\u00c6')
+            + list('\u00eb\u00df\u00f1\u00bf\u00ab\u00bb') + ['\t', '\n']
+        ))
+        self.char_to_id = {c: i for i, c in enumerate(self.chars)}
+        self.id_to_char = {i: c for i, c in enumerate(self.chars)}
+        self.vocab_size = len(self.chars)
+
+    def encode(self, text):
+        return [self.char_to_id.get(c, self.char_to_id.get(' ', 0)) for c in text]
+
+    def decode(self, ids):
+        return ''.join(self.id_to_char.get(i, ' ') for i in ids)
+
+    @classmethod
+    def load(cls, path):
+        tok = cls.__new__(cls)
+        with open(path, 'r', encoding='utf-8') as f:
+            data = json.load(f)
+        tok.chars = data['chars']
+        tok.char_to_id = {c: i for i, c in enumerate(tok.chars)}
+        tok.id_to_char = {i: c for i, c in enumerate(tok.chars)}
+        tok.vocab_size = data['vocab_size']
+        return tok
+
+
+# ─── Dataset ─────────────────────────────────────────────────────
+class TokenDataset(Dataset):
+    def __init__(self, tokens, seq_len):
+        self.tokens = tokens
+        self.seq_len = seq_len
+
+    def __len__(self):
+        return max(0, (len(self.tokens) - 1) // self.seq_len)
+
+    def __getitem__(self, idx):
+        start = idx * self.seq_len
+        end = start + self.seq_len + 1
+        chunk = self.tokens[start:end]
+        x = chunk[:-1].clone().detach().long()
+        y = chunk[1:].clone().detach().long()
+        return x, y
+
+
+def load_combined_data(data_dir, tokenizer):
+    all_tokens = []
+    
+    aicl_10x = os.path.join(data_dir, 'aicl_10x.pt')
+    if os.path.exists(aicl_10x):
+        t = torch.load(aicl_10x, map_location='cpu', weights_only=True)
+        t = t.long().clamp(0, tokenizer.vocab_size - 1)
+        all_tokens.append(t)
+        print(f"  AICL 10x: {len(t):,} tokens")
+    
+    for p in [os.path.join(data_dir, '..', 'checkpoints', 'train_ids.pt'),
+              os.path.join(data_dir, 'train_ids.pt'),
+              '/root/CogNet/checkpoints/train_ids.pt']:
+        if os.path.exists(p):
+            t = torch.load(p, map_location='cpu', weights_only=True)
+            t = t.long().clamp(0, tokenizer.vocab_size - 1)
+            all_tokens.append(t)
+            print(f"  Train data: {len(t):,} tokens")
+            break
+    
+    aicl_raw = os.path.join(data_dir, 'aicl_all.pt')
+    if os.path.exists(aicl_raw):
+        t = torch.load(aicl_raw, map_location='cpu', weights_only=True)
+        t = t.long().clamp(0, tokenizer.vocab_size - 1)
+        all_tokens.append(t)
+        print(f"  AICL raw: {len(t):,} tokens")
+    
+    if not all_tokens:
+        raise FileNotFoundError("No data found!")
+    
+    combined = torch.cat(all_tokens, dim=0)
+    print(f"  COMBINED: {len(combined):,} tokens")
+    return combined
+
+
+# ─── CPU-Offloaded Optimizer ────────────────────────────────────
+class CPUAdamW:
+    """AdamW with optimizer states on CPU. Vectorized for speed."""
+    def __init__(self, params, lr=1e-3, betas=(0.9, 0.95), eps=1e-8, weight_decay=0.1):
+        self.params = list(params)
+        self.lr = lr
+        self.beta1, self.beta2 = betas
+        self.eps = eps
+        self.wd = weight_decay
+        self.step_count = 0
+        
+        # Flat buffers for vectorized ops on CPU
+        # Group all params into one flat tensor for fast batch update
+        self.group_sizes = [p.numel() for p in self.params]
+        self.group_offsets = []
+        off = 0
+        for s in self.group_sizes:
+            self.group_offsets.append(off)
+            off += s
+        total = off
+        
+        self.m_flat = torch.zeros(total, dtype=torch.float32, device='cpu')
+        self.v_flat = torch.zeros(total, dtype=torch.float32, device='cpu')
+    
+    def zero_grad(self):
+        for p in self.params:
+            if p.grad is not None:
+                p.grad.zero_()
+    
+    @torch.no_grad()
+    def step(self):
+        self.step_count += 1
+        bc1 = 1 - self.beta1 ** self.step_count
+        bc2 = 1 - self.beta2 ** self.step_count
+        
+        # Gather all grads and params into flat CPU tensors
+        g_parts = []
+        p_parts = []
+        for p in self.params:
+            if p.grad is not None:
+                g_parts.append(p.grad.detach().reshape(-1).float().cpu())
+                p_parts.append(p.data.reshape(-1).float().cpu())
+            else:
+                g_parts.append(torch.zeros(p.numel(), dtype=torch.float32, device='cpu'))
+                p_parts.append(p.data.reshape(-1).float().cpu())
+        
+        g_flat = torch.cat(g_parts)
+        p_flat = torch.cat(p_parts)
+        
+        # Vectorized AdamW update
+        self.m_flat.mul_(self.beta1).add_(g_flat, alpha=1 - self.beta1)
+        self.v_flat.mul_(self.beta2).addcmul_(g_flat, g_flat, value=1 - self.beta2)
+        
+        m_hat = self.m_flat / bc1
+        v_hat = self.v_flat / bc2
+        
+        update = m_hat / (v_hat.sqrt() + self.eps)
+        if self.wd > 0:
+            update.add_(p_flat, alpha=self.wd)
+        
+        p_flat.add_(update, alpha=-self.lr)
+        
+        # Scatter back to GPU params
+        offset = 0
+        for p in self.params:
+            sz = p.numel()
+            p.data.copy_(p_flat[offset:offset+sz].reshape(p.shape).to(p.device, dtype=p.dtype))
+            offset += sz
+
+
+def get_lr(step, warmup, max_steps, lr, min_lr):
+    if step < warmup:
+        return lr * (step + 1) / warmup
+    decay = max(0.0, (max_steps - step) / (max_steps - warmup))
+    return min_lr + (lr - min_lr) * decay
+
+
+# ─── Main ────────────────────────────────────────────────────────
+def main():
+    print("=" * 60)
+    print("  CogNet-1B Training (CPU Offload)")
+    print("=" * 60)
+    
+    device = torch.device('cuda')
+    print(f"  Device: {device}")
+    print(f"  GPU: {torch.cuda.get_device_name(0)}")
+    print(f"  VRAM: {torch.cuda.get_device_properties(0).total_memory / 1e9:.1f} GB")
+    
+    torch.backends.cuda.matmul.allow_tf32 = True
+    torch.backends.cudnn.allow_tf32 = True
+    
+    # Tokenizer
+    tok_path = '/root/CogNet/tokenizer_v3.json'
+    if not os.path.exists(tok_path):
+        tok_path = '/root/CogNet/data_1b/tokenizer_v3.json'
+    tokenizer = CharTokenizer.load(tok_path)
+    print(f"  Tokenizer: {tokenizer.vocab_size} chars")
+    
+    # Data
+    print("\n  Loading data...")
+    combined = load_combined_data('/root/CogNet/data_1b', tokenizer)
+    n = len(combined)
+    split = int(n * 0.95)
+    train_tokens = combined[:split]
+    val_tokens = combined[split:]
+    print(f"  Train: {len(train_tokens):,} | Val: {len(val_tokens):,}")
+    
+    train_ds = TokenDataset(train_tokens, SEQ_LEN)
+    val_ds = TokenDataset(val_tokens, SEQ_LEN)
+    train_loader = DataLoader(train_ds, batch_size=BATCH_SIZE, shuffle=True,
+                              num_workers=2, pin_memory=True, drop_last=True)
+    val_loader = DataLoader(val_ds, batch_size=BATCH_SIZE, shuffle=False,
+                            num_workers=1, pin_memory=True, drop_last=True)
+    
+    # Build 1B model
+    print("\n  Building CogNet-1B...")
+    model = CogNet1B(
+        vocab_size=tokenizer.vocab_size,
+        hidden_dim=HIDDEN_DIM,
+        num_blocks=NUM_BLOCKS,
+        num_channels=NUM_CHANNELS,
+        channel_dim=CHANNEL_DIM,
+        ff_dim=FF_DIM,
+        max_seq_len=MAX_SEQ_LEN,
+        working_slots=WORKING_SLOTS,
+        episodic_slots=EPISODIC_SLOTS,
+        semantic_slots=SEMANTIC_SLOTS,
+    )
+    
+    total_params = sum(p.numel() for p in model.parameters())
+    print(f"  Parameters: {total_params:,} ({total_params/1e9:.2f}B)")
+    
+    model.to(device)
+    print(f"  Model VRAM: {torch.cuda.memory_allocated()/1e9:.2f} GB")
+    
+    # CPU-offloaded optimizer — saves 8GB VRAM!
+    optimizer = CPUAdamW(model.parameters(), lr=LR,
+                          betas=(0.9, 0.95), weight_decay=WEIGHT_DECAY)
+    print("  Optimizer: CPU-offloaded AdamW")
+    
+    # Checkpoint dir
+    ckpt_dir = '/root/CogNet/checkpoints_1b'
+    os.makedirs(ckpt_dir, exist_ok=True)
+    
+    # Resume
+    start_step = 0
+    best_val = float('inf')
+    ckpt_path = os.path.join(ckpt_dir, 'cognet_1b_best.pt')
+    if os.path.exists(ckpt_path):
+        print(f"  Resuming from checkpoint...")
+        ckpt = torch.load(ckpt_path, map_location=device, weights_only=False)
+        model.load_state_dict(ckpt['model_state_dict'])
+        start_step = ckpt.get('step', 0) + 1
+        best_val = ckpt.get('val_loss', float('inf'))
+        # Re-init optimizer with new params
+        optimizer = CPUAdamW(model.parameters(), lr=LR,
+                              betas=(0.9, 0.95), weight_decay=WEIGHT_DECAY)
+        print(f"  Resumed step {start_step}, val_loss={best_val:.4f}")
+    
+    print(f"\n  Effective BS: {BATCH_SIZE * GRAD_ACCUM}")
+    print(f"  Max steps: {MAX_STEPS}")
+    print("=" * 60)
+    
+    model.train()
+    train_iter = iter(train_loader)
+    global_step = start_step
+    t0 = time.time()
+    log_loss = 0.0
+    log_count = 0
+    
+    while global_step < MAX_STEPS:
+        optimizer.zero_grad()
+        accum_loss = 0.0
+        
+        for micro in range(GRAD_ACCUM):
+            try:
+                x, y = next(train_iter)
+            except StopIteration:
+                train_iter = iter(train_loader)
+                x, y = next(train_iter)
+            
+            x = x.to(device, non_blocking=True)
+            y = y.to(device, non_blocking=True)
+            
+            with torch.amp.autocast('cuda', dtype=torch.float16):
+                result = model(x)
+                loss = F.cross_entropy(
+                    result['logits'].view(-1, tokenizer.vocab_size),
+                    y.view(-1)
+                ) / GRAD_ACCUM
+            
+            loss.backward()
+            accum_loss += loss.item()
+        
+        # Clip gradients
+        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
+        
+        # Update LR
+        lr = get_lr(global_step, WARMUP_STEPS, MAX_STEPS, LR, MIN_LR)
+        optimizer.lr = lr
+        
+        # Step (optimizer on CPU)
+        optimizer.step()
+        
+        global_step += 1
+        log_loss += accum_loss
+        log_count += 1
+        
+        if global_step % LOG_EVERY == 0:
+            avg = log_loss / log_count
+            ppl = math.exp(min(avg, 10))
+            elapsed = time.time() - t0
+            spm = LOG_EVERY / (elapsed / 60) if elapsed > 0 else 0
+            vram = torch.cuda.max_memory_allocated() / 1e9
+            
+            print(f"Step {global_step:>6d} | Loss: {avg:.4f} | PPL: {ppl:.2f} | "
+                  f"LR: {lr:.6f} | {spm:.0f} spm | VRAM: {vram:.1f}GB")
+            
+            log_loss = 0.0
+            log_count = 0
+            t0 = time.time()
+        
+        if global_step % EVAL_EVERY == 0:
+            model.eval()
+            val_loss = 0.0
+            vc = 0
+            with torch.no_grad():
+                for vx, vy in val_loader:
+                    vx = vx.to(device, non_blocking=True)
+                    vy = vy.to(device, non_blocking=True)
+                    with torch.amp.autocast('cuda', dtype=torch.float16):
+                        vr = model(vx)
+                        vl = F.cross_entropy(vr['logits'].view(-1, tokenizer.vocab_size), vy.view(-1))
+                    val_loss += vl.item()
+                    vc += 1
+                    if vc >= 50:
+                        break
+            
+            avg_val = val_loss / max(vc, 1)
+            val_ppl = math.exp(min(avg_val, 10))
+            print(f"  [EVAL] step={global_step} val_loss={avg_val:.4f} val_ppl={val_ppl:.2f}")
+            
+            # Sample generation
+            prompts = ["The ", "Once upon a time", "Bonjour ", "def ", "# AICL\nGoal:\n"]
+            for p in prompts:
+                ids = tokenizer.encode(p)
+                inp = torch.tensor([ids], dtype=torch.long, device=device)
+                with torch.no_grad(), torch.amp.autocast('cuda', dtype=torch.float16):
+                    out = model(inp)
+                    probs = F.softmax(out['logits'][0, -1] / 0.7, dim=-1)
+                    top5 = torch.topk(probs, 5)
+                    chars = [tokenizer.decode([t.item()]) for t in top5.indices]
+                print(f'  "{p}" -> {chars}')
+            
+            if avg_val < best_val:
+                best_val = avg_val
+            
+            model.train()
+            t0 = time.time()
+        
+        if global_step % SAVE_EVERY == 0:
+            ckpt = {
+                'model_state_dict': model.state_dict(),
+                'step': global_step,
+                'val_loss': best_val,
+                'hidden_dim': HIDDEN_DIM,
+                'num_blocks': NUM_BLOCKS,
+                'num_channels': NUM_CHANNELS,
+                'channel_dim': CHANNEL_DIM,
+                'ff_dim': FF_DIM,
+                'max_seq_len': MAX_SEQ_LEN,
+                'working_slots': WORKING_SLOTS,
+                'episodic_slots': EPISODIC_SLOTS,
+                'semantic_slots': SEMANTIC_SLOTS,
+                'vocab_size': tokenizer.vocab_size,
+            }
+            torch.save(ckpt, os.path.join(ckpt_dir, 'cognet_1b_best.pt'))
+            
+            # FP16 checkpoint
+            fp16 = {k: v.half() for k, v in model.state_dict().items()}
+            torch.save({
+                'model_state_dict': fp16,
+                'step': global_step, 'val_loss': best_val,
+                'hidden_dim': HIDDEN_DIM, 'num_blocks': NUM_BLOCKS,
+                'num_channels': NUM_CHANNELS, 'channel_dim': CHANNEL_DIM,
+                'ff_dim': FF_DIM, 'max_seq_len': MAX_SEQ_LEN,
+                'working_slots': WORKING_SLOTS, 'episodic_slots': EPISODIC_SLOTS,
+                'semantic_slots': SEMANTIC_SLOTS, 'vocab_size': tokenizer.vocab_size,
+            }, os.path.join(ckpt_dir, 'cognet_1b_fp16.pt'))
+            
+            sz = os.path.getsize(os.path.join(ckpt_dir, 'cognet_1b_best.pt')) / 1e9
+            print(f"  [CKPT] step={global_step} val_loss={best_val:.4f} size={sz:.1f}GB")
+    
+    print(f"\nDONE! Step {global_step}, best val_loss: {best_val:.4f}")
+
+
+if __name__ == "__main__":
+    main()