tinymind-native-8b-remote-handoff/bundle/train/phimind_full_train.py

"""Full RTX 3090 training pipeline for Φ-Mind.

Phases:
  1. Build/load BPE tokenizer (65k vocab, Thai+EN)
  2. Stream real data (distilled + forged + public)
  3. Train with BF16, torch.compile, gradient checkpointing
  4. Export checkpoint + metrics

Usage:
  python train/phimind_full_train.py --size small --steps 100000
  python train/phimind_full_train.py --size base  --steps 200000 --resume checkpoints/phimind/last.pt
"""

from __future__ import annotations

import argparse
import json
import math
import random
import time
from pathlib import Path
from typing import Iterator

import torch
import torch.nn as nn
import torch.nn.functional as F

# Optional: fast tokenizer
try:
    from tokenizers import Tokenizer as HFTokenizer
    _HF_TOK = True
except ImportError:
    _HF_TOK = False

from model.phimind import PhiMindConfig, PhiMindModel, count_params


# ---------------------------------------------------------------------------
# Tokenizer wrapper
# ---------------------------------------------------------------------------

class PhiMindTokenizer:
    """Wraps HuggingFace fast BPE tokenizer; falls back to UTF-8 bytes."""

    PAD, UNK, BOS, EOS = 0, 1, 2, 3

    def __init__(self, tokenizer_path: str | Path | None = None):
        self._hf: HFTokenizer | None = None
        self.vocab_size = 256 + 4  # fallback

        if tokenizer_path and Path(tokenizer_path).exists() and _HF_TOK:
            self._hf = HFTokenizer.from_file(str(tokenizer_path))
            self.vocab_size = self._hf.get_vocab_size()

    def encode(self, text: str, max_len: int = 2048) -> list[int]:
        if self._hf is not None:
            ids = self._hf.encode(text).ids
        else:
            usable = self.vocab_size - 4
            ids = [self.BOS] + [4 + (b % usable) for b in text.encode("utf-8")] + [self.EOS]
        return ids[:max_len]

    def decode(self, ids: list[int]) -> str:
        if self._hf is not None:
            return self._hf.decode(ids)
        chars = []
        for i in ids:
            if i in (self.BOS, self.EOS, self.PAD):
                continue
            chars.append(chr(max(0, i - 4)))
        return "".join(chars)


# ---------------------------------------------------------------------------
# Data streaming
# ---------------------------------------------------------------------------

def _iter_jsonl(path: Path) -> Iterator[dict]:
    try:
        with open(path, encoding="utf-8") as f:
            for line in f:
                line = line.strip()
                if line:
                    try:
                        yield json.loads(line)
                    except json.JSONDecodeError:
                        pass
    except OSError:
        pass


def _record_to_text(row: dict) -> str:
    """Convert any QA/chat record to a training string."""
    q = str(row.get("question", row.get("prompt", row.get("input", ""))))
    a = str(row.get("answer", row.get("response", row.get("output", ""))))
    think = str(row.get("thinking", ""))
    system = str(row.get("system", ""))

    parts = ["<bos>"]
    if system:
        parts.append(f"<system>{system}</system>\n")
    if think:
        parts.append(f"<user>{q}</user>\n<assistant><think>{think}</think>\n{a}<eos>")
    else:
        parts.append(f"<user>{q}</user>\n<assistant>{a}<eos>")
    return "".join(parts)


def stream_training_data(data_dir: Path, shuffle_seed: int = 42) -> Iterator[str]:
    """Yield training texts from all available JSONL files, shuffled."""
    paths = sorted(data_dir.glob("**/*.jsonl"))
    random.seed(shuffle_seed)
    random.shuffle(paths)

    for path in paths:
        rows = list(_iter_jsonl(path))
        random.shuffle(rows)
        for row in rows:
            text = _record_to_text(row)
            if len(text) > 20:
                yield text


