#!/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()