amachine.am_transformers.am_checkpointer

  1import json
  2import shutil
  3import time
  4from dataclasses import asdict
  5from pathlib import Path
  6from typing import Any
  7import os
  8
  9import torch
 10from safetensors.torch import load_file as safetensors_load_file
 11from safetensors.torch import save_file as safetensors_save_file
 12from transformers.trainer_utils import load_sharded_checkpoint
 13
 14class _CallableSentinel:
 15    """Placeholder for unserializable Python functions found in state dicts."""
 16    pass
 17
 18class CheckpointManager:
 19
 20    def __init__(self, config, model, optimizer, trainer_state, device):
 21        self.config         = config
 22        self.model          = model
 23        self.optimizer      = optimizer
 24        self.trainer_state  = trainer_state
 25        self.device         = device
 26
 27    def _unwrap_model(self):
 28        return self.model._orig_mod if hasattr(self.model, "_orig_mod") else self.model
 29
 30    def save(self, step: int) -> None:
 31        ckpt = Path(self.config.output_dir) / f"step_{step:07d}"
 32        ckpt.mkdir(parents=True, exist_ok=True)
 33        t_save = time.time()
 34
 35        # Copy tokenizer / model config JSON files so the checkpoint is self-contained
 36        for json_file in Path(self.config.model_dir).glob("*.json"):
 37            shutil.copy2(json_file, ckpt / json_file.name)
 38
 39        self._unwrap_model().save_pretrained(str(ckpt), safe_serialization=True)
 40        self._save_optimizer(ckpt)
 41        self._save_trainer(ckpt)
 42        self._save_rng(ckpt)
 43        self.config.save(ckpt / "training_config.json")
 44
 45        print(f"Checkpoint saved: {ckpt}  ({time.time() - t_save:.1f}s)")
 46
 47    def load(self, path: str) -> None:
 48        
 49        ckpt = Path(path)
 50        print(f"Resuming from {ckpt}.")
 51
 52        result = load_sharded_checkpoint(
 53            self._unwrap_model(), str(ckpt), strict=True, prefer_safe=True
 54        )
 55        if result.missing_keys:
 56            print(f"  Warning: missing keys:    {result.missing_keys}")
 57        if result.unexpected_keys:
 58            print(f"  Warning: unexpected keys: {result.unexpected_keys}")
 59
 60        self._load_optimizer(ckpt)
 61        self._load_trainer(ckpt)
 62
 63        if not self.config.resume_reset_data:
 64            self._load_rng(ckpt)
 65
 66    # Optimizer
 67
 68    def _encode_node(self, obj: Any, tensors: dict[str, torch.Tensor]) -> dict[str, Any]:
 69        """Recursively encodes optimizer state into a strict, explicitly tagged AST schema."""
 70        if isinstance(obj, torch.Tensor):
 71            if obj.layout != torch.strided:
 72                raise TypeError(f"Unsupported tensor layout {obj.layout}. Safetensors requires strided tensors.")
 73            
 74            t_id = f"t_{len(tensors):08d}"
 75            # Detach removes autograd graph refs; contiguous prepares for Safetensors
 76            tensors[t_id] = obj.detach().cpu().contiguous()
 77            return {"t": "tensor", "v": t_id}
 78            
 79        elif isinstance(obj, dict):
 80            # Encode as list of pairs to preserve non-string keys (e.g., int parameter IDs)
 81            return {"t": "dict", "v": [[self._encode_node(k, tensors), self._encode_node(v, tensors)] for k, v in obj.items()]}
 82            
 83        elif isinstance(obj, list):
 84            return {"t": "list", "v": [self._encode_node(x, tensors) for x in obj]}
 85            
 86        elif isinstance(obj, torch.Size):
 87            return {"t": "torch_size", "v": [self._encode_node(x, tensors) for x in obj]}
 88            
 89        elif isinstance(obj, tuple):
 90            return {"t": "tuple", "v": [self._encode_node(x, tensors) for x in obj]}
 91            
 92        elif isinstance(obj, torch.dtype):
 93            # str(torch.float32) -> "torch.float32"
 94            return {"t": "torch_dtype", "v": str(obj)}
 95            
 96        elif isinstance(obj, torch.device):
 97            # str(torch.device('cuda:0')) -> "cuda:0"
 98            return {"t": "torch_device", "v": str(obj)}
 99            
