| """Resumable GPU training with fixed validation masks and atomic recovery state.""" |
| from contextlib import nullcontext |
| from dataclasses import asdict |
| import json |
| import math |
| from pathlib import Path |
| import shutil |
| import time |
|
|
| import numpy as np |
| import torch |
| from torch.nn import functional as F |
| from .data import load_arrays |
| from .model import Config, Denoiser, PAD, BOS, corrupt |
|
|
|
|
| def amp(device, precision): |
| return torch.autocast("cuda", dtype=torch.bfloat16) if device == "cuda" and precision == "bf16" else nullcontext() |
|
|
|
|
| def atomic_save(value, path): |
| temporary = path.with_suffix(".tmp") |
| torch.save(value, temporary) |
| temporary.replace(path) |
|
|
|
|
| @torch.inference_mode() |
| def evaluate(model, array, frequencies, *, device, precision, limit=256, batch_size=8): |
| model.eval() |
| reports = [] |
| for rate in (0.15, 0.5, 0.85): |
| rng = torch.Generator(device=device).manual_seed(2718) |
| total = correct = top5 = 0 |
| loss = baseline_loss = 0. |
| baseline_correct = 0 |
| for start in range(0, min(limit, len(array)), batch_size): |
| clean = torch.tensor(np.array(array[start:min(start+batch_size,limit)]), dtype=torch.long, device=device) |
| noisy, selected, noise = corrupt(clean, rng, mode="random", rate=rate) |
| with amp(device, precision): |
| logits = model(noisy, noise, selected) |
| targets = clean[selected] |
| logits = logits.float() |
| loss += F.cross_entropy(logits, targets, reduction="sum").item() |
| correct += logits.argmax(-1).eq(targets).sum().item() |
| top5 += logits.topk(5, -1).indices.eq(targets[:, None]).any(-1).sum().item() |
| total += targets.numel() |
| baseline_loss += -frequencies.to(device)[targets].log().sum().item() |
| baseline_correct += targets.eq(int(frequencies.argmax())).sum().item() |
| reports.append({"mask_rate": rate, "targets": total, "cross_entropy": loss / total, |
| "accuracy": correct / total, "top5_accuracy": top5 / total, |
| "unigram_cross_entropy": baseline_loss / total, "unigram_accuracy": baseline_correct / total}) |
| return reports |
|
|
|
|
| def train(args): |
| if min(args.steps, args.batch_size, args.accumulation, args.eval_every, args.threads, args.eval_samples) < 1: |
| raise ValueError("Training counts must be positive") |
| torch.set_num_threads(args.threads) |
| torch.manual_seed(args.seed) |
| device = args.device |
| if device == "cuda" and not torch.cuda.is_available(): |
| raise ValueError("CUDA/ROCm unavailable; choose --device cpu") |
| if args.precision == "bf16" and device == "cuda" and not torch.cuda.is_bf16_supported(): |
| raise ValueError("BF16 unsupported; choose --precision fp32") |
| arrays, manifest = load_arrays(args.data) |
| config = Config(vocab_size=manifest["vocab_size"], width=args.width, heads=args.heads, layers=args.layers, |
| max_length=manifest["max_length"]) |
| out = Path(args.output) |
| out.mkdir(parents=True, exist_ok=True) |
| if (out / "last.pt").exists() and not args.resume: |
| raise ValueError("Run already exists; use --resume or a new output directory") |
| run_config = {k: v for k, v in vars(args).items() if k not in ("resume", "command")} |
| model = Denoiser(config).to(device) |
| lineage = None |
| if getattr(args, "initialize_from", None): |
| if args.resume: |
| raise ValueError("Choose resume OR initialize-from") |
| source = Path(args.initialize_from) |
| lineage = json.loads((source / "training_info.json").read_text()) |
| if lineage["manifest"] != manifest or lineage.get("pretrained_weights") is not False: |
| raise ValueError("Initialization must use this corpus and a from-scratch DiffuThink checkpoint") |
| prior = Denoiser.from_pretrained(source) |
| if prior.config != config: |
| raise ValueError("Initialization model configuration differs") |
| model.load_state_dict(prior.state_dict()) |
| del prior |
| optimizer = torch.optim.AdamW(model.parameters(), lr=args.lr, betas=(0.9, 0.95), weight_decay=0.1) |
| sampling = torch.Generator().manual_seed(args.seed) |
| noise_rng = torch.Generator(device=device).manual_seed(args.seed + 1) |
| counts = np.bincount(np.asarray(arrays["train"]).reshape(-1), minlength=config.vocab_size).astype(np.float64) + 1 |
| counts[[PAD, BOS]] = 1 |
| frequencies = torch.tensor(counts / counts.sum(), dtype=torch.float32) |
| step, seen_tokens, elapsed, best = 0, 0, 0., float("inf") |
| if args.resume: |
| if not (out / "best" / "model.safetensors").exists(): |
| raise ValueError("Resume in the original output directory containing best/ and last.pt") |
| state = torch.load(args.resume, map_location="cpu", weights_only=True) |
| if state["manifest"] != manifest or state["model_config"] != asdict(config): |
| raise ValueError("Resume data/model mismatch") |
| for key in ("seed", "batch_size", "accumulation", "lr", "steps", "warmup", "precision", "device"): |
| if state["run_config"][key] != run_config[key]: |
| raise ValueError(f"Resume requires unchanged {key}; schedule must remain identical") |
| model.load_state_dict(state["model"]) |
| optimizer.load_state_dict(state["optimizer"]) |
| sampling.set_state(state["sampling_rng"]) |
| noise_rng.set_state(state["noise_rng"]) |
| torch.set_rng_state(state["torch_rng"]) |
| if device == "cuda": |
| torch.cuda.set_rng_state(state["cuda_rng"], device) |
| step, seen_tokens, elapsed, best = state["step"], state["seen_tokens"], state["elapsed"], state["best"] |
| lineage = state.get("lineage") |
| del state |
| else: |
| initial = evaluate(model, arrays["validation"], frequencies, device=device, precision=args.precision, limit=args.eval_samples) |
| (out / "initial_validation.json").write_text(json.dumps(initial, indent=2), encoding="utf-8") |
| (out / "run_config.json").write_text(json.dumps(run_config, indent=2), encoding="utf-8") |
| (out / "manifest.json").write_text(json.dumps(manifest, indent=2), encoding="utf-8") |
| start = time.perf_counter() |
| parameters = sum(p.numel() for p in model.parameters()) |
| print(json.dumps({"parameters": parameters, "device": device, "precision": args.precision, "start_step": step, "target_steps": args.steps}), flush=True) |
| log = out / "metrics.jsonl" |
| def save_state(): |
| atomic_save({"model": model.state_dict(), "optimizer": optimizer.state_dict(), |
| "model_config": asdict(config), "manifest": manifest, "run_config": run_config, |
| "step": step, "seen_tokens": seen_tokens, "elapsed": elapsed + time.perf_counter() - start, |
| "best": best, "sampling_rng": sampling.get_state(), "noise_rng": noise_rng.get_state(), |
| "lineage": lineage, |
| "torch_rng": torch.get_rng_state(), |
| "cuda_rng": torch.cuda.get_rng_state() if device == "cuda" else None}, out / "last.pt") |
| try: |
| while step < args.steps: |
| model.train() |
| lr_scale = min(1., (step + 1) / max(1, args.warmup)) |
| progress = max(0., (step - args.warmup) / max(1, args.steps - args.warmup)) |
| lr = args.lr * lr_scale * (0.1 + 0.9 * 0.5 * (1 + math.cos(math.pi * progress))) |
| for group in optimizer.param_groups: |
| group["lr"] = lr |
| optimizer.zero_grad(set_to_none=True) |
| loss_value = 0. |
| for _ in range(args.accumulation): |
| indices = torch.randint(len(arrays["train"]), (args.batch_size,), generator=sampling).numpy() |
| clean = torch.tensor(np.array(arrays["train"][indices]), dtype=torch.long, device=device) |
| noisy, selected, noise = corrupt(clean, noise_rng) |
| with amp(device, args.precision): |
| logits = model(noisy, noise, selected) |
| loss = F.cross_entropy(logits.float(), clean[selected]) |
| if not torch.isfinite(loss): |
| raise RuntimeError("Non-finite loss; previous recovery checkpoint is preserved") |
| (loss / args.accumulation).backward() |
| loss_value += loss.item() / args.accumulation |
| seen_tokens += int(clean.ne(PAD).sum()) |
| gradient = torch.nn.utils.clip_grad_norm_(model.parameters(), 1.) |
| if not torch.isfinite(gradient): |
| raise RuntimeError("Non-finite gradient; previous checkpoint is preserved") |
| optimizer.step() |
| step += 1 |
| if step == 1 or step % 25 == 0: |
| wall = elapsed + time.perf_counter() - start |
| row = {"step": step, "loss": loss_value, "lr": lr, "gradient_norm": float(gradient), |
| "seen_tokens": seen_tokens, "seconds": wall, "tokens_per_second": seen_tokens / wall} |
| with log.open("a", encoding="utf-8") as f: |
| f.write(json.dumps(row) + "\n") |
| print(json.dumps(row), flush=True) |
| if step % args.eval_every == 0 or step == args.steps: |
| metrics = evaluate(model, arrays["validation"], frequencies, device=device, precision=args.precision, limit=args.eval_samples) |
| score = sum(r["cross_entropy"] for r in metrics) / len(metrics) |
| with log.open("a", encoding="utf-8") as f: |
| f.write(json.dumps({"step": step, "validation": metrics, "validation_loss": score}) + "\n") |
| if score < best: |
| best = score |
| model.save_pretrained(out / "best") |
| shutil.copy2(Path(args.data) / "tokenizer.json", out / "best" / "tokenizer.json") |
| (out / "best" / "training_info.json").write_text(json.dumps({"step": step, "validation_loss": score, |
| "parameters": parameters, "seen_tokens": seen_tokens, "pretrained_weights": False, |
| "manifest": manifest, "previous_phase": lineage}, indent=2), encoding="utf-8") |
| save_state() |
| print(json.dumps({"step": step, "validation_loss": score, "best": best}), flush=True) |
| except KeyboardInterrupt: |
| |
| print("Interrupted. Resume from the previous completed evaluation checkpoint.", flush=True) |
| return |
| model = Denoiser.from_pretrained(out / "best", device) |
| report = {"status": "completed", "model_config": asdict(config), "parameters": parameters, |
| "training": run_config, "seen_tokens": seen_tokens, "seconds": elapsed + time.perf_counter() - start, |
| "torch_version": str(torch.__version__), "device_name": torch.cuda.get_device_name() if device == "cuda" else "cpu", |
| "best": json.loads((out / "best" / "training_info.json").read_text()), |
| "test": evaluate(model, arrays["test"], frequencies, device=device, precision=args.precision, limit=args.test_samples), |
| "limitations": ["English synthetic stories, not a general chatbot", "No comparison to v1 byte accuracy: tokenizer and corpus differ", |
| "Custom masked-denoising objective, not a claimed diffusion likelihood bound"]} |
| (out / "report.json").write_text(json.dumps(report, indent=2), encoding="utf-8") |
| print(f"Training complete: {out / 'report.json'}", flush=True) |
|
|
