spb/forge Public MIT

1// Author: Simon-Pierre Boucher — contact@spboucher.ai2#include "train/checkpoint.h"34#include <cstdio>5#include <cstdlib>6#include <cstring>7#include <unordered_map>8#include <unordered_set>910namespace forge::train {1112namespace {1314constexpr uint32_t kMagic = 0x45475246; // "FRGE"15constexpr uint32_t kVersion = 1;1617[[noreturn]] void die(const std::string& msg) {18    std::fprintf(stderr, "forge/checkpoint: %s\n", msg.c_str());19    std::abort();20}2122void write_bytes(FILE* f, const void* p, size_t n) {23    if (std::fwrite(p, 1, n, f) != n) die("write failed");24}25void read_bytes(FILE* f, void* p, size_t n) {26    if (std::fread(p, 1, n, f) != n) die("read failed (truncated checkpoint?)");27}28template <typename T>29void write_pod(FILE* f, T v) { write_bytes(f, &v, sizeof(T)); }30template <typename T>31T read_pod(FILE* f) {32    T v;33    read_bytes(f, &v, sizeof(T));34    return v;35}36void write_str(FILE* f, const std::string& s) {37    write_pod<uint32_t>(f, uint32_t(s.size()));38    write_bytes(f, s.data(), s.size());39}40std::string read_str(FILE* f) {41    const uint32_t n = read_pod<uint32_t>(f);42    std::string s(n, '\0');43    read_bytes(f, s.data(), n);44    return s;45}4647// Tied params appear under several names; store each storage once (first48// name wins) and load by name into whichever Var carries it.49std::vector<std::pair<std::string, Var>> dedupe(50    const std::vector<std::pair<std::string, Var>>& named) {51    std::vector<std::pair<std::string, Var>> out;52    std::unordered_set<const void*> seen;53    for (const auto& [name, p] : named)54        if (seen.insert(p.id()).second) out.emplace_back(name, p);55    return out;56}5758} // namespace5960void save_checkpoint(const std::string& path,61                     const std::vector<std::pair<std::string, Var>>& named_params,62                     AdamW* opt, const CheckpointData& meta) {63    const std::string tmp = path + ".tmp";64    FILE* f = std::fopen(tmp.c_str(), "wb");65    if (!f) die("cannot open " + tmp);6667    write_pod<uint32_t>(f, kMagic);68    write_pod<uint32_t>(f, kVersion);69    write_pod<int64_t>(f, meta.step);70    write_pod<int64_t>(f, opt ? opt->t() : 0);71    write_pod<uint64_t>(f, meta.rng_state);72    write_str(f, meta.config_json);7374    const auto params = dedupe(named_params);75    write_pod<uint32_t>(f, uint32_t(params.size()));76    for (const auto& [name, p] : params) {77        write_str(f, name);78        const auto& shape = p.value().shape();79        write_pod<uint32_t>(f, uint32_t(shape.size()));80        for (int64_t d : shape) write_pod<int64_t>(f, d);81        write_bytes(f, p.value().raw(), p.value().nbytes());82    }8384    const bool has_opt = opt && !opt->m().empty();85    write_pod<uint32_t>(f, has_opt ? uint32_t(opt->m().size()) : 0u);86    if (has_opt) {87        for (size_t i = 0; i < opt->m().size(); ++i) {88            write_bytes(f, opt->m()[i].raw(), opt->m()[i].nbytes());89            write_bytes(f, opt->v()[i].raw(), opt->v()[i].nbytes());90        }91    }9293    std::fclose(f);94    if (std::rename(tmp.c_str(), path.c_str()) != 0) die("rename failed: " + path);95}9697CheckpointData load_checkpoint(const std::string& path,98                               const std::vector<std::pair<std::string, Var>>& named_params,99                               AdamW* opt) {100    FILE* f = std::fopen(path.c_str(), "rb");101    if (!f) die("cannot open " + path);102    if (read_pod<uint32_t>(f) != kMagic) die("bad magic: " + path);103    if (read_pod<uint32_t>(f) != kVersion) die("unsupported version: " + path);104105    CheckpointData meta;106    meta.step = read_pod<int64_t>(f);107    const int64_t opt_t = read_pod<int64_t>(f);108    meta.rng_state = read_pod<uint64_t>(f);109    meta.config_json = read_str(f);110111    std::unordered_map<std::string, Var> by_name;112    for (const auto& [name, p] : dedupe(named_params)) by_name.emplace(name, p);113114    const uint32_t n_params = read_pod<uint32_t>(f);115    if (n_params != by_name.size()) die("parameter count mismatch");116    for (uint32_t i = 0; i < n_params; ++i) {117        const std::string name = read_str(f);118        const uint32_t ndim = read_pod<uint32_t>(f);119        std::vector<int64_t> shape(ndim);120        for (auto& d : shape) d = read_pod<int64_t>(f);121        auto it = by_name.find(name);122        if (it == by_name.end()) die("unknown parameter in checkpoint: " + name);123        if (it->second.value().shape() != shape) die("shape mismatch: " + name);124        read_bytes(f, it->second.value().raw(), it->second.value().nbytes());125    }126127    const uint32_t n_opt = read_pod<uint32_t>(f);128    if (n_opt > 0 && opt) {129        opt->ensure_state();130        if (opt->m().size() != n_opt) die("optimizer state count mismatch");131        for (uint32_t i = 0; i < n_opt; ++i) {132            read_bytes(f, opt->m()[i].raw(), opt->m()[i].nbytes());133            read_bytes(f, opt->v()[i].raw(), opt->v()[i].nbytes());134        }135        opt->set_t(opt_t);136    } else if (n_opt > 0) {137        // skip optimizer state138        std::fseek(f, 0, SEEK_END);139    }140141    std::fclose(f);142    return meta;143}144145std::string read_checkpoint_config(const std::string& path) {146    FILE* f = std::fopen(path.c_str(), "rb");147    if (!f) die("cannot open " + path);148    if (read_pod<uint32_t>(f) != kMagic) die("bad magic: " + path);149    if (read_pod<uint32_t>(f) != kVersion) die("unsupported version: " + path);150    read_pod<int64_t>(f);151    read_pod<int64_t>(f);152    read_pod<uint64_t>(f);153    const std::string cfg = read_str(f);154    std::fclose(f);155    return cfg;156}157158} // namespace forge::train159