def build_dataset(
    data_dir: Path,
    tokenizer: PhiMindTokenizer,
    max_seq_len: int,
    max_tokens: int = 50_000_000,
    seed: int = 42,
) -> list[torch.Tensor]:
    """Load and tokenize all data into a list of token tensors."""
    sequences: list[torch.Tensor] = []
    total_tokens = 0
    for text in stream_training_data(data_dir, seed):
        ids = tokenizer.encode(text, max_seq_len)
        if len(ids) < 4:
            continue
        sequences.append(torch.tensor(ids, dtype=torch.long))
        total_tokens += len(ids)
        if total_tokens >= max_tokens:
            break
    return sequences


# ---------------------------------------------------------------------------
# Model configs per size
# ---------------------------------------------------------------------------

def build_model_config(size: str, vocab_size: int) -> PhiMindConfig:
    configs = {
        "tiny": PhiMindConfig(
            vocab_size=vocab_size, dim=256, n_layers=8, max_seq_len=2048,
            hrr_local_window=128, soliton_n_modes=64, rg_period=4,
        ),
        "small": PhiMindConfig(
            vocab_size=vocab_size, dim=512, n_layers=12, max_seq_len=4096,
            hrr_local_window=256, soliton_n_modes=128, rg_period=4,
        ),
        "base": PhiMindConfig(
            vocab_size=vocab_size, dim=1024, n_layers=16, max_seq_len=8192,
            hrr_local_window=512, soliton_n_modes=256, rg_period=4,
        ),
    }
    if size not in configs:
        raise ValueError(f"size must be one of {list(configs)}")
    return configs[size]


# ---------------------------------------------------------------------------
# Training helpers
# ---------------------------------------------------------------------------

def _collate_batch(
    sequences: list[torch.Tensor],
    max_len: int,
    pad_id: int = 0,
) -> tuple[torch.Tensor, torch.Tensor]:
    length = min(max(s.numel() for s in sequences), max_len)
    B = len(sequences)
    input_ids = torch.full((B, length), pad_id, dtype=torch.long)
    labels = torch.full((B, length), -100, dtype=torch.long)
    for i, seq in enumerate(sequences):
        n = min(seq.numel(), length)
        input_ids[i, :n] = seq[:n]
        labels[i, :n] = seq[:n]
    return input_ids, labels


def _causal_loss(logits: torch.Tensor, labels: torch.Tensor) -> torch.Tensor:
    return F.cross_entropy(
        logits[:, :-1].reshape(-1, logits.size(-1)),
        labels[:, 1:].reshape(-1),
        ignore_index=-100,
    )


# ---------------------------------------------------------------------------
# Main trainer
# ---------------------------------------------------------------------------

