| 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
|
|
|
|
|
| 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,
|
| 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()
|
|
|
