scripts/train_1b_gpu.py

#!/usr/bin/env python
"""GPU training script for Fractus 1B with progressive growth.

Loads the best CPU-trained checkpoint (palier 3, d=768) and grows it to 1B,
then trains on GPU with all optimizations active.

Usage (on a GPU machine):
    python scripts/train_1b_gpu.py \\
        --checkpoint checkpoints/fractus_palier3.pt \\
        --tokens 1760000000 \\
        --batch-size 4 \\
        --seq-len 64 \\
        --lr 1e-4 \\
        --accumulation-steps 4 \\
        --bf16

Expected on RTX 3090 (24GB):
    - Forward+backward per token: ~0.5ms → ~2000 tok/s
    - With sparse MoE (2/128): ~3000 tok/s effective
    - With PGSU (4/16 layers): ~1.5x more → ~4500 tok/s
    - 1.76B tokens at 4000 tok/s ≈ 5 days

With progressive growth warm start:
    - Palier 4 starts from palier 3 weights (d=768 trained)
    - Needs ~1/4 of Chinchilla to converge → ~440M tokens
    - 440M at 4000 tok/s ≈ 30 hours → ~1.5 days
"""
import argparse, os, sys, time, math, json
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
import torch
import torch.nn.functional as F
from fractus.continuous_engine import ContinuousThoughtEngine
from fractus.grow import grow_cte
from fractus.tokenizer import FractusTokenizer

CORPUS = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))),
                      "data", "quality_corpus.pt")

# 1B target config (white paper config K).
TARGET_1B = dict(
    d_model=1280, n_heads=20, d_head=64, n_levels=2,
    n_oscillators=16, coupling_rank=8,
    n_experts=128, top_k=2, expert_d_ff=2048, siren_rank=64,
)


def load_checkpoint(engine, ckpt_path):
    """Load weights from a checkpoint, ignoring buffer mismatches."""
    ckpt = torch.load(ckpt_path, weights_only=False, map_location="cpu")
    model_sd = ckpt["model_state"]
    own_sd = engine.state_dict()
    for key, val in model_sd.items():
        if key in own_sd and own_sd[key].shape == val.shape:
            own_sd[key] = val
    engine.load_state_dict(own_sd)
    cfg = ckpt.get("config", {})
    print(f"  Loaded checkpoint: {cfg.get('palier', '?')}, "
          f"d={cfg.get('d_model', '?')}, E={cfg.get('n_experts', '?')}", flush=True)
    return engine


def train_1b_gpu(engine, tokens, n_tokens, lr, batch_size, seq_len,
                 accumulation_steps, use_bf16, pgsu_active, log_every=500):
    """Train the 1B model on GPU with all optimizations.

    Optimizations active:
    - tick_chunk_train: head on last position only (seq_len x less head FLOPs)
    - Sparse MoE low-rank: only top-2/128 experts computed (64x less MoE work)
    - Gradient accumulation: fewer optimizer steps
    - bf16 AMP: 2x on all matmuls, halved memory
    - PGSU: 4/16 layers active per step (if enabled)
    """
    device = next(engine.parameters()).device
    vocab = engine.vocab_size
    dtype = torch.bfloat16 if use_bf16 else torch.float32

    opt = torch.optim.AdamW(engine.parameters(), lr=lr, weight_decay=0.01)

    # PGSU setup.
    pgsu = None
    if pgsu_active:
        try:
            from fractus1B.pgsu import PGSU
            # Note: PGSU works on Fractus1B (16 blocks). For the CTE (single MoE),
            # PGSU is a no-op. This is here for when we switch to Fractus1B.
            print("  PGSU: available for Fractus1B (not used on CTE)", flush=True)
        except Exception:
            pass

    engine.train()
    engine.reset_thought(batch_size=1)

    t0 = time.time()
    total_loss = 0.0
    total_correct = 0
    total_n = 0
    chunk_idx = 0
    opt.zero_grad()

    g = torch.Generator().manual_seed(42)
    n = tokens.numel()

    for start in range(0, min(n_tokens, n - seq_len - 1), seq_len):
        chunk = tokens[start:start + seq_len].unsqueeze(0).to(device)
        target = tokens[start + seq_len].to(device)

        if use_bf16:
            with torch.autocast(device_type="cuda", dtype=dtype):
                last_logits = engine.tick_chunk_train(chunk)
                loss = F.cross_entropy(last_logits, target.unsqueeze(0)) / accumulation_steps
        else:
            last_logits = engine.tick_chunk_train(chunk)
            loss = F.cross_entropy(last_logits, target.unsqueeze(0)) / accumulation_steps

        loss.backward()

        total_loss += loss.item() * accumulation_steps
        pred = last_logits.argmax(dim=-1)
        total_correct += (pred == target.unsqueeze(0)).sum().item()
        total_n += 1
        chunk_idx += 1

        if chunk_idx % accumulation_steps == 0:
            torch.nn.utils.clip_grad_norm_(engine.parameters(), 1.0)
            opt.step()
            opt.zero_grad()

        if chunk_idx % log_every == 0:
            processed = chunk_idx * seq_len
            elapsed = time.time() - t0
            rate = processed / max(elapsed, 1)
            avg = total_loss / max(total_n, 1)
            acc = total_correct / max(total_n, 1)
            ppl = math.exp(min(avg, 20))
            mem_gb = torch.cuda.max_memory_allocated() / 1e9 if torch.cuda.is_available() else 0
            print(f"  {processed:>10,}/{n_tokens:,} loss={avg:.3f} ppl={ppl:.1f} "
                  f"acc={acc:.3f} {rate:.0f} tok/s "
                  f"GPU_mem={mem_gb:.1f}GB", flush=True)

    # Final remainder.
    if chunk_idx % accumulation_steps != 0:
        torch.nn.utils.clip_grad_norm_(engine.parameters(), 1.0)
        opt.step()

    elapsed = time.time() - t0
    avg_loss = total_loss / max(total_n, 1)
    ppl = math.exp(min(avg_loss, 20))
    print(f"\n  DONE: loss={avg_loss:.3f} ppl={ppl:.1f} "
          f"({chunk_idx * seq_len:,} tokens in {elapsed/3600:.1f}h, "
          f"{chunk_idx * seq_len / max(elapsed,1):.0f} tok/s)", flush=True)
    return engine


