Load the latest checkpoint and return the global step.
(model, optimizer, checkpoint_dir, device)
| 195 | |
| 196 | |
| 197 | def load_checkpoint(model, optimizer, checkpoint_dir, device): |
| 198 | """Load the latest checkpoint and return the global step.""" |
| 199 | checkpoint_steps = [ |
| 200 | int(d.name) |
| 201 | for d in checkpoint_dir.iterdir() |
| 202 | if d.is_dir() and d.name.isdigit() and not d.name.startswith("tmp_") |
| 203 | ] |
| 204 | |
| 205 | if not checkpoint_steps: |
| 206 | raise FileNotFoundError(f"No checkpoints found in {checkpoint_dir}") |
| 207 | |
| 208 | latest_step = max(checkpoint_steps) |
| 209 | ckpt_dir = checkpoint_dir / f"{latest_step}" |
| 210 | |
| 211 | # Clear memory before loading checkpoints |
| 212 | if torch.cuda.is_available(): |
| 213 | torch.cuda.empty_cache() |
| 214 | gc.collect() |
| 215 | log_memory_usage(device, latest_step, "before_loading_checkpoint") |
| 216 | |
| 217 | try: |
| 218 | # Load model state with error handling |
| 219 | logging.info("Loading model state...") |
| 220 | safetensors_path = ckpt_dir / "model.safetensors" |
| 221 | |
| 222 | if safetensors_path.exists(): |
| 223 | model_to_load = model.module if isinstance(model, torch.nn.parallel.DistributedDataParallel) else model |
| 224 | safetensors.torch.load_model(model_to_load, safetensors_path, device=str(device)) |
| 225 | logging.info("Loaded model state from safetensors format") |
| 226 | else: |
| 227 | raise FileNotFoundError(f"No model checkpoint found at {ckpt_dir}") |
| 228 | |
| 229 | torch.cuda.empty_cache() |
| 230 | gc.collect() |
| 231 | log_memory_usage(device, latest_step, "after_loading_model") |
| 232 | |
| 233 | # Load optimizer state with error handling |
| 234 | logging.info("Loading optimizer state...") |
| 235 | optimizer_path = ckpt_dir / "optimizer.pt" |
| 236 | |
| 237 | if optimizer_path.exists(): |
| 238 | optimizer_state_dict = torch.load(optimizer_path, map_location=device, weights_only=False) |
| 239 | logging.info("Loaded optimizer state from pt format") |
| 240 | else: |
| 241 | raise FileNotFoundError(f"No optimizer checkpoint found at {ckpt_dir}") |
| 242 | |
| 243 | optimizer.load_state_dict(optimizer_state_dict) |
| 244 | del optimizer_state_dict |
| 245 | torch.cuda.empty_cache() |
| 246 | gc.collect() |
| 247 | log_memory_usage(device, latest_step, "after_loading_optimizer") |
| 248 | |
| 249 | # Load metadata |
| 250 | logging.info("Loading metadata...") |
| 251 | metadata = torch.load(ckpt_dir / "metadata.pt", map_location=device, weights_only=False) |
| 252 | global_step = metadata.get("global_step", latest_step) |
| 253 | del metadata |
| 254 | torch.cuda.empty_cache() |
no test coverage detected