| from __future__ import annotations | |
| from dataclasses import asdict | |
| import os | |
| from pathlib import Path | |
| from typing import Any, Optional | |
| import urllib.request | |
| import torch | |
| from conveyer.config import ConveyerConfig, _merge_dataclass | |
| from conveyer.model import ConveyerDiffusionLM | |
| HF_DEFAULT_CKPT_URL = "https://huggingface.co/sumitk/diffusion-conveyer-llm/resolve/main/best.pt" | |
| def _move_to_cpu(obj: Any) -> Any: | |
| if isinstance(obj, torch.Tensor): | |
| return obj.detach().to(device="cpu", non_blocking=False) | |
| if isinstance(obj, dict): | |
| return {k: _move_to_cpu(v) for k, v in obj.items()} | |
| if isinstance(obj, list): | |
| return [_move_to_cpu(v) for v in obj] | |
| if isinstance(obj, tuple): | |
| return tuple(_move_to_cpu(v) for v in obj) | |
| return obj | |
| def _atomic_write_payload(path: Path, payload: dict[str, Any]) -> None: | |
| """ | |
| Crash-safe atomic checkpoint writer on Windows NTFS: | |
| 1. If `path` already exists, rotates `path -> path.prev` BEFORE writing so an abrupt BSOD/power loss | |
| never leaves both the current and previous checkpoint with un-flushed NTFS data blocks. | |
| 2. Writes `payload` to `path.tmp`, calls `f.flush()` + `os.fsync(f.fileno())` to force NTFS data blocks | |
| to physical NVMe storage BEFORE updating the directory metadata via `os.replace`. | |
| """ | |
| path.parent.mkdir(parents=True, exist_ok=True) | |
| tmp_p = path.with_suffix(path.suffix + ".tmp") | |
| prev_p = path.with_suffix(path.suffix + ".prev") | |
| with open(tmp_p, "wb") as f: | |
| torch.save(payload, f) | |
| f.flush() | |
| os.fsync(f.fileno()) | |
| if path.exists(): | |
| try: | |
| os.replace(path, prev_p) | |
| except OSError: | |
| pass | |
| os.replace(tmp_p, path) | |
| def save_checkpoint( | |
| path: str | Path, | |
| model: ConveyerDiffusionLM, | |
| cfg: ConveyerConfig, | |
| step: int, | |
| optimizer: Optional[torch.optim.Optimizer] = None, | |
| metrics: Optional[dict[str, Any]] = None, | |
| best_val_loss: Optional[float] = None, | |
| elapsed_sec: float = 0.0, | |
| ) -> None: | |
| """ | |
| Copies `model` and `optimizer` state dicts to CPU RAM first (avoiding 350 MB of synchronous | |
| WDDM GPU->CPU DMA transfers inside `torch.save` that can trigger `dxgmms2.sys` BSODs), | |
| then writes atomically with `os.fsync` and `.prev` backup rotation. | |
| """ | |
| if torch.cuda.is_available(): | |
| torch.cuda.synchronize() | |
| cpu_model_state = _move_to_cpu(model.state_dict()) | |
| payload: dict[str, Any] = { | |
| "model_state": cpu_model_state, | |
| "config": asdict(cfg), | |
| "step": step, | |
| "metrics": metrics or {}, | |
| "best_val_loss": best_val_loss, | |
| "elapsed_sec": elapsed_sec, | |
| } | |
| if optimizer is not None: | |
| payload["optimizer_state"] = _move_to_cpu(optimizer.state_dict()) | |
| _atomic_write_payload(Path(path), payload) | |
| def load_checkpoint_payload(path: str | Path, map_location: Any = "cpu") -> tuple[dict[str, Any], Path]: | |
| """ | |
| Loads checkpoint dict from `path`, automatically falling back to `path.prev` or the latest | |
| valid `ckpt_step_*.pt` in the same directory if `path` was corrupted by a power loss or BSOD. | |
| """ | |
| p = Path(path) | |
| candidates: list[Path] = [p, p.with_suffix(p.suffix + ".prev")] | |
| if p.parent.exists(): | |
| milestones = sorted(p.parent.glob("ckpt_step_*.pt"), key=lambda x: x.stat().st_mtime, reverse=True) | |
| for m in milestones: | |
| if m not in candidates: | |
| candidates.append(m) | |
| m_prev = m.with_suffix(m.suffix + ".prev") | |
| if m_prev.exists() and m_prev not in candidates: | |
| candidates.append(m_prev) | |
| last_err: Optional[Exception] = None | |
| for cand in candidates: | |
| if not cand.exists() or cand.stat().st_size == 0: | |
| continue | |
| try: | |
| data = torch.load(cand, map_location=map_location, weights_only=False) | |
| if isinstance(data, dict) and "model_state" in data: | |
| return data, cand | |
| except Exception as e: | |
| last_err = e | |
| print(f"[Checkpoint] Warning: skipping unreadable checkpoint {cand.name} ({e})", flush=True) | |
| if p.name == "best.pt" and "tinystories" in str(p).replace("\\", "/"): | |
| try: | |
| p.parent.mkdir(parents=True, exist_ok=True) | |
| tmp_dl = p.with_suffix(".pt.dl") | |
| print(f"[Checkpoint] Downloading pretrained checkpoint from {HF_DEFAULT_CKPT_URL} ...", flush=True) | |
| req = urllib.request.Request(HF_DEFAULT_CKPT_URL, headers={"User-Agent": "Mozilla/5.0"}) | |
| with urllib.request.urlopen(req, timeout=120) as resp, open(tmp_dl, "wb") as f: | |
| while True: | |
| chunk = resp.read(4 * 1024 * 1024) | |
| if not chunk: | |
| break | |
| f.write(chunk) | |
| os.replace(tmp_dl, p) | |
| data = torch.load(p, map_location=map_location, weights_only=False) | |
| if isinstance(data, dict) and "model_state" in data: | |
| return data, p | |
| except Exception as e: | |
| last_err = e | |
| raise RuntimeError(f"No valid checkpoint found for {p} (last error: {last_err})") | |
| def load_checkpoint( | |
| path: str | Path, | |
| device: Optional[torch.device] = None, | |
| ) -> tuple[ConveyerDiffusionLM, ConveyerConfig, dict[str, Any]]: | |
| if device is None: | |
| device = torch.device("cuda" if torch.cuda.is_available() else "cpu") | |
| ckpt, used_path = load_checkpoint_payload(path, map_location=device) | |
| if used_path != Path(path): | |
| print(f"[Checkpoint] Recovered from fallback checkpoint {used_path}", flush=True) | |
| cfg = ConveyerConfig() | |
| if "config" in ckpt: | |
| _merge_dataclass(cfg, ckpt["config"]) | |
| model = ConveyerDiffusionLM(cfg.model).to(device) | |
| model.load_state_dict(ckpt["model_state"]) | |
| return model, cfg, ckpt | |