def evaluate(engine, tokens, n_eval=1000, seq_len=64):
    """Quick eval: perplexity on holdout."""
    engine.eval()
    device = next(engine.parameters()).device
    holdout = tokens[-n_eval:]
    total_nll, n = 0.0, 0
    engine.reset_thought(batch_size=1)
    for s in range(0, min(len(holdout) - seq_len - 1, n_eval), seq_len):
        chunk = holdout[s:s + seq_len].unsqueeze(0).to(device)
        target = holdout[s + seq_len].to(device)
        with torch.no_grad():
            logits = engine.tick_chunk_train(chunk)
            nll = F.cross_entropy(logits, target.unsqueeze(0))
        total_nll += nll.item()
        n += 1
    avg = total_nll / max(n, 1)
    return avg, math.exp(min(avg, 20))


def generate_sample(engine, tok, prompt_text, n_tokens=60):
    """Greedy decode from prompt."""
    engine.eval()
    engine.reset_thought(batch_size=1)
    device = next(engine.parameters()).device
    ids = tok.encode(prompt_text)
    for t in ids:
        engine.tick(torch.tensor([t], device=device))
    cur = torch.tensor([ids[-1]], device=device)
    generated = list(ids)
    for _ in range(n_tokens):
        with torch.no_grad():
            logits, _ = engine.tick(cur)
        nxt = int(logits.argmax(-1).item())
        generated.append(nxt)
        cur = torch.tensor([nxt], device=device)
    return tok.decode(generated)