class FullTrainer:
    def __init__(
        self,
        cfg: PhiMindConfig,
        sequences: list[torch.Tensor],
        out_dir: Path,
        train_steps: int = 100_000,
        batch_size: int = 4,
        grad_accum: int = 8,
        lr: float = 3e-4,
        warmup_steps: int = 2000,
        eval_interval: int = 1000,
        save_interval: int = 5000,
        seed: int = 42,
        resume: str | None = None,
        use_compile: bool = True,
    ):
        torch.manual_seed(seed)
        self.out_dir = out_dir
        self.out_dir.mkdir(parents=True, exist_ok=True)
        self.train_steps = train_steps
        self.batch_size = batch_size
        self.grad_accum = grad_accum
        self.warmup_steps = warmup_steps
        self.eval_interval = eval_interval
        self.save_interval = save_interval
        self.cfg = cfg

        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        self.dtype = torch.bfloat16 if self.device.type == "cuda" else torch.float32

        self.model = PhiMindModel(cfg).to(self.device)
        if self.device.type == "cuda" and use_compile:
            try:
                self.model = torch.compile(self.model)  # type: ignore[assignment]
                print("torch.compile enabled")
            except Exception as e:
                print(f"torch.compile skipped: {e}")

        self.optimizer = torch.optim.AdamW(
            self.model.parameters(), lr=lr,
            weight_decay=0.01, betas=(0.9, 0.95),
        )
        self.scaler = torch.amp.GradScaler(enabled=(self.dtype == torch.bfloat16))

        n_train = max(1, int(len(sequences) * 0.95))
        self.train_seqs = sequences[:n_train]
        self.eval_seqs = sequences[n_train:] or sequences[:min(64, len(sequences))]

        self.start_step = 0
        self.lr = lr
        if resume and Path(resume).exists():
            self._load_checkpoint(resume)

    def _load_checkpoint(self, path: str) -> None:
        ckpt = torch.load(path, map_location=self.device, weights_only=False)
        state = ckpt.get("model_state", ckpt)
        # Handle compiled model prefix
        if hasattr(self.model, "_orig_mod"):
            self.model._orig_mod.load_state_dict(state, strict=False)
        else:
            self.model.load_state_dict(state, strict=False)
        if "optimizer_state" in ckpt:
            self.optimizer.load_state_dict(ckpt["optimizer_state"])
        self.start_step = int(ckpt.get("step", 0))
        print(f"Resumed from step {self.start_step}")

    def _lr_scale(self, step: int) -> float:
        if step < self.warmup_steps:
            return (step + 1) / max(self.warmup_steps, 1)
        progress = (step - self.warmup_steps) / max(self.train_steps - self.warmup_steps, 1)
        return max(0.05, 0.5 * (1 + math.cos(math.pi * progress)))

    @torch.no_grad()
    def _eval_loss(self) -> float:
        self.model.eval()
        total, count = 0.0, 0
        idxs = list(range(len(self.eval_seqs)))
        random.shuffle(idxs)
        for i in idxs[:min(20, len(idxs))]:
            seq = self.eval_seqs[i]
            ids, lbls = _collate_batch([seq], self.cfg.max_seq_len)
            ids, lbls = ids.to(self.device), lbls.to(self.device)
            with torch.amp.autocast(device_type=self.device.type, dtype=self.dtype):
                out = self.model(ids)
                loss = _causal_loss(out["logits"], lbls)
            if torch.isfinite(loss):
                total += loss.item()
                count += 1
        self.model.train()
        return total / max(count, 1)

    def _save(self, step: int, eval_loss: float, tag: str = "last") -> Path:
        path = self.out_dir / f"phimind_{tag}.pt"
        model_state = (
            self.model._orig_mod.state_dict()
            if hasattr(self.model, "_orig_mod")
            else self.model.state_dict()
        )
        torch.save({
            "step": step,
            "model_state": model_state,
            "optimizer_state": self.optimizer.state_dict(),
            "model_cfg": self.cfg,
            "eval_loss": eval_loss,
        }, path)
        return path

    def train(self) -> dict:
        print(f"\nΦ-Mind Training")
        print(f"  Device: {self.device} | dtype: {self.dtype}")
        print(f"  Params: {count_params(self.model)}")
        print(f"  Train seqs: {len(self.train_seqs):,} | Eval seqs: {len(self.eval_seqs):,}")
        print(f"  Steps: {self.train_steps:,} | Batch: {self.batch_size} × accum {self.grad_accum}")

        self.model.train()
        self.optimizer.zero_grad()

        t0 = time.perf_counter()
        log: list[dict] = []
        best_eval = float("inf")
        accum_loss = 0.0
        micro = 0

        for step in range(self.start_step, self.train_steps):
            scale = self._lr_scale(step)
            for pg in self.optimizer.param_groups:
                pg["lr"] = self.lr * scale

            idx = random.randrange(len(self.train_seqs))
            batch = [
                self.train_seqs[(idx + i) % len(self.train_seqs)]
                for i in range(self.batch_size)
            ]
            ids, lbls = _collate_batch(batch, self.cfg.max_seq_len)
            ids, lbls = ids.to(self.device), lbls.to(self.device)

            with torch.amp.autocast(device_type=self.device.type, dtype=self.dtype):
                out = self.model(ids)
                loss = _causal_loss(out["logits"], lbls) / self.grad_accum

            self.scaler.scale(loss).backward()
            accum_loss += loss.item()
            micro += 1

            if micro % self.grad_accum == 0:
                self.scaler.unscale_(self.optimizer)
                gn = torch.nn.utils.clip_grad_norm_(self.model.parameters(), 1.0)
                self.scaler.step(self.optimizer)
                self.scaler.update()
                self.optimizer.zero_grad()

                train_loss = accum_loss * self.grad_accum
                accum_loss = 0.0

                if (step + 1) % 50 == 0:
                    elapsed = time.perf_counter() - t0
                    tokens_sec = (step + 1) * self.batch_size * self.cfg.max_seq_len / elapsed
                    print(
                        f"step {step+1:6d} | loss {train_loss:.4f} | "
                        f"lr {self.lr * scale:.2e} | gn {float(gn):.2f} | "
                        f"{tokens_sec:.0f} tok/s"
                    )
                    log.append({
                        "step": step + 1, "train_loss": train_loss,
                        "lr": self.lr * scale, "grad_norm": float(gn),
                    })

            if (step + 1) % self.eval_interval == 0:
                ev = self._eval_loss()
                ppl = math.exp(min(ev, 20))
                print(f"  [eval] step {step+1} | loss {ev:.4f} | ppl {ppl:.1f}")
                if ev < best_eval:
                    best_eval = ev
                    self._save(step + 1, ev, "best")

            if (step + 1) % self.save_interval == 0:
                self._save(step + 1, best_eval, "last")

        final_eval = self._eval_loss()
        final_path = self._save(self.train_steps, final_eval, "final")

        result = {
            "train_steps": self.train_steps,
            "final_eval_loss": final_eval,
            "best_eval_loss": best_eval,
            "final_perplexity": math.exp(min(final_eval, 20)),
            "checkpoint": str(final_path),
            "elapsed_s": time.perf_counter() - t0,
            "param_count": count_params(self.model),
            "log": log,
        }
        (self.out_dir / "train_result.json").write_text(
            json.dumps(result, ensure_ascii=False, indent=2), encoding="utf-8"
        )
        return result


