spb/forge Public MIT

1// Author: Simon-Pierre Boucher — contact@spboucher.ai2#pragma once34#include "core/autograd.h"5#include "train/optimizer.h"67#include <cstdint>8#include <string>9#include <utility>10#include <vector>1112namespace forge::train {1314// Binary checkpoint: model params (by name), optimizer moments + t, step,15// RNG state, and the config JSON — everything needed for exact resume16// (CLAUDE.md training-loop requirements). Format:17//   magic "FRGE" u32 | version u32 | step i64 | opt_t i64 | rng u6418//   | config_json (u32 len + bytes)19//   | n_params u32 | per param: name (u32+bytes), ndim u32, dims i64[],20//     f32 data21//   | n_opt u32 | per tracked param: m f32[], v f32[]   (0 if never stepped)22struct CheckpointData {23    int64_t step = 0;24    uint64_t rng_state = 0;25    std::string config_json;26};2728void save_checkpoint(const std::string& path,29                     const std::vector<std::pair<std::string, Var>>& named_params,30                     AdamW* opt, const CheckpointData& meta);3132// Loads params by NAME into an already-constructed model (shape-checked).33// opt may be null (inference). Returns the stored metadata.34CheckpointData load_checkpoint(const std::string& path,35                               const std::vector<std::pair<std::string, Var>>& named_params,36                               AdamW* opt);3738// Reads just the config JSON (to construct the model before loading).39std::string read_checkpoint_config(const std::string& path);4041} // namespace forge::train42