from __future__ import annotations import argparse import json import math import os from pathlib import Path import subprocess import time from typing import Any, Optional import torch from torch.utils.data import DataLoader, Subset from conveyer.checkpoint import load_checkpoint_payload, save_checkpoint from conveyer.config import ConveyerConfig, load_config from conveyer.data import TokenizerProtocol, build_dataset_and_tokenizer, collate_conveyer_batch from conveyer.infer import generate_conveyer from conveyer.losses import compute_unrolled_conveyer_loss from conveyer.model import ConveyerDiffusionLM def _query_gpu_temp_c() -> Optional[int]: try: out = subprocess.check_output( ["nvidia-smi", "--query-gpu=temperature.gpu", "--format=csv,noheader,nounits"], timeout=2.0, stderr=subprocess.DEVNULL, ) return int(out.decode("utf-8", errors="ignore").strip().splitlines()[0]) except Exception: return None def get_lr(step: int, max_steps: int, base_lr: float, warmup_steps: int) -> float: if step < warmup_steps: return base_lr * float(step + 1) / float(max(1, warmup_steps)) progress = float(step - warmup_steps) / float(max(1, max_steps - warmup_steps)) progress = min(max(progress, 0.0), 1.0) return base_lr * (0.1 + 0.9 * 0.5 * (1.0 + math.cos(math.pi * progress))) @torch.no_grad() def evaluate_validation( model: ConveyerDiffusionLM, val_loader: DataLoader[dict[str, torch.Tensor]], cfg: ConveyerConfig, device: torch.device, pad_id: int = 0, max_batches: int = 25, ) -> dict[str, Any]: model.eval() B = model.belt_size loss_sum = 0.0 s0_loss_sum = 0.0 s0_acc_sum = 0.0 lm_loss_sum = 0.0 lm_acc_sum = 0.0 belt_acc_sum = 0.0 per_slot_acc_sum = torch.zeros(B, device=device) per_slot_ent_sum = torch.zeros(B, device=device) count = 0 for i, batch in enumerate(val_loader): if i >= max_batches: break seq = batch["input_ids"].to(device) plen = batch["prompt_lengths"].to(device) t_eval = max(1, int(plen.min().item()) - min(4, cfg.train.unroll_steps // 2)) use_amp = device.type == "cuda" and torch.cuda.is_bf16_supported() with torch.autocast(device_type=device.type, dtype=torch.bfloat16 if use_amp else torch.float32, enabled=use_amp): m = compute_unrolled_conveyer_loss( model=model, seq_tokens=seq, cfg=cfg, prompt_lengths=plen, tick_start=t_eval, pad_token_id=pad_id, ) loss_sum += float(m.loss.item()) s0_loss_sum += m.slot0_loss s0_acc_sum += m.slot0_acc lm_loss_sum += m.prefix_lm_loss lm_acc_sum += m.prefix_lm_acc belt_acc_sum += m.belt_mean_acc per_slot_acc_sum += m.per_slot_acc per_slot_ent_sum += m.per_slot_entropy count += 1 model.train() denom = max(count, 1) s0_ce = s0_loss_sum / denom lm_ce = lm_loss_sum / denom return { "val_loss": loss_sum / denom, "val_slot0_ce": s0_ce, "val_slot0_ppl": math.exp(min(s0_ce, 20.0)), "val_slot0_acc": s0_acc_sum / denom, "val_lm_ce": lm_ce, "val_lm_ppl": math.exp(min(lm_ce, 20.0)), "val_lm_acc": lm_acc_sum / denom, "val_belt_acc": belt_acc_sum / denom, "val_per_slot_acc": (per_slot_acc_sum / denom).cpu().tolist(), "val_per_slot_entropy": (per_slot_ent_sum / denom).cpu().tolist(), } def _prune_old_milestones(out_dir: Path, keep_last: int = 8) -> None: milestones = sorted(out_dir.glob("ckpt_step_*.pt"), key=lambda p: p.stat().st_mtime) if len(milestones) > keep_last: for old_p in milestones[: len(milestones) - keep_last]: try: old_p.unlink() prev_p = old_p.with_suffix(old_p.suffix + ".prev") if prev_p.exists(): prev_p.unlink() except OSError: pass @torch.no_grad() def _sample_preview( model: ConveyerDiffusionLM, tokenizer: TokenizerProtocol, device: torch.device, prompt_text: str, max_new_tokens: int = 32, ) -> str: model.eval() p_ids = torch.tensor([tokenizer.encode(prompt_text.strip())], dtype=torch.long, device=device) out_ids, _ = generate_conveyer( model=model, prompt_tokens=p_ids, max_new_tokens=max_new_tokens, prompt_overlap=10, temperature=0.0, eos_token_id=tokenizer.eos_id, ) gen_part = tokenizer.decode(out_ids[0, p_ids.size(1) :].tolist()) model.train() return gen_part.replace("\n", " ").strip() def _sanitize_jsonl(log_path: Path) -> None: """Strips any trailing NUL bytes left by an abrupt OS crash so the log remains valid JSONL.""" if not log_path.exists(): return raw = log_path.read_bytes().replace(b"\x00", b"") lines = [ln for ln in raw.splitlines() if ln.strip().startswith(b"{") and ln.strip().endswith(b"}")] log_path.write_bytes(b"\n".join(lines) + (b"\n" if lines else b"")) def train_model( cfg: ConveyerConfig, init_ckpt: str | Path | None = None, resume_ckpt: str | Path | None = None, ) -> dict[str, Any]: torch.manual_seed(cfg.train.seed) device = torch.device("cuda" if torch.cuda.is_available() else "cpu") train_ds, val_ds, tokenizer = build_dataset_and_tokenizer(cfg) if cfg.train.overfit_batches > 0: n_sub = min(len(train_ds), cfg.train.overfit_batches * cfg.train.batch_size) train_ds = Subset(train_ds, list(range(n_sub))) val_ds = train_ds collate_fn = lambda b: collate_conveyer_batch(b, pad_id=tokenizer.pad_id) train_loader = DataLoader( train_ds, batch_size=cfg.train.batch_size, shuffle=(cfg.train.overfit_batches == 0), collate_fn=collate_fn, drop_last=False, ) val_loader = DataLoader( val_ds, batch_size=cfg.train.batch_size, shuffle=False, collate_fn=collate_fn, drop_last=False, ) model = ConveyerDiffusionLM(cfg.model).to(device) optimizer = torch.optim.AdamW( model.parameters(), lr=cfg.train.lr, weight_decay=cfg.train.weight_decay, betas=(0.9, 0.95), ) out_dir = Path(cfg.train.out_dir) out_dir.mkdir(parents=True, exist_ok=True) log_jsonl = out_dir / "train_log.jsonl" _sanitize_jsonl(log_jsonl) start_step = 1 best_val_loss = float("inf") prior_elapsed = 0.0 # Resolve resume vs warm-start initialization (with automatic fallback to .prev or ckpt_step_*.pt) resume_target: Optional[Path] = None if resume_ckpt is not None: if str(resume_ckpt).lower() == "auto": resume_target = out_dir / "last.pt" else: resume_target = Path(resume_ckpt) if resume_target is not None: try: ckpt_data, used_path = load_checkpoint_payload(resume_target, map_location="cpu") model.load_state_dict(ckpt_data["model_state"], strict=False) if "optimizer_state" in ckpt_data: try: optimizer.load_state_dict(ckpt_data["optimizer_state"]) except Exception as e: print(f"[ConveyerTrain] Warning: could not restore optimizer state ({e})", flush=True) start_step = int(ckpt_data.get("step", 0)) + 1 best_val_loss = float( ckpt_data.get("best_val_loss") or ckpt_data.get("metrics", {}).get("val_loss", float("inf")) ) prior_elapsed = float(ckpt_data.get("elapsed_sec", 0.0)) print( f"[ConveyerTrain] Resumed from {used_path} at step {start_step} " f"(best_val_loss={best_val_loss:.4f}, prior_elapsed={prior_elapsed/3600:.2f}h)", flush=True, ) del ckpt_data if device.type == "cuda": torch.cuda.empty_cache() except Exception as e: print(f"[ConveyerTrain] No readable resume checkpoint in {out_dir} ({e}); starting from step 1.", flush=True) elif init_ckpt is not None and Path(init_ckpt).exists(): ckpt_data, used_path = load_checkpoint_payload(init_ckpt, map_location="cpu") model.load_state_dict(ckpt_data["model_state"], strict=False) print(f"[ConveyerTrain] Warm-started weights from {used_path}", flush=True) del ckpt_data use_amp = device.type == "cuda" and torch.cuda.is_bf16_supported() amp_dtype = torch.bfloat16 if use_amp else torch.float32 mode_desc = ( "1-Pass AR Baseline" if cfg.model.causal_belt_only else ( "64-Slot Conveyor (Soft-Vocab Only)" if not cfg.model.use_thought_highway else "64-Slot Dual-Channel Conveyor (Soft-Vocab + Gated Thought)" ) ) max_hours = float(getattr(cfg.train, "max_hours", 0.0)) save_every = max(1, int(getattr(cfg.train, "save_every", 500))) milestone_every = save_every * 2 print( f"[ConveyerTrain] Mode: {mode_desc} | Params: {model.num_parameters()/1e6:.2f}M | " f"Device: {device} | Train stories: {len(train_ds):,} | Val stories: {len(val_ds):,} | " f"Unroll U={cfg.train.unroll_steps} | Belt B={cfg.model.belt_size} | " f"max_steps={cfg.train.max_steps:,} | max_hours={max_hours:.1f}h | save_every={save_every}", flush=True, ) train_iter = iter(train_loader) last_val_metrics: dict[str, Any] = {} t0 = time.time() model.train() for step in range(start_step, cfg.train.max_steps + 1): try: batch = next(train_iter) except StopIteration: train_iter = iter(train_loader) batch = next(train_iter) lr = get_lr(step - 1, cfg.train.max_steps, cfg.train.lr, cfg.train.warmup_steps) for pg in optimizer.param_groups: pg["lr"] = lr seq = batch["input_ids"].to(device) plen = batch["prompt_lengths"].to(device) optimizer.zero_grad(set_to_none=True) with torch.autocast(device_type=device.type, dtype=amp_dtype, enabled=use_amp): metrics = compute_unrolled_conveyer_loss( model=model, seq_tokens=seq, cfg=cfg, prompt_lengths=plen, pad_token_id=tokenizer.pad_id, ) metrics.loss.backward() if cfg.train.grad_clip > 0.0: torch.nn.utils.clip_grad_norm_(model.parameters(), cfg.train.grad_clip) optimizer.step() cooldown_ms = int(getattr(cfg.train, "step_cooldown_ms", 0)) if cooldown_ms > 0: if device.type == "cuda": torch.cuda.synchronize() time.sleep(cooldown_ms / 1000.0) total_elapsed = prior_elapsed + (time.time() - t0) if step % cfg.train.log_every == 0 or step == start_step: acc_back = float(metrics.per_slot_acc[-1].item()) acc_mid = float(metrics.per_slot_acc[model.belt_size // 2].item()) steps_done = max(1, step - start_step + 1) sec_per_step = (time.time() - t0) / steps_done gpu_temp = _query_gpu_temp_c() if device.type == "cuda" else None temp_str = f" | {gpu_temp}C" if gpu_temp is not None else "" print( f"step {step:6d}/{cfg.train.max_steps} | loss {metrics.loss.item():.4f} | " f"lm_ce {metrics.prefix_lm_loss:.3f} (acc {metrics.prefix_lm_acc:.3f}) | " f"s0_ce {metrics.slot0_loss:.3f} (acc {metrics.slot0_acc:.3f}) | " f"belt(0/m/63) {metrics.slot0_acc:.2f}/{acc_mid:.2f}/{acc_back:.2f} | " f"pause {metrics.mean_plan_pause*100:.0f}% | lr {lr:.2e} | " f"{sec_per_step:.2f}s/it{temp_str} | {total_elapsed/3600:.2f}h", flush=True, ) max_temp = int(getattr(cfg.train, "max_gpu_temp_c", 80)) if gpu_temp is not None and max_temp > 0 and gpu_temp >= max_temp: torch.cuda.synchronize() while True: print( f" [THERMAL GUARD] GPU at {gpu_temp}C (>= {max_temp}C limit); pausing 8s to let fans cool GPU...", flush=True, ) time.sleep(8.0) gpu_temp = _query_gpu_temp_c() if gpu_temp is None or gpu_temp <= max_temp - 5: print(f" [THERMAL GUARD] GPU cooled to {gpu_temp}C; resuming training.", flush=True) break is_time_up = max_hours > 0.0 and (total_elapsed >= max_hours * 3600.0) is_eval_step = (step % cfg.train.eval_every == 0) or (step == cfg.train.max_steps) or is_time_up is_save_step = (step % save_every == 0) or is_eval_step if is_eval_step: if device.type == "cuda": torch.cuda.empty_cache() val_m = evaluate_validation( model=model, val_loader=val_loader, cfg=cfg, device=device, pad_id=tokenizer.pad_id, ) last_val_metrics = val_m curve = val_m["val_per_slot_acc"] B = len(curve) print( f" [VAL step {step:6d} | {total_elapsed/3600:.2f}h] " f"val_loss {val_m['val_loss']:.4f} | " f"real_s0_ppl {val_m['val_slot0_ppl']:.2f} | real_s0_acc {val_m['val_slot0_acc']:.3f} | " f"real_lm_ppl {val_m['val_lm_ppl']:.2f} | real_lm_acc {val_m['val_lm_acc']:.3f} | " f"crystallization(63->32->8->0): " f"{curve[B-1]:.2f}->{curve[B//2]:.2f}->{curve[min(8, B-1)]:.2f}->{curve[0]:.2f}", flush=True, ) if cfg.data.dataset == "tinystories": sample_p = ( "Once upon a time, Lily and a little bunny walked into the green forest " "to find the golden key near the locked treasure chest. Suddenly a big lion" ) preview = _sample_preview(model, tokenizer, device, sample_p, max_new_tokens=28) print(f" [SAMPLE] ...Suddenly a big lion -> {preview}", flush=True) else: preview = "" if val_m["val_loss"] <= best_val_loss: best_val_loss = val_m["val_loss"] save_checkpoint( out_dir / "best.pt", model, cfg, step=step, optimizer=None, # Keep best.pt lightweight (~115 MB) for fast inference & fewer disk writes metrics=val_m, best_val_loss=best_val_loss, elapsed_sec=total_elapsed, ) log_entry = { "step": step, "elapsed_hours": round(total_elapsed / 3600.0, 4), "val_loss": round(val_m["val_loss"], 4), "val_slot0_ce": round(val_m["val_slot0_ce"], 4), "val_slot0_ppl": round(val_m["val_slot0_ppl"], 2), "val_slot0_acc": round(val_m["val_slot0_acc"], 4), "val_lm_ce": round(val_m["val_lm_ce"], 4), "val_lm_ppl": round(val_m["val_lm_ppl"], 2), "val_lm_acc": round(val_m["val_lm_acc"], 4), "val_belt_acc": round(val_m["val_belt_acc"], 4), "sample_lion": preview, } with open(log_jsonl, "a", encoding="utf-8") as lf: lf.write(json.dumps(log_entry) + "\n") lf.flush() os.fsync(lf.fileno()) if is_save_step: save_checkpoint( out_dir / "last.pt", model, cfg, step=step, optimizer=optimizer, metrics=last_val_metrics, best_val_loss=best_val_loss, elapsed_sec=total_elapsed, ) if step % milestone_every == 0: save_checkpoint( out_dir / f"ckpt_step_{step:06d}.pt", model, cfg, step=step, optimizer=optimizer, metrics=last_val_metrics, best_val_loss=best_val_loss, elapsed_sec=total_elapsed, ) _prune_old_milestones(out_dir, keep_last=8) if device.type == "cuda": torch.cuda.empty_cache() if is_time_up: print( f"[ConveyerTrain] Reached max_hours={max_hours:.2f}h budget at step {step}. " f"Checkpoints saved to {out_dir}.", flush=True, ) break return last_val_metrics def main() -> None: parser = argparse.ArgumentParser(description="Train Conveyor-Belt Diffusion LLM") parser.add_argument("--config", type=str, default="configs/synthetic.yaml", help="Path to YAML config") parser.add_argument("--init", type=str, default=None, help="Optional checkpoint (.pt) to warm-start weights from") parser.add_argument( "--resume", type=str, default=None, help="Checkpoint path (.pt) or 'auto' to resume full training state (model + optimizer + step) after a crash", ) parser.add_argument("--max_hours", type=float, default=None, help="Wall-clock training budget in hours (e.g. 24)") parser.add_argument("--overfit", type=int, default=0, help="Overfit on N fixed batches for verification") parser.add_argument("--no_thought_highway", action="store_true", help="Disable continuous thought highway (alpha_k = 0)") parser.add_argument("--no_belt_bptt", action="store_true", help="Detach conveyor belt state between unrolled ticks (W=1)") parser.add_argument("--ar_baseline", action="store_true", help="Train 1-pass causal Autoregressive baseline") parser.add_argument("overrides", nargs="*", help="Config overrides in section.key=value format") args = parser.parse_args() cfg = load_config(args.config, args.overrides) if args.max_hours is not None: cfg.train.max_hours = args.max_hours if args.overfit > 0: cfg.train.overfit_batches = args.overfit if args.no_thought_highway: cfg.train.no_thought_highway = True cfg.model.use_thought_highway = False if args.no_belt_bptt: cfg.train.no_belt_bptt = True if args.ar_baseline: cfg.train.ar_baseline = True cfg.model.causal_belt_only = True train_model(cfg, init_ckpt=args.init, resume_ckpt=args.resume) if __name__ == "__main__": main()