# ---------------------------------------------------------------------------
# CLI
# ---------------------------------------------------------------------------

def main() -> None:
    ap = argparse.ArgumentParser(description="Train Φ-Mind on RTX 3090")
    ap.add_argument("--size", default="small", choices=["tiny", "small", "base"])
    ap.add_argument("--steps", type=int, default=100_000)
    ap.add_argument("--batch", type=int, default=4)
    ap.add_argument("--accum", type=int, default=8)
    ap.add_argument("--lr", type=float, default=3e-4)
    ap.add_argument("--warmup", type=int, default=2000)
    ap.add_argument("--data-dir", default="data/filtered")
    ap.add_argument("--tokenizer", default="data/tokenizer/tokenizer.json")
    ap.add_argument("--out-dir", default="checkpoints/phimind")
    ap.add_argument("--max-tokens", type=int, default=50_000_000)
    ap.add_argument("--eval-interval", type=int, default=1000)
    ap.add_argument("--save-interval", type=int, default=5000)
    ap.add_argument("--resume", default=None)
    ap.add_argument("--no-compile", action="store_true")
    ap.add_argument("--seed", type=int, default=42)
    args = ap.parse_args()

    tok = PhiMindTokenizer(args.tokenizer)
    print(f"Tokenizer: vocab_size={tok.vocab_size:,}")

    cfg = build_model_config(args.size, tok.vocab_size)
    print(f"Model: Φ-Mind-{args.size} | dim={cfg.dim} | layers={cfg.n_layers}")

    data_dir = Path(args.data_dir)
    print(f"Loading data from {data_dir} ...")
    sequences = build_dataset(data_dir, tok, cfg.max_seq_len, args.max_tokens, args.seed)
    if not sequences:
        print("WARNING: no data found — using synthetic fallback")
        sequences = [
            torch.randint(4, tok.vocab_size, (cfg.max_seq_len // 4,))
            for _ in range(100)
        ]
    print(f"Loaded {len(sequences):,} sequences")

    trainer = FullTrainer(
        cfg=cfg,
        sequences=sequences,
        out_dir=Path(args.out_dir),
        train_steps=args.steps,
        batch_size=args.batch,
        grad_accum=args.accum,
        lr=args.lr,
        warmup_steps=args.warmup,
        eval_interval=args.eval_interval,
        save_interval=args.save_interval,
        seed=args.seed,
        resume=args.resume,
        use_compile=not args.no_compile,
    )
    result = trainer.train()

    print(f"\nDone. Final perplexity: {result['final_perplexity']:.2f}")
    print(f"Checkpoint: {result['checkpoint']}")


if __name__ == "__main__":
    main()