Upload runpod_train_1b.py with huggingface_hub

runpod_train_1b.py

@@ -0,0 +1,1393 @@
+#!/usr/bin/env python3
+"""
+CogNet-1B — RunPod All-in-One Training Script (AICL + HuggingFace Datasets Edition)
+=============================================================
+1. Downloads real HuggingFace datasets (wikitext, codeparrot, fineweb, oscar, c4, alpaca, the-stack)
+2. Downloads pre-tokenized data from CogNet HF repo (if available)
+3. Clones AICL repo from GitHub (github.com/AFKmoney/AICL)
+4. Converts AICL datasets + examples + source + spec to training tokens
+5. Tokenizes HF scripts + config files
+6. Generates supplementary synthetic data (code + English + French)
+7. Merges ALL data into final training dataset
+8. Builds the 1B model (~1.06B params)
+9. Trains with BF16 mixed precision on RTX PRO 6000
+10. Auto-saves checkpoints to network volume
+11. Auto-pushes to HuggingFace when training completes
+12. Graceful shutdown on SIGTERM/SIGINT (saves checkpoint)
+13. Auto-resume from latest checkpoint if found
+
+Data mix: ~60% HuggingFace datasets + 15% AICL (10x repeated) + 5% scripts + 20% synthetic
+
+Usage:
+    python3 runpod_train_1b.py [--max-steps 200000] [--batch-size 8] [--skip-data-prep]
+
+CHECKPOINTS ARE SAVED TO /workspace/CogNet/checkpoints_1b/
+This should be on a RunPod Network Volume so data survives pod restarts!
+"""
+
+import os
+import sys
+import time
+import math
+import json
+import random
+import string
+import logging
+import argparse
+import signal
+import subprocess
+from pathlib import Path
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.utils.data import Dataset, DataLoader
+
+# ──────────────────────────────────────────────────────────
+#  Paths — /workspace is the RunPod Network Volume
+# ──────────────────────────────────────────────────────────
+WORKSPACE = os.environ.get('COGNET_WORKSPACE', '/workspace/CogNet')
+DATA_DIR = os.path.join(WORKSPACE, 'data_1b')
+CKPT_DIR = os.path.join(WORKSPACE, 'checkpoints_1b')
+LOG_FILE = os.path.join(WORKSPACE, 'train_1b.log')
+TOKENIZER_PATH = os.path.join(WORKSPACE, 'tokenizer_v3.json')
+HF_REPO = 'thefinalboss/CogNet-1B'
+HF_TOKEN = os.environ.get('HF_TOKEN', '')
+
+# ──────────────────────────────────────────────────────────
+#  Logging
+# ──────────────────────────────────────────────────────────
+os.makedirs(WORKSPACE, exist_ok=True)
+os.makedirs(CKPT_DIR, exist_ok=True)
+os.makedirs(DATA_DIR, exist_ok=True)
+
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s [%(levelname)s] %(message)s',
+    handlers=[
+        logging.StreamHandler(sys.stdout),
+        logging.FileHandler(LOG_FILE, mode='a'),
+    ],
+)
+log = logging.getLogger('cognet1b')
+
+# ──────────────────────────────────────────────────────────
+#  Graceful shutdown — save checkpoint on SIGTERM/SIGINT
+# ──────────────────────────────────────────────────────────
+shutdown_requested = False
+def handle_signal(signum, frame):
+    global shutdown_requested
+    log.warning(f'Received signal {signum}, will save checkpoint and exit after current step...')
+    shutdown_requested = True
+
+signal.signal(signal.SIGTERM, handle_signal)
+signal.signal(signal.SIGINT, handle_signal)
+
+
+# ═══════════════════════════════════════════════════════════
+#  AICL Repo URL
+# ═══════════════════════════════════════════════════════════
+AICL_REPO_URL = 'https://github.com/AFKmoney/AICL.git'
+AICL_LOCAL = os.path.join(WORKSPACE, 'aicl_repo')
+AICL_REPEAT = int(os.environ.get('AICL_REPEAT', '10'))
+
+
+# ═══════════════════════════════════════════════════════════
+#  Step 1: Clone AICL repo from GitHub
+# ═══════════════════════════════════════════════════════════
+def clone_aicl_repo():
+    """Clone the AICL GitHub repository for dataset conversion."""
+    if os.path.isdir(os.path.join(AICL_LOCAL, '.git')):
+        log.info(f'AICL repo already exists at {AICL_LOCAL}')
+        return
+    log.info('Cloning AICL repo from GitHub...')
+    subprocess.run(['git', 'clone', AICL_REPO_URL, AICL_LOCAL], check=True)
+    log.info(f'AICL repo cloned to {AICL_LOCAL}')
+
+
+# ═══════════════════════════════════════════════════════════
+#  Step 2: Extract text from AICL repo
+# ═══════════════════════════════════════════════════════════
+def extract_aicl_jsonl(repo_path):
+    """Extract text from JSONL dataset files."""
+    texts = []
+    datasets_dir = os.path.join(repo_path, 'datasets')
+    if not os.path.isdir(datasets_dir):
+        log.warning('No datasets/ directory in AICL repo')
+        return texts
+
+    import glob as glob_mod
+    jsonl_files = sorted(glob_mod.glob(os.path.join(datasets_dir, '*.jsonl')))
+    log.info(f'  Found {len(jsonl_files)} JSONL files')
+
+    total = 0
+    for jf in jsonl_files:
+        count = 0
+        with open(jf, 'r', encoding='utf-8') as f:
+            for line in f:
+                line = line.strip()
+                if not line:
+                    continue
+                try:
+                    entry = json.loads(line)
+                except json.JSONDecodeError:
+                    continue
+                if 'code' in entry:
+                    texts.append(entry['code'])
+                    count += 1
+                elif 'completion' in entry:
+                    instr = entry.get('instruction', '')
+                    if instr:
+                        texts.append(f"# Instruction:\n{instr}\n\n# Completion:\n{entry['completion']}")
+                    else:
+                        texts.append(entry['completion'])
+                    count += 1
+                if 'snippets' in entry:
+                    for snip in entry['snippets']:
+                        if isinstance(snip, dict) and 'completion' in snip:
+                            texts.append(snip['completion'])
+                            count += 1
+                        elif isinstance(snip, str):
+                            texts.append(snip)
+                            count += 1
+        total += count
+        log.info(f'    {os.path.basename(jf)}: {count} entries')
+    log.info(f'  JSONL total: {total} entries')
+    return texts
+
+
+def extract_aicl_examples(repo_path):
+    """Extract .aicl example files."""
+    import glob as glob_mod
+    texts = []
+    examples_dir = os.path.join(repo_path, 'examples')
+    if not os.path.isdir(examples_dir):
+        return texts
+    aicl_files = sorted(glob_mod.glob(os.path.join(examples_dir, '**/*.aicl'), recursive=True))
+    for af in aicl_files:
+        try:
+            with open(af, 'r', encoding='utf-8') as f:
+                content = f.read()
+            if content.strip():
+                texts.append(f"# === AICL Example: {os.path.basename(af)} ===\n{content}")
+        except Exception:
+            pass
+    log.info(f'  .aicl examples: {len(texts)} files')
+    return texts
+
+
+def extract_aicl_source(repo_path):
+    """Extract source code from src/, tools/, scripts/."""
+    import glob as glob_mod
+    texts = []
+    code_dirs = ['src', 'tools', 'scripts']
+    code_exts = {'.py', '.ts', '.tsx', '.js', '.jsx', '.mjs', '.json', '.prisma'}
+    for cdir in code_dirs:
+        full_dir = os.path.join(repo_path, cdir)
+        if not os.path.isdir(full_dir):
+            continue
+        count = 0
+        for ext in code_exts:
+            for cf in sorted(glob_mod.glob(os.path.join(full_dir, f'**/*{ext}'), recursive=True)):
+                if 'node_modules' in cf or '.next' in cf or '__pycache__' in cf:
+                    continue
+                try:
+                    with open(cf, 'r', encoding='utf-8') as f:
+                        content = f.read()
+                    if len(content.strip()) > 50:
+                        texts.append(f"# === Source: {os.path.relpath(cf, repo_path)} ===\n{content}")
+                        count += 1
+                except Exception:
+                    pass
+        log.info(f'  {cdir}/: {count} files')
+    return texts
+
+
+def extract_aicl_spec_docs(repo_path):
+    """Extract spec, docs, README, tests."""
+    import glob as glob_mod
+    texts = []
+    # Spec
+    for f in sorted(glob_mod.glob(os.path.join(repo_path, 'spec', '*'))):
+        try:
+            with open(f, 'r', encoding='utf-8') as fh:
+                content = fh.read()
+            if content.strip():
+                texts.append(f"# === AICL Spec: {os.path.relpath(f, repo_path)} ===\n{content}")
+        except Exception:
+            pass
+    # Docs
+    for f in sorted(glob_mod.glob(os.path.join(repo_path, 'docs', '*'))):
+        try:
+            with open(f, 'r', encoding='utf-8') as fh:
+                content = fh.read()
+            if content.strip():
+                texts.append(f"# === AICL Docs: {os.path.relpath(f, repo_path)} ===\n{content}")
+        except Exception:
+            pass
+    # README
+    readme = os.path.join(repo_path, 'README.md')
+    if os.path.isfile(readme):
+        try:
+            with open(readme, 'r', encoding='utf-8') as f:
+                texts.append(f"# === AICL README ===\n{f.read()}")
+        except Exception:
+            pass
+    # Tests
+    for f in sorted(glob_mod.glob(os.path.join(repo_path, 'tests', '*.py'))):
+        try:
+            with open(f, 'r', encoding='utf-8') as fh:
+                content = fh.read()
+            if len(content.strip()) > 100:
+                texts.append(f"# === AICL Tests: {os.path.relpath(f, repo_path)} ===\n{content}")
+        except Exception:
+            pass
+    log.info(f'  Spec/docs/tests: {len(texts)} files')
+    return texts
+
+
+# ═══════════════════════════════════════════════════════════
+#  Step 3: Data Preparation — Full Pipeline
+# ═══════════════════════════════════════════════════════════
+def prepare_data(skip=False):
+    """
+    Full data preparation pipeline:
+    A. Download real HuggingFace datasets (wikitext, code, web, French)
+    B. Download pre-tokenized data from CogNet HF repo (if available)
+    C. Clone AICL repo from GitHub → convert to tokens
+    D. Load scripts from HuggingFace and convert to tokens
+    E. Generate synthetic data (code + English + French)
+    F. Merge everything + shuffle → train_merged.pt
+    """
+    if skip:
+        log.info('Skipping data preparation (--skip-data-prep)')
+        return
+
+    # Check if merged data already exists
+    train_path = os.path.join(DATA_DIR, 'train_merged.pt')
+    if os.path.exists(train_path):
+        size_mb = os.path.getsize(train_path) / 1e6
+        log.info(f'Training data already exists: {train_path} ({size_mb:.0f} MB)')
+        return
+
+    # Log existing partial files (but DON'T skip — we may need to add HF datasets)
+    pt_files = list(Path(DATA_DIR).glob('*.pt'))
+    if pt_files:
+        total_mb = sum(f.stat().st_size for f in pt_files) / 1e6
+        log.info(f'Found {len(pt_files)} existing tokenized files ({total_mb:.0f} MB) in {DATA_DIR}')
+        for f in pt_files:
+            log.info(f'  {f.name} ({f.stat().st_size / 1e6:.0f} MB)')
+
+    log.info('='*60)
+    log.info('DATA PREPARATION — Full Pipeline')
+    log.info('A. HuggingFace datasets (wikitext, code, web, French)')
+    log.info('B. CogNet HF repo data (pre-tokenized)')
+    log.info('C. AICL repo conversion (GitHub → tokens)')
+    log.info('D. HuggingFace scripts → tokens')
+    log.info('E. Synthetic data generation')
+    log.info('F. Merge + shuffle')
+    log.info('='*60)
+
+    tokenizer = CharTokenizer()
+    all_tensors = []  # List of tensors to merge
+
+    # ──────────────────────────────────────────────────────
+    #  Part A: Download REAL HuggingFace datasets
+    # ──────────────────────────────────────────────────────
+    log.info('\n--- Part A: Downloading HuggingFace datasets ---')
+    subprocess.run([sys.executable, '-m', 'pip', 'install', 'datasets', '-q'], check=False)
+
+    hf_datasets_path = os.path.join(DATA_DIR, 'hf_datasets_tokens.pt')
+    if os.path.exists(hf_datasets_path):
+        log.info(f'  HF datasets tokens already exist: {hf_datasets_path}')
+        t = torch.load(hf_datasets_path, map_location='cpu', weights_only=True).long()
+        all_tensors.append(t)
+        log.info(f'  Loaded: {len(t):,} tokens')
+    else:
+        try:
+            from datasets import load_dataset
+
+            hf_ids = []  # All token IDs from HF datasets
+            total_chars = 0
+
+            # --- A1: Wikitext-103-raw (English Wikipedia, ~500MB text) ---
+            log.info('  A1: Loading wikitext-103-raw-v1...')
+            try:
+                wt = load_dataset('wikitext', 'wikitext-103-raw-v1', split='train')
+                wt_texts = [x['text'] for x in wt if x['text'] and len(x['text'].strip()) > 20]
+                wt_chars = sum(len(t) for t in wt_texts)
+                log.info(f'  wikitext: {len(wt_texts):,} docs, {wt_chars:,} chars ({wt_chars/1e6:.0f} MB)')
+                for text in wt_texts:
+                    hf_ids.extend(tokenizer.encode(text))
+                total_chars += wt_chars
+            except Exception as e:
+                log.warning(f'  wikitext failed: {e}')
+
+            # --- A2: CodeParrot-clean (Python code) ---
+            log.info('  A2: Loading codeparrot/codeparrot-clean...')
+            try:
+                cp = load_dataset('codeparrot/codeparrot-clean', split='train', streaming=True)
+                cp_chars = 0
+                cp_docs = 0
+                for example in cp:
+                    code = example.get('content', '') or example.get('text', '')
+                    if len(code.strip()) > 100:
+                        hf_ids.extend(tokenizer.encode(code))
+                        cp_chars += len(code)
+                        cp_docs += 1
+                    if cp_chars > 300_000_000:  # 300M chars max
+                        break
+                log.info(f'  codeparrot: {cp_docs:,} files, {cp_chars:,} chars ({cp_chars/1e6:.0f} MB)')
+                total_chars += cp_chars
+            except Exception as e:
+                log.warning(f'  codeparrot failed: {e}')
+
+            # --- A3: FineWeb (high quality web text, sample 5M chars) ---
+            log.info('  A3: Loading HuggingFaceFW/fineweb (sample)...')
+            try:
+                fw = load_dataset('HuggingFaceFW/fineweb', split='train', streaming=True)
+                fw_chars = 0
+                fw_docs = 0
+                for example in fw:
+                    text = example.get('text', '')
+                    if len(text.strip()) > 50:
+                        hf_ids.extend(tokenizer.encode(text))
+                        fw_chars += len(text)
+                        fw_docs += 1
+                    if fw_chars > 500_000_000:  # 500M chars max
+                        break
+                log.info(f'  fineweb: {fw_docs:,} docs, {fw_chars:,} chars ({fw_chars/1e6:.0f} MB)')
+                total_chars += fw_chars
+            except Exception as e:
+                log.warning(f'  fineweb failed: {e}')
+
+            # --- A4: OSCAR French subset (French web text) ---
+            log.info('  A4: Loading oscar (French subset)...')
+            try:
+                oscar_fr = load_dataset('oscar', 'unshuffled_deduplicated_fr', split='train', streaming=True, trust_remote_code=True)
+                fr_chars = 0
+                fr_docs = 0
+                for example in oscar_fr:
+                    text = example.get('text', '')
+                    if len(text.strip()) > 50:
+                        hf_ids.extend(tokenizer.encode(text))
+                        fr_chars += len(text)
+                        fr_docs += 1
+                    if fr_chars > 100_000_000:  # 100M chars max
+                        break
+                log.info(f'  oscar-fr: {fr_docs:,} docs, {fr_chars:,} chars ({fr_chars/1e6:.0f} MB)')
+                total_chars += fr_chars
+            except Exception as e:
+                log.warning(f'  oscar-fr failed: {e}')
+
+            # --- A5: The Stack Smol (multi-language code) ---
+            log.info('  A5: Loading bigcode/the-stack-smol...')
+            try:
+                stack = load_dataset('bigcode/the-stack-smol', split='train', streaming=True, trust_remote_code=True)
+                stack_chars = 0
+                stack_docs = 0
+                for example in stack:
+                    code = example.get('content', '') or example.get('text', '')
+                    if len(code.strip()) > 100:
+                        hf_ids.extend(tokenizer.encode(code))
+                        stack_chars += len(code)
+                        stack_docs += 1
+                    if stack_chars > 200_000_000:  # 200M chars max
+                        break
+                log.info(f'  the-stack-smol: {stack_docs:,} files, {stack_chars:,} chars ({stack_chars/1e6:.0f} MB)')
+                total_chars += stack_chars
+            except Exception as e:
+                log.warning(f'  the-stack-smol failed: {e}')
+
+            # --- A6: Alpaca-cleaned (instruction following) ---
+            log.info('  A6: Loading yahma/alpaca-cleaned...')
+            try:
+                alpaca = load_dataset('yahma/alpaca-cleaned', split='train')
+                alp_texts = []
+                for x in alpaca:
+                    instr = x.get('instruction', '')
+                    inp = x.get('input', '')
+                    out = x.get('output', '')
+                    text = f"### Instruction:\n{instr}\n"
+                    if inp:
+                        text += f"### Input:\n{inp}\n"
+                    text += f"### Response:\n{out}\n"
+                    alp_texts.append(text)
+                alp_chars = sum(len(t) for t in alp_texts)
+                log.info(f'  alpaca: {len(alp_texts):,} instructions, {alp_chars:,} chars')
+                for text in alp_texts:
+                    hf_ids.extend(tokenizer.encode(text))
+                total_chars += alp_chars
+            except Exception as e:
+                log.warning(f'  alpaca failed: {e}')
+
+            # --- A7: C4 (Colossal Clean Crawled Corpus, English sample) ---
+            log.info('  A7: Loading c4 (en, sample)...')
+            try:
+                c4 = load_dataset('c4', 'en', split='train', streaming=True)
+                c4_chars = 0
+                c4_docs = 0
+                for example in c4:
+                    text = example.get('text', '')
+                    if len(text.strip()) > 100:
+                        hf_ids.extend(tokenizer.encode(text))
+                        c4_chars += len(text)
+                        c4_docs += 1
+                    if c4_chars > 300_000_000:  # 300M chars max
+                        break
+                log.info(f'  c4-en: {c4_docs:,} docs, {c4_chars:,} chars ({c4_chars/1e6:.0f} MB)')
+                total_chars += c4_chars
+            except Exception as e:
+                log.warning(f'  c4 failed: {e}')
+
+            # Save all HF dataset tokens
+            log.info(f'  Total HF datasets: {len(hf_ids):,} tokens, {total_chars:,} chars')
+            if hf_ids:
+                hf_tensor = torch.tensor(hf_ids, dtype=torch.long)
+                torch.save(hf_tensor, hf_datasets_path)
+                log.info(f'  Saved HF datasets tokens: {hf_datasets_path} ({os.path.getsize(hf_datasets_path)/1e6:.0f} MB)')
+                all_tensors.append(hf_tensor)
+                del hf_ids, hf_tensor
+            else:
+                log.warning('  No HF dataset tokens collected!')
+
+        except ImportError:
+            log.warning('  datasets library not available, skipping HF datasets')
+        except Exception as e:
+            log.warning(f'  HF datasets loading failed: {e}')
+
+    # ──────────────────────────────────────────────────────
+    #  Part B: Download pre-tokenized data from CogNet HF repo
+    # ──────────────────────────────────────────────────────
+    log.info('\n--- Part B: Downloading CogNet HF repo data ---')
+    try:
+        from huggingface_hub import hf_hub_download, list_repo_files
+
+        # Install huggingface_hub if needed
+        subprocess.run([sys.executable, '-m', 'pip', 'install', 'huggingface_hub', '-q'], check=False)
+
+        # List data files in HF repo
+        repo_files = list_repo_files(HF_REPO, token=HF_TOKEN)
+        data_files = [f for f in repo_files if f.startswith('data/') and f.endswith('.pt')]
+        script_files = [f for f in repo_files if f.endswith('.py') or f.endswith('.json')]
+        log.info(f'HF repo: {len(data_files)} data files, {len(script_files)} script files')
+
+        # Download any pre-tokenized .pt data files
+        for df in data_files:
+            fname = os.path.basename(df)
+            local_path = os.path.join(DATA_DIR, fname)
+            if os.path.exists(local_path):
+                log.info(f'  Already have {fname}, skipping')
+                continue
+            log.info(f'  Downloading {df}...')
+            try:
+                hf_hub_download(HF_REPO, df, local_dir=DATA_DIR, token=HF_TOKEN)
+                size_mb = os.path.getsize(local_path) / 1e6
+                log.info(f'  Downloaded {fname} ({size_mb:.0f} MB)')
+            except Exception as e:
+                log.warning(f'  Failed to download {df}: {e}')
+
+        # Download scripts for tokenization
+        hf_scripts_dir = os.path.join(WORKSPACE, 'hf_scripts')
+        os.makedirs(hf_scripts_dir, exist_ok=True)
+        for sf in script_files:
+            if sf.startswith('data/'):
+                continue  # Skip data dir scripts
+            try:
+                hf_hub_download(HF_REPO, sf, local_dir=hf_scripts_dir, token=HF_TOKEN)
+            except Exception:
+                pass
+        log.info(f'  Downloaded scripts to {hf_scripts_dir}')
+
+    except Exception as e:
+        log.warning(f'HF download failed (non-fatal): {e}')
+        log.info('Will proceed without HF pre-tokenized data')
+
+    # Collect any downloaded .pt files (excluding hf_datasets_tokens.pt already loaded)
+    for pt_file in sorted(Path(DATA_DIR).glob('*.pt')):
+        if pt_file.name == 'hf_datasets_tokens.pt':
+            continue  # Already loaded in Part A
+        if pt_file.name in ['aicl_tokens.pt', 'hf_scripts_tokens.pt', 'synthetic_tokens.pt', 'train_merged.pt']:
+            continue  # Will be loaded later
+        log.info(f'  Loading {pt_file.name}...')
+        t = torch.load(str(pt_file), map_location='cpu', weights_only=True).long()
+        all_tensors.append(t)
+        log.info(f'  {pt_file.name}: {len(t):,} tokens')
+
+    # ──────────────────────────────────────────────────────
+    #  Part C: Clone AICL repo from GitHub → convert to tokens
+    # ──────────────────────────────────────────────────────
+    aicl_path = os.path.join(DATA_DIR, 'aicl_tokens.pt')
+    if os.path.exists(aicl_path):
+        log.info(f'\n--- Part C: AICL tokens already exist, loading {aicl_path} ---')
+        aicl_tensor = torch.load(aicl_path, map_location='cpu', weights_only=True).long()
+        log.info(f'  AICL: {len(aicl_tensor):,} tokens ({os.path.getsize(aicl_path)/1e6:.0f} MB)')
+        all_tensors.append(aicl_tensor)
+    else:
+        log.info('\n--- Part C: AICL Repo Conversion ---')
+        clone_aicl_repo()
+
+        log.info('Extracting text from AICL repo...')
+        aicl_texts = []
+        aicl_texts.extend(extract_aicl_jsonl(AICL_LOCAL))
+        aicl_texts.extend(extract_aicl_examples(AICL_LOCAL))
+        aicl_texts.extend(extract_aicl_source(AICL_LOCAL))
+        aicl_texts.extend(extract_aicl_spec_docs(AICL_LOCAL))
+
+        aicl_chars = sum(len(t) for t in aicl_texts)
+        log.info(f'AICL extracted: {len(aicl_texts):,} chunks, {aicl_chars:,} chars ({aicl_chars/1e6:.1f} MB)')
+
+        # Repeat AICL data to increase weight in training
+        aicl_texts_repeated = aicl_texts * AICL_REPEAT
+        aicl_chars_rep = sum(len(t) for t in aicl_texts_repeated)
+        log.info(f'AICL after {AICL_REPEAT}x repeat: {len(aicl_texts_repeated):,} chunks, {aicl_chars_rep:,} chars ({aicl_chars_rep/1e6:.1f} MB)')
+
+        # Tokenize AICL
+        log.info('Tokenizing AICL data...')
+        aicl_ids = []
+        for i, text in enumerate(aicl_texts_repeated):
+            aicl_ids.extend(tokenizer.encode(text))
+            if (i + 1) % 5000 == 0:
+                log.info(f'  AICL tokenized {i+1}/{len(aicl_texts_repeated)} chunks ({len(aicl_ids):,} tokens)')
+        log.info(f'AICL tokens: {len(aicl_ids):,}')
+
+        aicl_tensor = torch.tensor(aicl_ids, dtype=torch.long)
+        torch.save(aicl_tensor, aicl_path)
+        log.info(f'AICL data saved: {aicl_path} ({os.path.getsize(aicl_path)/1e6:.0f} MB)')
+        all_tensors.append(aicl_tensor)
+        del aicl_texts, aicl_texts_repeated, aicl_ids  # Free memory
+
+    # ──────────────────────────────────────────────────────
+    #  Part D: Tokenize HF scripts + config files
+    # ──────────────────────────────────────────────────────
+    log.info('\n--- Part D: HuggingFace Scripts → Tokens ---')
+    script_texts = []
+    hf_scripts_dir = os.path.join(WORKSPACE, 'hf_scripts')
+    if os.path.isdir(hf_scripts_dir):
+        import glob as glob_mod
+        for ext in ['.py', '.json', '.md']:
+            for sf in sorted(glob_mod.glob(os.path.join(hf_scripts_dir, f'**/*{ext}'), recursive=True)):
+                if 'node_modules' in sf or '.git' in sf:
+                    continue
+                try:
+                    with open(sf, 'r', encoding='utf-8') as f:
+                        content = f.read()
+                    if len(content.strip()) > 50:
+                        rel = os.path.relpath(sf, hf_scripts_dir)
+                        script_texts.append(f"# === HF Script: {rel} ===\n{content}")
+                except Exception:
+                    pass
+
+    # Also tokenize CogNet training scripts themselves (the model learns its own code!)
+    for script_name in ['runpod_train_1b.py', 'cognet_1b.py', 'train_1b_final.py',
+                        'gen_data_1b.py', 'cognet_data_prep.py', 'cognet_hf_push.py',
+                        'build_aicl_dataset.py']:
+        sp = os.path.join(hf_scripts_dir, script_name)
+        if os.path.exists(sp):
+            try:
+                with open(sp, 'r', encoding='utf-8') as f:
+                    content = f.read()
+                if len(content.strip()) > 100:
+                    script_texts.append(f"# === CogNet Script: {script_name} ===\n{content}")
+            except Exception:
+                pass
+
+    if script_texts:
+        log.info(f'Found {len(script_texts)} HF scripts to tokenize')
+        script_ids = []
+        for text in script_texts:
+            script_ids.extend(tokenizer.encode(text))
+        script_tensor = torch.tensor(script_ids, dtype=torch.long)
+        script_path = os.path.join(DATA_DIR, 'hf_scripts_tokens.pt')
+        torch.save(script_tensor, script_path)
+        log.info(f'HF scripts: {len(script_ids):,} tokens saved to {script_path}')
+        # Repeat 3x to give scripts some weight
+        all_tensors.append(script_tensor.repeat(3))
+        del script_texts, script_ids, script_tensor
+    else:
+        log.info('No HF scripts found to tokenize')
+
+    # ──────────────────────────────────────────────────────
+    #  Part E: Synthetic data (code + English + French)
+    # ──────────────────────────────────────────────────────
+    syn_path = os.path.join(DATA_DIR, 'synthetic_tokens.pt')
+    if os.path.exists(syn_path):
+        log.info(f'\n--- Part E: Synthetic tokens already exist, loading {syn_path} ---')
+        syn_tensor = torch.load(syn_path, map_location='cpu', weights_only=True).long()
+        log.info(f'  Synthetic: {len(syn_tensor):,} tokens ({os.path.getsize(syn_path)/1e6:.0f} MB)')
+        all_tensors.append(syn_tensor)
+    else:
+        log.info('\n--- Part E: Synthetic Data Generation ---')
+        target_synthetic_chars = 50_000_000  # 50M chars (reduced — real HF datasets provide most data now)
+
+        func_names = ['process','compute','transform','validate','parse','encode','decode',
+                      'train','predict','analyze','optimize','initialize','configure','execute',
+                      'render','compile','evaluate','generate','filter','sort','merge','split']
+        cls_names = ['Model','Processor','Handler','Manager','Engine','Pipeline','Service',
+                     'Client','Server','Agent','Router','Builder','Factory','Observer']
+        params = ['x','y','data','input','value','config','params','options','state','context']
+
+        py_templates = [
+            "def {f}({p1}, {p2}):\n    \"\"\"Process {p1} with {p2}.\"\"\"\n    result = {p1} + {p2}\n    return result\n\n",
+            "class {cls}:\n    def __init__(self, {p1}):\n        self.{p1} = {p1}\n\n    def process(self, {p2}):\n        return self.{p1} * {p2}\n\n",
+            "for {var} in range({n}):\n    {p1}({var})\n\n",
+            "try:\n    {var} = {func}({p1})\nexcept ValueError as e:\n    {var} = None\n\n",
+            "{var} = [{expr} for {var2} in {iter} if {cond}]\n",
+            "async def {f}({p1}):\n    result = await process({p1})\n    return result\n\n",
+            "import {mod}\nfrom {mod2} import {imp}\n\n{var} = {expr}\nif {cond}:\n    {stmt}\n",
+        ]
+
+        en_sentences = [
+            "The quick brown fox jumps over the lazy dog. ",
+            "Science tells us that the universe is vast and beautiful. ",
+            "CogNet is a non-transformer language model with cognitive routing and memory. ",
+            "The architecture processes sequences in linear time making it efficient for long contexts. ",
+            "Knowledge is power and understanding is the key to wisdom. ",
+            "The future of artificial intelligence is bright and full of possibilities. ",
+            "Mathematics is the language in which the universe speaks to us. ",
+            "Innovation distinguishes between a leader and a follower. ",
+            "AICL is a structured programming language for building complex software systems. ",
+            "Cognitive memory enables the model to store and retrieve information across different time scales. ",
+            "The working memory holds recent information, episodic memory stores experiences, and semantic memory encodes general knowledge. ",
+            "Channel routing allows different aspects of the input to be processed by specialized pathways. ",
+        ]
+
+        fr_sentences = [
+            "Bonjour le monde est beau et la science est merveilleuse. ",
+            "Le futur de l'intelligence artificielle est prometteur et rempli de defis passionnants. ",
+            "CogNet est un modele de langage non-transformateur avec routage cognitif et memoire. ",
+            "La connaissance est le pouvoir et la comprehension est la cle de la sagesse. ",
+            "AICL est un langage de programmation structure pour construire des systemes logiciels complexes. ",
+            "L'apprentissage automatique transforme notre facon de resoudre les problemes. ",
+        ]
+
+        syn_ids = []
+        chars_gen = 0
+        rng = random.Random(42)
+
+        while chars_gen < target_synthetic_chars:
+            texts = []
+            for _ in range(400):
+                t = rng.choice(py_templates)
+                try:
+                    text = t.format(
+                        f=rng.choice(func_names), cls=rng.choice(cls_names),
+                        p1=rng.choice(params), p2=rng.choice(params),
+                        var=rng.choice(params[:5]), var2=rng.choice(params[:5]),
+                        n=rng.randint(1, 100), func=rng.choice(func_names),
+                        mod=rng.choice(['torch','numpy','json','os','sys','typing']),
+                        mod2=rng.choice(['torch.nn','numpy','json','pathlib','collections']),
+                        imp=rng.choice(['Model','Dataset','DataLoader','Config','Handler']),
+                        expr=rng.choice(['x**2','x+1','x*2','len(x)','sum(data)']),
+                        iter=rng.choice(['range(10)','data','items','values','batch']),
+                        cond=rng.choice(['x > 0','x is not None','x % 2 == 0','len(x) > 0']),
+                        stmt=rng.choice(['result = x * 2','continue','break','process(x)']),
+                    )
+                    texts.append(text)
+                except (KeyError, IndexError):
+                    texts.append(f"x = {rng.randint(1,100)}\nresult = x * 2\n\n")
+            for _ in range(800):
+                texts.append(rng.choice(en_sentences))
+            for _ in range(400):
+                texts.append(rng.choice(fr_sentences))
+
+            batch_text = ''.join(texts)
+            batch_ids = tokenizer.encode(batch_text)
+            syn_ids.extend(batch_ids)
+            chars_gen += len(batch_text)
+
+            if len(syn_ids) % 50_000_000 < 10_000_000:
+                pct = chars_gen / target_synthetic_chars * 100
+                log.info(f'  Synthetic: {chars_gen/1e6:.0f}M chars ({pct:.0f}%)')
+
+        syn_tensor = torch.tensor(syn_ids, dtype=torch.long)
+        torch.save(syn_tensor, syn_path)
+        log.info(f'Synthetic data saved: {syn_path} ({os.path.getsize(syn_path)/1e6:.0f} MB, {len(syn_ids):,} tokens)')
+        all_tensors.append(syn_tensor)
+        del syn_ids  # Free memory
+
+    # ──────────────────────────────────────────────────────
+    #  Part F: Merge ALL data + shuffle
+    # ──────────────────────────────────────────────────────
+    log.info('\n--- Part F: Merging ALL datasets ---')
+    log.info(f'Datasets to merge: {len(all_tensors)}')
+    for i, t in enumerate(all_tensors):
+        log.info(f'  [{i}] {len(t):,} tokens')
+
+    merged = torch.cat(all_tensors, dim=0)
+    log.info(f'Total tokens before shuffle: {len(merged):,}')
+
+    # Shuffle to mix AICL + HF + scripts + synthetic
+    log.info(f'Shuffling {len(merged):,} tokens...')
+    perm = torch.randperm(len(merged))
+    merged = merged[perm]
+    del perm, all_tensors  # Free memory
+
+    merged_path = os.path.join(DATA_DIR, 'train_merged.pt')
+    torch.save(merged, merged_path)
+    size_mb = os.path.getsize(merged_path) / 1e6
+    log.info(f'Merged dataset saved: {merged_path} ({size_mb:.0f} MB, {len(merged):,} tokens)')
+
+    del merged  # Free memory
+
+    log.info(f'Data preparation complete. Files in {DATA_DIR}:')
+    for f in sorted(Path(DATA_DIR).glob('*.pt')):
+        log.info(f'  {f.name} ({f.stat().st_size / 1e6:.0f} MB)')
+
+
+# ═══════════════════════════════════════════════════════════
+#  Character Tokenizer (136 vocab)
+# ═══════════════════════════════════════════════════════════
+class CharTokenizer:
+    """Character-level tokenizer: 4 special + 132 printable/French chars = 136."""
+
+    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
+
+        chars = list(range(32, 127))
+        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)
+        log.info(f'Tokenizer saved to {path}')
+
+    @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 — supports multiple data formats
+# ═══════════════════════════════════════════════════════════
+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):
+            # Handle int16 or other dtypes from merge step
+            tokens = tokens.long()
+        self.tokens = tokens
+        self.seq_len = seq_len
+        log.info(f'Dataset: {len(self.tokens):,} tokens, {len(self):,} sequences')
+
+    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]
+        return chunk[:-1], chunk[1:]
+
+
+class MultiFileDataset(Dataset):
+    """Load and concatenate multiple .pt token files."""
+    def __init__(self, data_dir, seq_len=512):
+        pt_files = sorted(Path(data_dir).glob('*.pt'))
+        if not pt_files:
+            raise FileNotFoundError(f'No .pt files found in {data_dir}')
+
+        log.info(f'Loading {len(pt_files)} data files...')
+        all_tokens = []
+        for f in pt_files:
+            log.info(f'  Loading {f.name}...')
+            t = torch.load(f, map_location='cpu', weights_only=True)
+            all_tokens.append(t.long())
+
+        self.tokens = torch.cat(all_tokens, dim=0)
+        self.seq_len = seq_len
+        log.info(f'Combined dataset: {len(self.tokens):,} tokens, {len(self):,} sequences')
+
+    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]
+        return chunk[:-1], chunk[1:]
+
+
+# ═══════════════════════════════════════════════════════════
+#  CogNet-1B Architecture (~1.06B params)
+# ═══════════════════════════════════════════════════════════
+class CognitiveMemory(nn.Module):
+    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
+
+        w_att = torch.matmul(x, self.working.t())
+        w_att = F.softmax(w_att / math.sqrt(D), dim=-1)
+        w_out = self.work_proj(torch.matmul(w_att, self.working))
+
+        e_att = torch.matmul(x, self.episodic.t())
+        e_att = F.softmax(e_att / math.sqrt(D), dim=-1)
+        e_out = self.epis_proj(torch.matmul(e_att, self.episodic))
+
+        s_att = torch.matmul(x, self.semantic.t())
+        s_att = F.softmax(s_att / math.sqrt(D), dim=-1)
+        s_out = self.seman_proj(torch.matmul(s_att, self.semantic))
+
+        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)
+
+        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):
+    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_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)
+
+        self.ff = nn.Sequential(
+            nn.Linear(dim, ff_dim),
+            nn.GELU(),
+            nn.Linear(ff_dim, dim),
+        )
+        self.norm2 = nn.LayerNorm(dim)
+
+        self.cog_mem = CognitiveMemory(dim, working_slots, episodic_slots, semantic_slots)
+        self.norm3 = nn.LayerNorm(dim)
+
+        self.router = nn.Linear(dim, n_channels)
+
+    def forward(self, x):
+        B, T, D = x.shape
+
+        route_logits = self.router(x)
+        route_weights = F.softmax(route_logits, dim=-1)
+
+        channel_outs = []
+        for i, proj in enumerate(self.channel_projs):
+            ch_x = F.gelu(proj(x))
+            w = route_weights[:, :, i:i+1]
+            channel_outs.append(ch_x * w)
+
+        merged = self.channel_merge(torch.cat(channel_outs, dim=-1))
+        x = self.norm1(x + merged)
+
+        x = self.norm2(x + self.ff(x))
+        x = self.norm3(self.cog_mem(x))
+        return x
+
+
+class CogNet1B(nn.Module):
+    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
+        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))
+
+
+# ═══════════════════════════════════════════════════════════
+#  Save checkpoint
+# ═══════════════════════════════════════════════════════════
+def save_checkpoint(model, optimizer, step, loss_val, best_loss, tokenizer, args, filename):
+    path = os.path.join(args.ckpt_dir, filename)
+    save_dict = {
+        'step': step,
+        'model_state_dict': model.state_dict(),
+        'optimizer_state_dict': optimizer.state_dict(),
+        'loss': loss_val,
+        '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, path)
+    return path
+
+
+# ═══════════════════════════════════════════════════════════
+#  HuggingFace Push
+# ═══════════════════════════════════════════════════════════
+def push_to_huggingface(ckpt_path, tokenizer):
+    """Push the best checkpoint to HuggingFace."""
+    try:
+        from huggingface_hub import HfApi
+        repo_id = os.environ.get('HF_REPO', HF_REPO)
+
+        api = HfApi()
+        api.create_repo(repo_id=repo_id, exist_ok=True, token=HF_TOKEN)
+
+        # Upload checkpoint
+        api.upload_file(path_or_fileobj=ckpt_path, path_in_repo='checkpoints/cognet_1b_best.pt',
+                       repo_id=repo_id, token=HF_TOKEN)
+        # Upload tokenizer
+        api.upload_file(path_or_fileobj=TOKENIZER_PATH, path_in_repo='checkpoints/tokenizer_v3.json',
+                       repo_id=repo_id, token=HF_TOKEN)
+
+        log.info(f'Pushed to HuggingFace: {repo_id}')
+    except Exception as e:
+        log.warning(f'HF push failed (non-fatal): {e}')
+
+
+# ═══════════════════════════════════════════════════════════
+#  Find training data
+# ═══════════════════════════════════════════════════════════
+def find_data_path(args):
+    """Find the best training data file/directory."""
+    # 1. Explicit path
+    if args.data_path and os.path.exists(args.data_path):
+        return args.data_path
+
+    # 2. train_merged.pt (from cognet_data_prep.py merge step)
+    merged = os.path.join(DATA_DIR, 'train_merged.pt')
+    if os.path.exists(merged):
+        return merged
+
+    # 3. aicl_10x.pt (from fallback generation)
+    aicl = os.path.join(DATA_DIR, 'aicl_10x.pt')
+    if os.path.exists(aicl):
+        return aicl
+
+    # 4. Multiple .pt files — use data directory
+    pt_files = list(Path(DATA_DIR).glob('*.pt'))
+    if len(pt_files) >= 2:
+        return DATA_DIR  # Use MultiFileDataset
+
+    # 5. Single .pt file
+    if pt_files:
+        return str(pt_files[0])
+
+    return None
+
+
+# ═══════════════════════════════════════════════════════════
+#  Main Training Loop
+# ═══════════════════════════════════════════════════════════
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--batch-size', type=int, default=4)
+    parser.add_argument('--grad-accum', type=int, default=8,
+                       help='Gradient accumulation steps (effective batch = batch-size * grad-accum)')
+    parser.add_argument('--seq-len', type=int, default=512)
+    parser.add_argument('--max-steps', type=int, default=100000)
+    parser.add_argument('--warmup-steps', type=int, default=2000)
+    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=2000)
+    parser.add_argument('--eval-every', type=int, default=500)
+    parser.add_argument('--log-every', type=int, default=50)
+    parser.add_argument('--data-path', type=str, default=None,
+                       help='Explicit path to training data .pt file or directory')
+    parser.add_argument('--ckpt-dir', type=str, default=CKPT_DIR)
+    parser.add_argument('--resume', type=str, default=None)
+    parser.add_argument('--bf16', action='store_true', default=True)
+    parser.add_argument('--skip-data-prep', action='store_true',
+                       help='Skip data preparation (data already exists)')
+    parser.add_argument('--compile', action='store_true', default=False,
+                       help='Use torch.compile on model')
+    args = parser.parse_args()
+
+    # ── Setup ──
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    log.info('='*60)
+    log.info('CogNet-1B Training — RunPod AICL Edition (RTX PRO 6000)')
+    log.info('='*60)
+    log.info(f'Device: {device}')
+    log.info(f'Workspace: {WORKSPACE}')
+    log.info(f'Checkpoint dir: {args.ckpt_dir}')
+    log.info(f'HuggingFace repo: {HF_REPO}')
+
+    if torch.cuda.is_available():
+        for i in range(torch.cuda.device_count()):
+            props = torch.cuda.get_device_properties(i)
+            log.info(f'GPU {i}: {props.name} — {props.total_memory/1e9:.1f} GB VRAM')
+
+    # ── Step 1: Prepare data (AICL conversion + synthetic) ──
+    prepare_data(skip=args.skip_data_prep)
+
+    # ── Step 2: Load tokenizer ──
+    tokenizer_paths = [
+        TOKENIZER_PATH,
+        os.path.join(DATA_DIR, 'tokenizer_v3.json'),
+        os.path.join(WORKSPACE, 'hf_scripts', 'tokenizer_v3.json'),
+    ]
+
+    tokenizer = None
+    for tp in tokenizer_paths:
+        if os.path.exists(tp):
+            tokenizer = CharTokenizer.load(tp)
+            log.info(f'Loaded tokenizer from {tp} (vocab={tokenizer.vocab_size})')
+            # Copy to expected location
+            if tp != TOKENIZER_PATH:
+                import shutil
+                shutil.copy2(tp, TOKENIZER_PATH)
+                log.info(f'Copied tokenizer to {TOKENIZER_PATH}')
+            break
+
+    if tokenizer is None:
+        tokenizer = CharTokenizer()
+        tokenizer.save(TOKENIZER_PATH)
+        log.info(f'Created tokenizer (vocab={tokenizer.vocab_size})')
+
+    # ── Step 3: Load dataset ──
+    data_path = find_data_path(args)
+    if data_path is None:
+        log.error('No training data found! Run without --skip-data-prep or provide --data-path')
+        sys.exit(1)
+
+    log.info(f'Loading training data from: {data_path}')
+
+    if os.path.isdir(data_path):
+        dataset = MultiFileDataset(data_path, args.seq_len)
+    else:
+        dataset = TokenDataset(data_path, args.seq_len)
+
+    dataloader = DataLoader(dataset, batch_size=args.batch_size, shuffle=True,
+                           num_workers=4, pin_memory=True, drop_last=True)
+
+    # ── Step 4: Build model ──
+    log.info('Building CogNet-1B (~1.06B params)...')
+    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)
+
+    total_params = sum(p.numel() for p in model.parameters())
+    log.info(f'Total parameters: {total_params:,} ({total_params/1e9:.2f}B)')
+
+    # Optional: torch.compile for speed
+    if args.compile:
+        log.info('Compiling model with torch.compile...')
+        model = torch.compile(model)
+
+    # ── Step 5: Optimizer ──
+    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,
+    )
+
+    use_bf16 = args.bf16 and torch.cuda.is_bf16_supported()
+    scaler = None if use_bf16 else torch.amp.GradScaler('cuda')
+    log.info(f'Mixed precision: {"BF16" if use_bf16 else "FP16+GradScaler"}')
+
+    # ── Resume ──
+    start_step = 0
+    best_loss = float('inf')
+
+    # Try explicit resume path first
+    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'])
+        if 'optimizer_state_dict' in ckpt:
+            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}')
+
+    # Try auto-resume from latest checkpoint
+    if start_step == 0:
+        latest = os.path.join(args.ckpt_dir, 'cognet_1b_latest.pt')
+        if os.path.exists(latest):
+            log.info(f'Auto-resuming from {latest}...')
+            ckpt = torch.load(latest, map_location=device, weights_only=False)
+            model.load_state_dict(ckpt['model_state_dict'])
+            if 'optimizer_state_dict' in ckpt:
+                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 any step checkpoints newer than latest
+    if start_step == 0:
+        step_ckpts = sorted(Path(args.ckpt_dir).glob('cognet_1b_step_*.pt'))
+        if step_ckpts:
+            # Find the one with the highest step number
+            latest_step_ckpt = max(
+                step_ckpts,
+                key=lambda p: int(p.stem.split('_')[-1])
+            )
+            log.info(f'Found step checkpoint: {latest_step_ckpt}')
+            ckpt = torch.load(str(latest_step_ckpt), map_location=device, weights_only=False)
+            model.load_state_dict(ckpt['model_state_dict'])
+            if 'optimizer_state_dict' in ckpt:
+                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}')
+
+    # ── Train ──
+    effective_batch = args.batch_size * args.grad_accum
+    log.info(f'Starting: step {start_step} -> {args.max_steps}')
+    log.info(f'Batch={args.batch_size} x GradAccum={args.grad_accum} = Effective {effective_batch}')
+    log.info(f'SeqLen={args.seq_len}, LR={args.min_lr}-{args.max_lr}')
+    log.info(f'Checkpoints: {args.ckpt_dir} (NETWORK VOLUME)')
+    log.info(f'Graceful shutdown: SIGTERM/SIGINT will save checkpoint')
+    log.info(f'Data: {data_path}')
+
+    model.train()
+    data_iter = iter(dataloader)
+    t0 = time.time()
+    loss_val = 0.0
+    accum_loss = 0.0
+
+    for step in range(start_step, args.max_steps):
+        if shutdown_requested:
+            log.info('Shutdown requested, saving checkpoint...')
+            save_checkpoint(model, optimizer, step, loss_val, best_loss, tokenizer, args,
+                          'cognet_1b_latest.pt')
+            log.info(f'Checkpoint saved at step {step}. Exiting.')
+            break
+
+        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
+
+        # ── Gradient accumulation ──
+        optimizer.zero_grad(set_to_none=True)
+        accum_loss = 0.0
+
+        for micro_step in range(args.grad_accum):
+            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)
+
+            if use_bf16:
+                with torch.amp.autocast('cuda', dtype=torch.bfloat16):
+                    logits, loss = model(x, y)
+                (loss / args.grad_accum).backward()
+            else:
+                with torch.amp.autocast('cuda', dtype=torch.float16):
+                    logits, loss = model(x, y)
+                scaler.scale(loss / args.grad_accum).backward()
+
+            accum_loss += loss.item()
+
+        # Step optimizer after accumulation
+        if use_bf16:
+            grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
+            optimizer.step()
+        else:
+            scaler.unscale_(optimizer)
+            grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
+            scaler.step(optimizer)
+            scaler.update()
+
+        loss_val = accum_loss / args.grad_accum
+
+        # 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 * effective_batch * args.seq_len
+            vram_used = torch.cuda.memory_allocated() / 1e9
+
+            log.info(
+                f'Step {step:>7d}/{args.max_steps} | '
+                f'Loss: {loss_val:.4f} | PPL: {math.exp(min(loss_val, 20)):.1f} | '
+                f'LR: {lr:.2e} | Grad: {grad_norm:.2f} | '
+                f'VRAM: {vram_used:.1f}GB | {tokens_per_sec:.0f} tok/s'
+            )
+            t0 = time.time()
+
+        # 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=150, temperature=0.8, top_k=50)
+                sample_text = tokenizer.decode(sample_ids[0].tolist())
+                log.info(f'--- Sample step {step} ---')
+                log.info(sample_text[:300])
+                log.info(f'--- End ---')
+            model.train()
+
+        # Save checkpoint
+        if step > 0 and step % args.save_every == 0:
+            save_checkpoint(model, optimizer, step, loss_val, best_loss, tokenizer, args,
+                          f'cognet_1b_step_{step}.pt')
+            save_checkpoint(model, optimizer, step, loss_val, best_loss, tokenizer, args,
+                          'cognet_1b_latest.pt')
+
+            if loss_val < best_loss:
+                best_loss = loss_val
+                save_checkpoint(model, optimizer, step, loss_val, best_loss, tokenizer, args,
+                              'cognet_1b_best.pt')
+                log.info(f'New best! loss={loss_val:.4f}')
+
+            # Clean up old step checkpoints (keep only last 3)
+            step_ckpts = sorted(Path(args.ckpt_dir).glob('cognet_1b_step_*.pt'))
+            if len(step_ckpts) > 3:
+                for old_ckpt in step_ckpts[:-3]:
+                    old_ckpt.unlink()
+                    log.info(f'Removed old checkpoint: {old_ckpt.name}')
+
+            log.info(f'Checkpoint saved at step {step} (loss={loss_val:.4f}, best={best_loss:.4f})')
+
+    else:
+        # Training completed normally (not broken by shutdown)
+        final_path = save_checkpoint(model, optimizer, args.max_steps, loss_val, best_loss,
+                                    tokenizer, args, 'cognet_1b_final.pt')
+        log.info(f'Training complete! Final loss: {loss_val:.4f}, Best: {best_loss:.4f}')
+
+        # Push to HuggingFace
+        best_path = os.path.join(args.ckpt_dir, 'cognet_1b_best.pt')
+        if os.path.exists(best_path):
+            push_to_huggingface(best_path, tokenizer)
+
+
+if __name__ == '__main__':
+    main()