| import os, torch, random | |
| from contextlib import contextmanager | |
| def set_seed(seed): | |
| random.seed(seed); torch.manual_seed(seed); torch.cuda.manual_seed_all(seed) | |
| def autocast_if_bf16(enabled=True): | |
| if enabled and torch.cuda.is_available() and torch.cuda.is_bf16_supported(): | |
| with torch.cuda.amp.autocast(dtype=torch.bfloat16): | |
| yield | |
| else: | |
| yield | |
| def save_checkpoint(model, opt, step, ckpt_dir): | |
| os.makedirs(ckpt_dir, exist_ok=True) | |
| path = os.path.join(ckpt_dir, f"oai_flash_step{step}.pt") | |
| torch.save({"model": model.state_dict(), "opt": opt.state_dict(), "step": step}, path) | |
| print("Saved:", path) | |
| def latest_checkpoint(ckpt_dir): | |
| if not os.path.exists(ckpt_dir): return None | |
| files = [f for f in os.listdir(ckpt_dir) if f.startswith("oai_flash_step")] | |
| if not files: return None | |
| files.sort(key=lambda x: int(x.split("step")[-1].split(".")[0])) | |
| return os.path.join(ckpt_dir, files[-1]) | |