File size: 5,995 Bytes
4e70c09 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 | 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
|