code/train_v3.py

#!/usr/bin/env python3
"""Berna R7 v3 - CLEAN training from scratch (fixed dataset)."""
import sys, json, time, argparse, shutil
from pathlib import Path
from datetime import datetime, timezone
import numpy as np
import torch
from torch.utils.data import DataLoader
# import bitsandbytes as bnb
from transformers import PretrainedConfig
from transformers.models.auto.configuration_auto import CONFIG_MAPPING

CODE_DIR = Path("/data/berna-r7-v3/core/code")
sys.path.insert(0, str(CODE_DIR))

from modeling_berna_v2 import BernaForCausalLM, BernaConfig
from cell_integration import CellSystem, build_cells_and_dnas
from dataset import BernaDataset

ROOT = Path("/data/berna-r7-v3/core")
CKPT_DIR = ROOT / "checkpoints"
LOG_DIR = ROOT / "logs"
LOG_FILE = LOG_DIR / "train.log"

SEQ_LEN = 2048
MICRO_BATCH = 1
GRAD_ACCUM = 128
LR = 3e-4
WARMUP = 500
TOTAL_STEPS = 10200
SAVE_EVERY = 100
LOG_EVERY = 50
KEEP_LAST = 3
CELL_MEASURE_EVERY = 100

CONFIG_MAPPING.register("berna", BernaConfig, exist_ok=True)


def log(msg):
    ts = datetime.now(timezone.utc).strftime("%Y-%m-%d %H:%M:%S")
    line = f"[{ts}] {msg}"
    print(line, flush=True)
    LOG_DIR.mkdir(parents=True, exist_ok=True)
    with open(LOG_FILE, "a") as f:
        f.write(line + "\n")


def load_config():
    with open(ROOT / "config" / "config.json") as f:
        d = json.load(f)
    return BernaConfig(**d), d


def lr_lambda(step):
    if step < WARMUP:
        return step / max(1, WARMUP)
    p = (step - WARMUP) / max(1, TOTAL_STEPS - WARMUP)
    p = min(p, 1.0)
    return 0.5 * (1.0 + np.cos(np.pi * p))


def save_checkpoint(model, optimizer, scheduler, step, tokens_seen, loss):
    CKPT_DIR.mkdir(parents=True, exist_ok=True)
    final = CKPT_DIR / f"step_{step:08d}"
    tmp = CKPT_DIR / f".tmp_{step:08d}"
    if tmp.exists():
        shutil.rmtree(tmp)
    tmp.mkdir()
    (tmp / "model").mkdir()
    model.save_pretrained(tmp / "model", safe_serialization=True)
    torch.save(optimizer.state_dict(), tmp / "optimizer.pt")
    torch.save(scheduler.state_dict(), tmp / "scheduler.pt")
    meta = {"step": step, "tokens_seen": tokens_seen, "loss": loss,
            "saved_at": datetime.now(timezone.utc).isoformat()}
    (tmp / "metadata.json").write_text(json.dumps(meta, indent=2))
    if final.exists():
        shutil.rmtree(final)
    tmp.rename(final)
    return final


def cleanup_old_checkpoints(keep=KEEP_LAST):
    if not CKPT_DIR.exists():
        return
    dirs = sorted([d for d in CKPT_DIR.iterdir()
                   if d.is_dir() and d.name.startswith("step_")])
    for d in dirs[:-keep]:
        shutil.rmtree(d)