100        elif isinstance(obj, (int, float, str, bool, type(None))):
101            return {"t": "primitive", "v": obj}
102            
103        elif callable(obj):
104            # PSGD/Kron put functions in param_groups. We flag them to merge back on load.
105            return {"t": "callable", "v": getattr(obj, "__name__", "anonymous_function")}
106            
107        else:
108            raise TypeError(f"Unsupported python type {type(obj).__name__} in optimizer state: {obj}")
109
110    def _decode_node(self, node: dict[str, Any], tensors: dict[str, torch.Tensor]) -> Any:
111        """Reconstructs the exact Python objects from the tagged AST schema."""
112        t = node["t"]
113        v = node["v"]
114        
115        if t == "tensor":       return tensors[v]
116        if t == "dict":         return {self._decode_node(k, tensors): self._decode_node(val, tensors) for k, val in v}
117        if t == "list":         return [self._decode_node(x, tensors) for x in v]
118        if t == "tuple":        return tuple(self._decode_node(x, tensors) for x in v)
119        if t == "torch_size":   return torch.Size(self._decode_node(x, tensors) for x in v)
120        if t == "torch_dtype":  return getattr(torch, v.split(".")[-1]) # "torch.float32" -> torch.float32
121        if t == "torch_device": return torch.device(v)
122        if t == "primitive":    return v
123        if t == "callable":     return _CallableSentinel()
124        
125        raise ValueError(f"Corrupt metadata: Unknown node type '{t}'")
126
127    def _restore_unserializables(self, rebuilt: Any, current: Any) -> Any:
128        """Walks the rebuilt state dict to inject functions from the live instantiated optimizer."""
129        if isinstance(rebuilt, dict) and isinstance(current, dict):
130            return {k: self._restore_unserializables(v, current.get(k)) for k, v in rebuilt.items()}
131        elif isinstance(rebuilt, list) and isinstance(current, list):
132            return [self._restore_unserializables(r, c) for r, c in zip(rebuilt, current)]
133        elif isinstance(rebuilt, tuple) and isinstance(current, tuple):
134            return tuple(self._restore_unserializables(r, c) for r, c in zip(rebuilt, current))
135        elif isinstance(rebuilt, _CallableSentinel):
136            return current  # Inject the live function back in
137        else:
138            return rebuilt
139
140    def _save_optimizer(self, ckpt: Path) -> None:
141        # Rank guard: Ensure only the main process writes to disk in normal DDP
142        if getattr(self, "global_rank", 0) != 0:
143            return
144
145        ckpt.mkdir(parents=True, exist_ok=True)
146        tensors: dict[str, torch.Tensor] = {}
147        state_dict = self.optimizer.state_dict()
148
149        meta = {
150            "__format_version__": 2,
151            "optimizer_state_dict": self._encode_node(state_dict, tensors)
152        }
153
154        sf_path = ckpt / "optimizer.safetensors"
155        meta_path = ckpt / "optimizer_meta.json"
156
157        tmp_sf = sf_path.with_suffix(".safetensors.tmp")
158        tmp_meta = meta_path.with_suffix(".json.tmp")
159
160        if tensors:
161            safetensors_save_file(tensors, tmp_sf)
162
163        with open(tmp_meta, "w") as f:
164            json.dump(meta, f, indent=2)
165            f.flush()
166            os.fsync(f.fileno())  # Guarantee disk persistence
167
168        # Atomic writes prevent corruption if the job is preempted mid-save
169        os.replace(tmp_meta, meta_path)
170        if tensors:
171            os.replace(tmp_sf, sf_path)
172        elif sf_path.exists():
173            # Clean up old safetensors if the optimizer structure changed to have no tensors
174            os.remove(sf_path)
175
176    def _load_optimizer(self, ckpt: Path) -> None:
177        meta_path = ckpt / "optimizer_meta.json"
178        sf_path   = ckpt / "optimizer.safetensors"
179
180        if not meta_path.exists():
181            print(f"  No optimizer metadata found at {meta_path} — starting fresh.")
182            return
183
184        with open(meta_path) as f:
185            meta = json.load(f)
186
187        if meta.get("__format_version__") != 2:
188            raise RuntimeError(f"Unsupported optimizer metadata format version {meta.get('__format_version__')} in {meta_path}")
189
190        # If sf_path doesn't exist, assume a tensor-less optimizer state
191        tensors = {}
192        if sf_path.exists():
193            tensors = safetensors_load_file(str(sf_path), device="cpu")
194
195        # 1. Decode the structural tree
196        rebuilt_state_dict = self._decode_node(meta["optimizer_state_dict"], tensors)
197        
198        # 2. Merge back live functions (needed for PSGD/Kron)
199        live_state_dict = self.optimizer.state_dict()
200        final_state_dict = self._restore_unserializables(rebuilt_state_dict, live_state_dict)
201
202        try:
203            self.optimizer.load_state_dict(final_state_dict)
204        except Exception as e:
205            raise RuntimeError(f"Failed to load optimizer state from {ckpt}: {e}") from e
206
207    # Trainer state
208
209    def _save_trainer(self, ckpt: Path) -> None:
210        state = asdict(self.trainer_state)
211        # Snapshot two extra fields for forward-compatibility checks on load
212        state["_data_path"]    = self.config.data
213        state["_was_compiled"] = self.config.compile
214        with open(ckpt / "trainer_state.json", "w") as f:
215            json.dump(state, f, indent=2)
216
217    def _load_trainer(self, ckpt: Path) -> None:
218        json_path = ckpt / "trainer_state.json"
219
220        if not json_path.exists():
221            print("  No trainer state found — all counters reset to defaults.")
222            return
223
224        with open(json_path) as f:
225            state = json.load(f)
226
227        # Restore all fields that exist on the current dataclass (forward-compatible)
228        for k, v in state.items():
229            if hasattr(self.trainer_state, k):
230                setattr(self.trainer_state, k, v)
231
232        # Warn if resuming onto a different dataset
233        saved_path = state.get("_data_path")
234        if saved_path and saved_path != self.config.data:
235            print(
236                f"  Warning: checkpoint trained on '{saved_path}', "
237                f"current data is '{self.config.data}'. "
238                f"Set resume_reset_data=True if this is intentional."
239            )
240
241        # Hard error on compile-mode mismatch (compiled vs eager weights are incompatible)
242        was_compiled = state.get("_was_compiled", False)
243        if was_compiled != self.config.compile:
244            raise ValueError(
245                f"Compile mode mismatch: checkpoint _was_compiled={was_compiled}, "
246                f"current compile={self.config.compile}."
247            )
248
249        print(
250            f"  Trainer state restored: step={self.trainer_state.global_step:,}, "
251            f"tokens={self.trainer_state.global_tokens:,}, "
252            f"loss_ema={self.trainer_state.spike_loss_ema}, "
253            f"batches_consumed={self.trainer_state.batches_consumed:,}"
254        )
255
256    # RNG state
257
258    def _save_rng(self, ckpt: Path) -> None:
259        tensors: dict[str, torch.Tensor] = {"torch_cpu": torch.get_rng_state()}
260
261        cuda_states = torch.cuda.get_rng_state_all() if torch.cuda.is_available() else []
262        for i, s in enumerate(cuda_states):
263            tensors[f"torch_cuda_{i}"] = s
264
265        safetensors_save_file(tensors, ckpt / "rng_state.safetensors")
266        with open(ckpt / "rng_meta.json", "w") as f:
267            json.dump({"num_cuda_devices": len(cuda_states)}, f)
268
269    def _load_rng(self, ckpt: Path) -> None:
270        sf_path   = ckpt / "rng_state.safetensors"
271        meta_path = ckpt / "rng_meta.json"
272
273        if not sf_path.exists() or not meta_path.exists():
274            print("  No RNG state found — skipping (data order may differ slightly on resume).")
275            return
276
277        with open(meta_path) as f:
278            meta = json.load(f)
279
280        tensors = safetensors_load_file(str(sf_path), device="cpu")
281        torch.set_rng_state(tensors["torch_cpu"])
282
283        if torch.cuda.is_available():
284            n_saved = meta["num_cuda_devices"]
285            n_avail = torch.cuda.device_count()
286            if n_saved != n_avail:
287                print(
288                    f"  Warning: checkpoint has {n_saved} CUDA device(s), "
289                    f"current run has {n_avail} — CUDA RNG state not restored."
290                )
291            else:
292                torch.cuda.set_rng_state_all(
293                    [tensors[f"torch_cuda_{i}"] for i in range(n_saved)]
294                )
295
296        print("  RNG state restored.")