def main():
    ap = argparse.ArgumentParser(description="Fractus 1B GPU Training")
    ap.add_argument("--checkpoint", type=str, default=None,
                    help="CPU-trained checkpoint to grow from (e.g. fractus_palier3.pt)")
    ap.add_argument("--tokens", type=int, default=440_000_000,
                    help="training tokens (default: 440M = ~1/4 Chinchilla for warm start)")
    ap.add_argument("--batch-size", type=int, default=4)
    ap.add_argument("--seq-len", type=int, default=64)
    ap.add_argument("--lr", type=float, default=1e-4)
    ap.add_argument("--accumulation-stays", type=int, default=4)
    ap.add_argument("--accumulation-steps", type=int, default=4)
    ap.add_argument("--bf16", action="store_true", default=True,
                    help="enable bf16 mixed precision (default: on)")
    ap.add_argument("--no-bf16", dest="bf16", action="store_false")
    ap.add_argument("--pgsu", action="store_true", help="enable PGSU (Fractus1B only)")
    ap.add_argument("--corpus", type=str, default=CORPUS)
    ap.add_argument("--seed", type=int, default=42)
    ap.add_argument("--eval-interval", type=int, default=50000,
                    help="evaluate PPL every N tokens")
    ap.add_argument("--save-interval", type=int, default=100000,
                    help="save checkpoint every N tokens")
    args = ap.parse_args()

    torch.manual_seed(args.seed)
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"=== Fractus 1B GPU Training ===", flush=True)
    print(f"Device: {device}", flush=True)
    if device.type == "cuda":
        print(f"GPU: {torch.cuda.get_device_name(0)}", flush=True)
        print(f"VRAM: {torch.cuda.get_device_properties(0).total_memory/1e9:.1f} GB", flush=True)
        print(f"bf16: {args.bf16}", flush=True)
    else:
        print("WARNING: no GPU detected — training will be extremely slow!", flush=True)

    torch.set_num_threads(os.cpu_count() or 6)

    # Load corpus.
    print(f"Loading corpus: {args.corpus}", flush=True)
    tokens = torch.load(args.corpus, weights_only=False).to(torch.int64)
    print(f"Corpus: {len(tokens):,} tokens", flush=True)

    # Build engine — either grow from checkpoint or build fresh 1B.
    if args.checkpoint and os.path.exists(args.checkpoint):
        print(f"\nLoading checkpoint: {args.checkpoint}", flush=True)
        # Build a model matching the checkpoint config, load, then grow.
        ckpt = torch.load(args.checkpoint, weights_only=False, map_location="cpu")
        cfg = ckpt["config"]
        engine = ContinuousThoughtEngine(
            vocab_size=50257,
            d_model=cfg["d_model"], n_heads=cfg["n_heads"],
            d_head=cfg.get("d_head", 64), n_levels=2,
            n_oscillators=8, coupling_rank=4,
            n_experts=cfg["n_experts"], top_k=2,
            expert_d_ff=cfg["expert_d_ff"],
            siren_rank=cfg["siren_rank"])
        engine = load_checkpoint(engine, args.checkpoint)

        # Grow to 1B target.
        print(f"\nGrowing to 1B target: d={TARGET_1B['d_model']}, "
              f"E={TARGET_1B['n_experts']}", flush=True)
        engine = grow_cte(engine, TARGET_1B)
        print(f"  Grown: d={engine.d_model}, E={engine.blocks[0].moe.n_experts}, "
              f"params={sum(p.numel() for p in engine.parameters()):,}", flush=True)
    else:
        print("\nBuilding 1B from scratch (no checkpoint)", flush=True)
        engine = ContinuousThoughtEngine(vocab_size=50257, **TARGET_1B)
        print(f"  params={sum(p.numel() for p in engine.parameters()):,}", flush=True)

    engine = engine.to(device)

    # Eval before training.
    nll_before, ppl_before = evaluate(engine, tokens)
    print(f"\nBefore: NLL={nll_before:.3f} PPL={ppl_before:.1f}", flush=True)

    # Train.
    print(f"\nTraining: {args.tokens:,} tokens, lr={args.lr}, "
          f"accum={args.accumulation_steps}, bf16={args.bf16}", flush=True)
    engine = train_1b_gpu(
        engine, tokens, args.tokens, lr=args.lr,
        batch_size=args.batch_size, seq_len=args.seq_len,
        accumulation_steps=args.accumulation_steps,
        use_bf16=args.bf16, pgsu_active=args.pgsu)

    # Eval after.
    nll_after, ppl_after = evaluate(engine, tokens)
    print(f"After:  NLL={nll_after:.3f} PPL={ppl_after:.1f}", flush=True)

    # Generation test.
    tok = FractusTokenizer.gpt2_compatible()
    print(f"\n=== Generation Test ===", flush=True)
    for prompt in ["The function", "def fractus", "import torch", "Hello, my name is"]:
        text = generate_sample(engine, tok, prompt, n_tokens=60)
        print(f'  "{text}"', flush=True)

    # Save.
    ckpt_path = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))),
                             "checkpoints", "fractus_1b_gpu.pt")
    os.makedirs(os.path.dirname(ckpt_path), exist_ok=True)
    torch.save({
        "model_state": engine.state_dict(),
        "config": {**TARGET_1B, "method": "gpu_progressive_growth"},
        "params": sum(p.numel() for p in engine.parameters()),
        "ppl": ppl_after,
    }, ckpt_path)
    print(f"\nSaved: {ckpt_path} ({os.path.getsize(ckpt_path)/1e9:.2f}GB)", flush=True)

    # Upload to HF.
    hf_token = os.environ.get("HF_TOKEN")
    if hf_token:
        try:
            from huggingface_hub import HfApi
            api = HfApi(token=hf_token)
            api.upload_file(
                path_or_fileobj=ckpt_path,
                path_in_repo="checkpoints/fractus_1b_gpu.pt",
                repo_id="thefinalboss/Fractus-1B", repo_type="model")
            print("Uploaded to HuggingFace: thefinalboss/Fractus-1B", flush=True)
        except Exception as e:
            print(f"HF upload failed: {e}", flush=True)


if __name__ == "__main__":
    main()