Upload train_1b_final.py with huggingface_hub

train_1b_final.py

@@ -0,0 +1,604 @@
+#!/usr/bin/env python3
+"""
+CogNet-1B Training Script — RTX 5090 32GB VRAM
+=================================================
+Full AdamW optimizer, BF16 mixed precision, gradient checkpointing.
+Architecture: ~1.02B parameters
+  - hidden_dim=2048, 16 blocks, 8 channels
+  - channel_dim=384, ff_dim=8192
+  - working_slots=128, episodic_slots=256, semantic_slots=512
+"""
+
+import os
+import sys
+import time
+import math
+import json
+import logging
+import argparse
+from pathlib import Path
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.utils.data import Dataset, DataLoader
+
+# ──────────────────────────────────────────────────────────
+#  Logging
+# ──────────────────────────────────────────────────────────
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s [%(levelname)s] %(message)s',
+    handlers=[
+        logging.StreamHandler(sys.stdout),
+        logging.FileHandler('/root/CogNet/train_1b.log', mode='a'),
+    ],
+)
+log = logging.getLogger('cognet1b')
+
+# ──────────────────────────────────────────────────────────
+#  Character Tokenizer
+# ──────────────────────────────────────────────────────────
+class CharTokenizer:
+    """Character-level tokenizer with 136 vocab (ASCII + French accents + special tokens)."""
+
+    def __init__(self, vocab_size=136):
+        self.vocab_size = vocab_size
+        self.pad_token_id = 0
+        self.unk_token_id = 1
+        self.bos_token_id = 2
+        self.eos_token_id = 3
+
+        # Build character set
+        chars = list(range(32, 127))  # ASCII printable
+        # French accented chars
+        french = [192,193,194,195,196,197,199,200,201,202,203,204,205,206,207,
+                  210,211,212,213,214,217,218,219,220,224,225,226,227,228,229,
+                  231,232,233,234,235,236,237,238,239,242,243,244,245,246,249,
+                  250,251,252,253,255]
+        chars.extend(french)
+
+        self.char_to_id = {self.pad_token_id: 0, self.unk_token_id: 1,
+                           self.bos_token_id: 2, self.eos_token_id: 3}
+        for i, c in enumerate(chars[:vocab_size - 4]):
+            self.char_to_id[c] = i + 4
+
+        self.id_to_char = {v: k for k, v in self.char_to_id.items()}
+
+    def encode(self, text):
+        ids = [self.bos_token_id]
+        for ch in text:
+            code = ord(ch)
+            ids.append(self.char_to_id.get(code, self.unk_token_id))
+        ids.append(self.eos_token_id)
+        return ids
+
+    def decode(self, ids):
+        chars = []
+        for i in ids:
+            if i in (self.pad_token_id, self.bos_token_id):
+                continue
+            if i == self.eos_token_id:
+                break
+            code = self.id_to_char.get(i, 0)
+            if code > 0:
+                chars.append(chr(code))
+        return ''.join(chars)
+
+    def save(self, path):
+        with open(path, 'w', encoding='utf-8') as f:
+            json.dump({
+                'vocab_size': self.vocab_size,
+                'char_to_id': {str(k): v for k, v in self.char_to_id.items()},
+            }, f, ensure_ascii=False, indent=2)
+
+    @classmethod
+    def load(cls, path):
+        with open(path, 'r', encoding='utf-8') as f:
+            data = json.load(f)
+        tok = cls.__new__(cls)
+        tok.vocab_size = data['vocab_size']
+        tok.char_to_id = {int(k): v for k, v in data['char_to_id'].items()}
+        tok.id_to_char = {v: k for k, v in tok.char_to_id.items()}
+        tok.pad_token_id = 0
+        tok.unk_token_id = 1
+        tok.bos_token_id = 2
+        tok.eos_token_id = 3
+        return tok
+
+
+# ──────────────────────────────────────────────────────────
+#  Dataset
+# ──────────────────────────────────────────────────────────
+class TokenDataset(Dataset):
+    def __init__(self, data_path, seq_len=512):
+        tokens = torch.load(data_path, map_location='cpu', weights_only=True)
+        if not isinstance(tokens, torch.LongTensor):
+            tokens = tokens.long()
+        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]
+        y = chunk[1:]
+        return x, y
+
+
+# ──────────────────────────────────────────────────────────
+#  CogNet-1B Architecture
+# ──────────────────────────────────────────────────────────
+class CognitiveMemory(nn.Module):
+    """Hierarchical cognitive memory: Working, Episodic, Semantic."""
+
+    def __init__(self, dim, working_slots=128, episodic_slots=256, semantic_slots=512):
+        super().__init__()
+        self.working = nn.Parameter(torch.randn(working_slots, dim) * 0.02)
+        self.episodic = nn.Parameter(torch.randn(episodic_slots, dim) * 0.02)
+        self.semantic = nn.Parameter(torch.randn(semantic_slots, dim) * 0.02)
+        self.gate_w = nn.Linear(dim * 3, 3)
+        self.work_proj = nn.Linear(dim, dim)
+        self.epis_proj = nn.Linear(dim, dim)
+        self.seman_proj = nn.Linear(dim, dim)
+
+    def forward(self, x):
+        B, T, D = x.shape
+        x_pooled = x.mean(dim=1)  # [B, D]
+
+        w_att = torch.matmul(x, self.working.t())  # [B, T, ws]
+        w_att = F.softmax(w_att / math.sqrt(D), dim=-1)
+        w_out = torch.matmul(w_att, self.working)   # [B, T, D]
+        w_out = self.work_proj(w_out)
+
+        e_att = torch.matmul(x, self.episodic.t())
+        e_att = F.softmax(e_att / math.sqrt(D), dim=-1)
+        e_out = torch.matmul(e_att, self.episodic)
+        e_out = self.epis_proj(e_out)
+
+        s_att = torch.matmul(x, self.semantic.t())
+        s_att = F.softmax(s_att / math.sqrt(D), dim=-1)
+        s_out = torch.matmul(s_att, self.semantic)
+        s_out = self.seman_proj(s_out)
+
+        gate_in = torch.cat([w_out.mean(1), e_out.mean(1), s_out.mean(1)], dim=-1)
+        gates = F.softmax(self.gate_w(gate_in), dim=-1)  # [B, 3]
+
+        combined = (gates[:, 0:1, None] * w_out +
+                    gates[:, 1:2, None] * e_out +
+                    gates[:, 2:3, None] * s_out)
+        return x + combined
+
+
+class CogNetBlock(nn.Module):
+    """CogNet block with cognitive routing across channels."""
+
+    def __init__(self, dim, n_channels=8, channel_dim=384, ff_dim=8192,
+                 working_slots=128, episodic_slots=256, semantic_slots=512):
+        super().__init__()
+        self.n_channels = n_channels
+        self.channel_dim = channel_dim
+
+        # Channel projections
+        self.channel_projs = nn.ModuleList([
+            nn.Linear(dim, channel_dim) for _ in range(n_channels)
+        ])
+        self.channel_merge = nn.Linear(n_channels * channel_dim, dim)
+        self.norm1 = nn.LayerNorm(dim)
+
+        # Feed-forward
+        self.ff = nn.Sequential(
+            nn.Linear(dim, ff_dim),
+            nn.GELU(),
+            nn.Linear(ff_dim, dim),
+        )
+        self.norm2 = nn.LayerNorm(dim)
+
+        # Cognitive memory
+        self.cog_mem = CognitiveMemory(dim, working_slots, episodic_slots, semantic_slots)
+        self.norm3 = nn.LayerNorm(dim)
+
+        # Routing gate
+        self.router = nn.Linear(dim, n_channels)
+
+    def forward(self, x):
+        B, T, D = x.shape
+
+        # Channel routing
+        route_logits = self.router(x)  # [B, T, n_ch]
+        route_weights = F.softmax(route_logits, dim=-1)  # [B, T, n_ch]
+
+        # Process each channel
+        channel_outs = []
+        for i, proj in enumerate(self.channel_projs):
+            ch_x = proj(x)  # [B, T, ch_dim]
+            ch_x = F.gelu(ch_x)
+            # Apply route weight
+            w = route_weights[:, :, i:i+1]  # [B, T, 1]
+            channel_outs.append(ch_x * w)
+
+        merged = torch.cat(channel_outs, dim=-1)  # [B, T, n_ch * ch_dim]
+        merged = self.channel_merge(merged)  # [B, T, D]
+
+        # Residual + norm
+        x = self.norm1(x + merged)
+
+        # FFN
+        x = self.norm2(x + self.ff(x))
+
+        # Cognitive memory
+        x = self.norm3(self.cog_mem(x))
+
+        return x
+
+
+class CogNet1B(nn.Module):
+    """CogNet-1B: ~1.06B parameter cognitive language model."""
+
+    def __init__(self, vocab_size=136, hidden_dim=2048, n_blocks=16,
+                 n_channels=8, channel_dim=384, ff_dim=8192,
+                 working_slots=128, episodic_slots=256, semantic_slots=512,
+                 seq_len=512):
+        super().__init__()
+        self.vocab_size = vocab_size
+        self.hidden_dim = hidden_dim
+        self.seq_len = seq_len
+
+        self.token_emb = nn.Embedding(vocab_size, hidden_dim)
+        self.pos_emb = nn.Embedding(seq_len, hidden_dim)
+
+        self.blocks = nn.ModuleList([
+            CogNetBlock(hidden_dim, n_channels, channel_dim, ff_dim,
+                       working_slots, episodic_slots, semantic_slots)
+            for _ in range(n_blocks)
+        ])
+
+        self.final_norm = nn.LayerNorm(hidden_dim)
+        self.head = nn.Linear(hidden_dim, vocab_size, bias=False)
+
+        # Weight tying
+        self.head.weight = self.token_emb.weight
+
+        self._init_weights()
+
+    def _init_weights(self):
+        for module in self.modules():
+            if isinstance(module, nn.Linear):
+                nn.init.normal_(module.weight, mean=0.0, std=0.02)
+                if module.bias is not None:
+                    nn.init.zeros_(module.bias)
+            elif isinstance(module, nn.Embedding):
+                nn.init.normal_(module.weight, mean=0.0, std=0.02)
+
+    def forward(self, input_ids, targets=None):
+        B, T = input_ids.shape
+        assert T <= self.seq_len, f"Sequence length {T} > max {self.seq_len}"
+
+        pos = torch.arange(T, device=input_ids.device).unsqueeze(0)
+        x = self.token_emb(input_ids) + self.pos_emb(pos)
+
+        for block in self.blocks:
+            x = block(x)
+
+        x = self.final_norm(x)
+        logits = self.head(x)
+
+        loss = None
+        if targets is not None:
+            loss = F.cross_entropy(
+                logits.view(-1, self.vocab_size),
+                targets.view(-1),
+                ignore_index=0,
+            )
+
+        return logits, loss
+
+    @torch.no_grad()
+    def generate(self, input_ids, max_new_tokens=200, temperature=0.8, top_k=50):
+        self.eval()
+        for _ in range(max_new_tokens):
+            idx_cond = input_ids[:, -self.seq_len:]
+            logits, _ = self(idx_cond)
+            logits = logits[:, -1, :] / temperature
+            if top_k > 0:
+                v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
+                logits[logits < v[:, [-1]]] = float('-inf')
+            probs = F.softmax(logits, dim=-1)
+            next_id = torch.multinomial(probs, num_samples=1)
+            input_ids = torch.cat([input_ids, next_id], dim=1)
+        return input_ids
+
+
+# ──────────────────────────────────────────────────────────
+#  Learning Rate Schedule
+# ──────────────────────────────────────────────────────────
+def get_cosine_lr(step, warmup_steps, max_steps, max_lr, min_lr):
+    if step < warmup_steps:
+        return max_lr * step / max(1, warmup_steps)
+    if step >= max_steps:
+        return min_lr
+    progress = (step - warmup_steps) / max(1, max_steps - warmup_steps)
+    return min_lr + 0.5 * (max_lr - min_lr) * (1 + math.cos(math.pi * progress))
+
+
+# ──────────────────────────────────────────────────────────
+#  Main Training Loop
+# ──────────────────────────────────────────────────────────
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--batch-size', type=int, default=8, help='Batch size (RTX 5090 32GB BF16)')
+    parser.add_argument('--seq-len', type=int, default=512)
+    parser.add_argument('--max-steps', type=int, default=200000)
+    parser.add_argument('--warmup-steps', type=int, default=4000)
+    parser.add_argument('--max-lr', type=float, default=1e-4)
+    parser.add_argument('--min-lr', type=float, default=1e-5)
+    parser.add_argument('--weight-decay', type=float, default=0.1)
+    parser.add_argument('--grad-clip', type=float, default=1.0)
+    parser.add_argument('--save-every', type=int, default=5000)
+    parser.add_argument('--eval-every', type=int, default=1000)
+    parser.add_argument('--log-every', type=int, default=100)
+    parser.add_argument('--data-path', type=str, default='/root/CogNet/data_1b/aicl_10x.pt')
+    parser.add_argument('--tokenizer-path', type=str, default='/root/CogNet/tokenizer_v3.json')
+    parser.add_argument('--ckpt-dir', type=str, default='/root/CogNet/checkpoints_1b')
+    parser.add_argument('--resume', type=str, default=None, help='Path to checkpoint to resume from')
+    parser.add_argument('--bf16', action='store_true', default=True, help='Use BF16 mixed precision')
+    parser.add_argument('--compile', action='store_true', default=False, help='torch.compile the model')
+    args = parser.parse_args()
+
+    # Setup
+    os.makedirs(args.ckpt_dir, exist_ok=True)
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
+    log.info(f'=== CogNet-1B Training (241GB VRAM Optimized) ===')
+    log.info(f'Device: {device}')
+
+    if torch.cuda.is_available():
+        for i in range(torch.cuda.device_count()):
+            props = torch.cuda.get_device_properties(i)
+            total_gb = props.total_memory / 1e9
+            log.info(f'GPU {i}: {props.name} — {total_gb:.1f} GB VRAM')
+
+    # Tokenizer
+    if os.path.exists(args.tokenizer_path):
+        tokenizer = CharTokenizer.load(args.tokenizer_path)
+        log.info(f'Loaded tokenizer from {args.tokenizer_path} (vocab={tokenizer.vocab_size})')
+    else:
+        tokenizer = CharTokenizer()
+        tokenizer.save(args.tokenizer_path)
+        log.info(f'Created and saved tokenizer (vocab={tokenizer.vocab_size})')
+
+    # Dataset
+    log.info(f'Loading data from {args.data_path}...')
+    dataset = TokenDataset(args.data_path, args.seq_len)
+    log.info(f'Dataset: {len(dataset):,} sequences of length {args.seq_len}')
+
+    dataloader = DataLoader(
+        dataset,
+        batch_size=args.batch_size,
+        shuffle=True,
+        num_workers=4,
+        pin_memory=True,
+        drop_last=True,
+    )
+
+    # Model
+    log.info('Building CogNet-1B model...')
+    model = CogNet1B(
+        vocab_size=tokenizer.vocab_size,
+        hidden_dim=2048,
+        n_blocks=16,
+        n_channels=8,
+        channel_dim=384,
+        ff_dim=8192,
+        working_slots=128,
+        episodic_slots=256,
+        semantic_slots=512,
+        seq_len=args.seq_len,
+    ).to(device)
+
+    # Count parameters
+    total_params = sum(p.numel() for p in model.parameters())
+    trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
+    log.info(f'Total parameters: {total_params:,} ({total_params/1e9:.2f}B)')
+    log.info(f'Trainable parameters: {trainable_params:,}')
+
+    # Optional: torch.compile
+    if args.compile:
+        log.info('Compiling model with torch.compile...')
+        model = torch.compile(model)
+
+    # Optimizer — FULL AdamW, no compromises with 241GB VRAM!
+    optimizer = torch.optim.AdamW(
+        model.parameters(),
+        lr=args.max_lr,
+        betas=(0.9, 0.95),
+        eps=1e-8,
+        weight_decay=args.weight_decay,
+        fused=True,  # Fused CUDA kernel for speed
+    )
+
+    # Mixed precision scaler
+    use_bf16 = args.bf16 and torch.cuda.is_bf16_supported()
+    scaler = None
+    if not use_bf16:
+        scaler = torch.amp.GradScaler('cuda')
+    log.info(f'Mixed precision: {"BF16" if use_bf16 else "FP16 with GradScaler"}')
+
+    # Resume from checkpoint
+    start_step = 0
+    best_loss = float('inf')
+    if args.resume and os.path.exists(args.resume):
+        log.info(f'Resuming from {args.resume}...')
+        ckpt = torch.load(args.resume, map_location=device, weights_only=False)
+        model.load_state_dict(ckpt['model_state_dict'])
+        optimizer.load_state_dict(ckpt['optimizer_state_dict'])
+        start_step = ckpt.get('step', 0)
+        best_loss = ckpt.get('best_loss', float('inf'))
+        log.info(f'Resumed at step {start_step}, best_loss={best_loss:.4f}')
+
+    # Also check for latest checkpoint in ckpt_dir
+    if start_step == 0:
+        latest = os.path.join(args.ckpt_dir, 'cognet_1b_latest.pt')
+        if os.path.exists(latest):
+            log.info(f'Found latest checkpoint: {latest}')
+            ckpt = torch.load(latest, map_location=device, weights_only=False)
+            model.load_state_dict(ckpt['model_state_dict'])
+            optimizer.load_state_dict(ckpt['optimizer_state_dict'])
+            start_step = ckpt.get('step', 0)
+            best_loss = ckpt.get('best_loss', float('inf'))
+            log.info(f'Resumed at step {start_step}, best_loss={best_loss:.4f}')
+
+    # Training loop
+    log.info(f'Starting training from step {start_step} to {args.max_steps}')
+    log.info(f'Batch size: {args.batch_size}, Seq len: {args.seq_len}')
+    log.info(f'LR range: {args.min_lr} -> {args.max_lr} -> {args.min_lr}')
+    log.info(f'Warmup: {args.warmup_steps} steps')
+
+    model.train()
+    data_iter = iter(dataloader)
+    t0 = time.time()
+
+    for step in range(start_step, args.max_steps):
+        # Get batch
+        try:
+            batch = next(data_iter)
+        except StopIteration:
+            data_iter = iter(dataloader)
+            batch = next(data_iter)
+
+        x, y = batch
+        x = x.to(device, non_blocking=True)
+        y = y.to(device, non_blocking=True)
+
+        # Learning rate
+        lr = get_cosine_lr(step, args.warmup_steps, args.max_steps, args.max_lr, args.min_lr)
+        for param_group in optimizer.param_groups:
+            param_group['lr'] = lr
+
+        # Forward + backward
+        optimizer.zero_grad(set_to_none=True)
+
+        if use_bf16:
+            with torch.amp.autocast('cuda', dtype=torch.bfloat16):
+                logits, loss = model(x, y)
+            loss.backward()
+        else:
+            with torch.amp.autocast('cuda', dtype=torch.float16):
+                logits, loss = model(x, y)
+            scaler.scale(loss).backward()
+
+        # Gradient clipping
+        if use_bf16:
+            grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
+        else:
+            scaler.unscale_(optimizer)
+            grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
+
+        # Optimizer step
+        if use_bf16:
+            optimizer.step()
+        else:
+            scaler.step(optimizer)
+            scaler.update()
+
+        # Logging
+        if step % args.log_every == 0:
+            elapsed = time.time() - t0
+            steps_per_sec = args.log_every / max(elapsed, 0.001)
+            tokens_per_sec = steps_per_sec * args.batch_size * args.seq_len
+
+            vram_used = torch.cuda.memory_allocated() / 1e9
+            vram_reserved = torch.cuda.memory_reserved() / 1e9
+
+            log.info(
+                f'Step {step:>7d}/{args.max_steps} | '
+                f'Loss: {loss.item():.4f} | '
+                f'LR: {lr:.2e} | '
+                f'Grad: {grad_norm:.2f} | '
+                f'VRAM: {vram_used:.1f}/{vram_reserved:.1f} GB | '
+                f'Speed: {tokens_per_sec:.0f} tok/s | '
+                f'{steps_per_sec:.1f} steps/s'
+            )
+            t0 = time.time()
+
+        # Evaluation / sample generation
+        if step > 0 and step % args.eval_every == 0:
+            model.eval()
+            with torch.no_grad():
+                prompt = torch.tensor([[tokenizer.bos_token_id]], device=device)
+                sample_ids = model.generate(prompt, max_new_tokens=100, temperature=0.8, top_k=50)
+                sample_text = tokenizer.decode(sample_ids[0].tolist())
+                log.info(f'--- Sample at step {step} ---')
+                log.info(sample_text[:300])
+                log.info(f'--- End sample ---')
+            model.train()
+
+        # Save checkpoint
+        if step > 0 and step % args.save_every == 0:
+            ckpt_path = os.path.join(args.ckpt_dir, f'cognet_1b_step_{step}.pt')
+            latest_path = os.path.join(args.ckpt_dir, 'cognet_1b_latest.pt')
+
+            is_best = loss.item() < best_loss
+            if is_best:
+                best_loss = loss.item()
+                best_path = os.path.join(args.ckpt_dir, 'cognet_1b_best.pt')
+
+            save_dict = {
+                'step': step,
+                'model_state_dict': model.state_dict(),
+                'optimizer_state_dict': optimizer.state_dict(),
+                'loss': loss.item(),
+                'best_loss': best_loss,
+                'config': {
+                    'vocab_size': tokenizer.vocab_size,
+                    'hidden_dim': 2048,
+                    'n_blocks': 16,
+                    'n_channels': 8,
+                    'channel_dim': 384,
+                    'ff_dim': 8192,
+                    'working_slots': 128,
+                    'episodic_slots': 256,
+                    'semantic_slots': 512,
+                    'seq_len': args.seq_len,
+                },
+            }
+
+            torch.save(save_dict, ckpt_path)
+            torch.save(save_dict, latest_path)
+            if is_best:
+                torch.save(save_dict, best_path)
+
+            log.info(f'Saved checkpoint: {ckpt_path} (loss={loss.item():.4f}, best={best_loss:.4f})')
+
+    # Final save
+    final_path = os.path.join(args.ckpt_dir, 'cognet_1b_final.pt')
+    save_dict = {
+        'step': args.max_steps,
+        'model_state_dict': model.state_dict(),
+        'optimizer_state_dict': optimizer.state_dict(),
+        'loss': loss.item(),
+        'best_loss': best_loss,
+        'config': {
+            'vocab_size': tokenizer.vocab_size,
+            'hidden_dim': 2048,
+            'n_blocks': 16,
+            'n_channels': 8,
+            'channel_dim': 384,
+            'ff_dim': 8192,
+            'working_slots': 128,
+            'episodic_slots': 256,
+            'semantic_slots': 512,
+            'seq_len': args.seq_len,
+        },
+    }
+    torch.save(save_dict, final_path)
+    log.info(f'Training complete! Final model saved to {final_path}')
+    log.info(f'Final loss: {loss.item():.4f}, Best loss: {best_loss:.4f}')
+
+
+if __name__ == '__main__':
+    main()