def train(args):
    LOG_DIR.mkdir(parents=True, exist_ok=True)
    log("=" * 60)
    log("Berna R7 v3 - CLEAN (fixed dataset, V2 K6, 1B from scratch)")
    log("=" * 60)

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    log(f"Device: {device}")
    if torch.cuda.is_available():
        log(f"GPU: {torch.cuda.get_device_name(0)}")

    cfg, cfg_dict = load_config()

    cell_sys = CellSystem(
        model_uuid="mother-001",
        base_dir=str(ROOT / "cells_runtime"),
        measure_every=CELL_MEASURE_EVERY,
        verbose=True,
    )
    cells_by_layer, dnas = build_cells_and_dnas(cfg_dict, model_uuid="mother-001")
    cell_sys.cells_by_layer = cells_by_layer
    cell_sys.dnas = dnas

    model = BernaForCausalLM(cfg, cells_by_layer=cells_by_layer, dnas=dnas).to(device)
    # TEST_LOSS_QUICK
    import numpy as np
    _test_eng = np.load('/data/berna-r7-data/tokenized/eng.npy', mmap_mode='r')
    _x = torch.from_numpy(_test_eng[100000:100512].astype(np.int64)).unsqueeze(0).to(device)
    _y = torch.from_numpy(_test_eng[100001:100513].astype(np.int64)).unsqueeze(0).to(device)
    with torch.no_grad():
        _out = model(input_ids=_x, labels=_y)
        print(f"QUICK_TEST_LOSS: {_out.loss.item():.4f}")
    
    n = sum(p.numel() for p in model.parameters())
    log(f"Model: {n:,} parameters")
    log(f"Cells: {sum(len(v) for v in cells_by_layer.values())} initial")

    model.gradient_checkpointing_enable()
    log("Gradient checkpointing: ENABLED")

    optimizer = torch.optim.AdamW(model.parameters(), lr=LR,
                                     betas=(0.9, 0.95), weight_decay=0.1)
    log("Optimizer: AdamW8bit")

    scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda)

    from cells.cell_manager import CellManager
    cell_sys.manager = CellManager(
        model_uuid="mother-001",
        base_dir=Path(str(ROOT / "cells_runtime" / "mother-001")),
    )
    all_cells = []
    for cs in cell_sys.cells_by_layer.values():
        all_cells.extend(cs)
    cell_sys.manager.cells = all_cells
    cell_sys.manager.dnas = cell_sys.dnas

    dataset = BernaDataset(seq_len=SEQ_LEN)
    loader = DataLoader(dataset, batch_size=MICRO_BATCH, shuffle=False,
                        num_workers=0, pin_memory=True)
    it = iter(loader)

    log(f"Total steps: {TOTAL_STEPS}")
    log(f"Effective batch: {MICRO_BATCH * GRAD_ACCUM} x {SEQ_LEN} = "
        f"{MICRO_BATCH * GRAD_ACCUM * SEQ_LEN:,} tokens/step")
    log(f"Total tokens planned: {TOTAL_STEPS * MICRO_BATCH * GRAD_ACCUM * SEQ_LEN:,}")
    log(f"Cell measurement: every {CELL_MEASURE_EVERY} steps")
    log("=" * 60)

    model.train()
    t0 = time.time()
    last_log_tokens = 0
    last_log_time = t0
    tokens_seen = 0
    start_step = 0

    if args.resume and CKPT_DIR.exists():
        dirs = sorted([d for d in CKPT_DIR.iterdir()
                       if d.is_dir() and d.name.startswith("step_")])
        if dirs:
            latest = dirs[-1]
            log(f"Resuming from: {latest}")
            from safetensors.torch import load_file
            sd = load_file(str(latest / "model" / "model.safetensors"))
            model.load_state_dict(sd, strict=False)
            opt_file = latest / "optimizer.pt"
            if opt_file.exists():
                log("SKIP: optimizer state (bnb incompatible)")
            sch_file = latest / "scheduler.pt"
            if sch_file.exists():
                scheduler.load_state_dict(torch.load(sch_file, map_location=device, weights_only=False))
            meta_file = latest / "metadata.json"
            if meta_file.exists():
                meta = json.loads(meta_file.read_text())
                start_step = meta["step"]
                tokens_seen = meta["tokens_seen"]
                log(f"Resumed at step {start_step}, tokens {tokens_seen:,}")

                # ===== NEW (2026-10-04): Restore cells from disk =====
                try:
                    if cell_sys.load_state():
                        cell_sys.sync_into_model(model)
                        n_active = sum(1 for c in cell_sys.manager.cells if c.state == "active")
                        log(f"[resume] cells: {len(cell_sys.manager.cells)} total, {n_active} active")
                    else:
                        log(f"[resume] WARN: no cells state on disk, using {len(all_cells)} initial")
                except Exception as e:
                    log(f"[resume] WARN: cells load failed: {type(e).__name__}: {e}")
                # ====================================================

    for step in range(start_step, TOTAL_STEPS):
        optimizer.zero_grad(set_to_none=True)
        total_loss = 0.0
        for _ in range(GRAD_ACCUM):
            try:
                x, y = next(it)
            except StopIteration:
                it = iter(loader)
                x, y = next(it)
            x = x.to(device, non_blocking=True)
            y = y.to(device, non_blocking=True)
            out = model(input_ids=x, labels=y)
            loss = out.loss / GRAD_ACCUM
            loss.backward()
            total_loss += loss.item()

        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        optimizer.step()
        scheduler.step()
        tokens_seen += MICRO_BATCH * GRAD_ACCUM * SEQ_LEN

        if (step + 1) % LOG_EVERY == 0:
            now = time.time()
            dt = now - last_log_time
            dtok = tokens_seen - last_log_tokens
            tps = dtok / max(dt, 1)
            lr = optimizer.param_groups[0]["lr"]
            vram = torch.cuda.memory_allocated() / 1e9 if torch.cuda.is_available() else 0
            log(f"step {step+1:>6} | loss {total_loss:.4f} | lr {lr:.2e} | "
                f"tok/s {int(tps)} | VRAM {vram:.1f}GB | tokens {tokens_seen:,}")
            last_log_tokens = tokens_seen
            last_log_time = now

        if (step + 1) % CELL_MEASURE_EVERY == 0:
            try:
                summary = cell_sys.maybe_step(model, step + 1, dataset=dataset)
                if summary is not None:
                    log(f"[cells] step {step+1}: active={summary['n_cells_active']}/{summary['n_cells_total']} "
                        f"splits_now={len(summary['splits'])}")
            except Exception as e:
                log(f"[cells] ERROR at step {step+1}: {type(e).__name__}: {e}")

        if (step + 1) % SAVE_EVERY == 0:
            p = save_checkpoint(model, optimizer, scheduler, step + 1,
                                 tokens_seen, total_loss)
            log(f"Checkpoint saved: {p.name}")
            if cell_sys.manager:
                cell_sys.manager.save()
                try:
                    cell_sys.save_narrative()
                except Exception as e:
                    log(f"[narrative] error: {e}")
            cleanup_old_checkpoints()

    log("Training complete!")


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--resume", action="store_true")
    args = ap.parse_args()
    train(args)


if __name__ == "__main__":
    main()