File size: 3,041 Bytes
d5e0d8f | 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 | """Atomic Predictor-v4 inference and resumable training checkpoints."""
from __future__ import annotations
import json
import os
import random
from pathlib import Path
from typing import Any
import torch
from safetensors.torch import load_file, save_file
def unwrap_model(model: torch.nn.Module) -> torch.nn.Module:
return model.module if hasattr(model, "module") else model
def trainable_state_dict(
model: torch.nn.Module,
*,
floating_dtype: torch.dtype | None = None,
) -> dict[str, torch.Tensor]:
model = unwrap_model(model)
trainable_names = {
name for name, parameter in model.named_parameters() if parameter.requires_grad
}
return {
name: tensor.detach()
.to(
device="cpu",
dtype=(
floating_dtype
if floating_dtype is not None and tensor.is_floating_point()
else tensor.dtype
),
)
.contiguous()
for name, tensor in model.state_dict().items()
if name in trainable_names
}
def save_predictor_weights(
model: torch.nn.Module,
path: str | Path,
*,
metadata: dict[str, Any],
floating_dtype: torch.dtype = torch.bfloat16,
) -> Path:
path = Path(path)
path.parent.mkdir(parents=True, exist_ok=True)
temporary = path.with_suffix(path.suffix + f".tmp.{os.getpid()}")
save_file(
trainable_state_dict(model, floating_dtype=floating_dtype),
str(temporary),
metadata={
"format": "self_forcing_predictor_v4",
"config": json.dumps(metadata, ensure_ascii=False, sort_keys=True),
},
)
os.replace(temporary, path)
return path
def load_predictor_weights(model: torch.nn.Module, path: str | Path) -> None:
state = load_file(str(path), device="cpu")
result = unwrap_model(model).load_state_dict(state, strict=False)
trainable = {
name
for name, parameter in unwrap_model(model).named_parameters()
if parameter.requires_grad
}
missing_trainable = sorted(trainable.intersection(result.missing_keys))
if result.unexpected_keys or missing_trainable:
raise RuntimeError(
"Predictor weight mismatch: "
f"unexpected={result.unexpected_keys}, "
f"missing_trainable={missing_trainable}"
)
def atomic_torch_save(payload: dict[str, Any], path: str | Path) -> Path:
path = Path(path)
path.parent.mkdir(parents=True, exist_ok=True)
temporary = path.with_suffix(path.suffix + f".tmp.{os.getpid()}")
torch.save(payload, temporary)
os.replace(temporary, path)
return path
def capture_rng_state() -> dict[str, Any]:
return {
"python": random.getstate(),
"torch_cpu": torch.get_rng_state(),
"torch_cuda": torch.cuda.get_rng_state(),
}
def restore_rng_state(state: dict[str, Any]) -> None:
random.setstate(state["python"])
torch.set_rng_state(state["torch_cpu"])
torch.cuda.set_rng_state(state["torch_cuda"])
|