diff --git a/include/nir_c_runtime/checkpoint.h b/include/nir_c_runtime/checkpoint.h index 6de7d4c..3a0f19b 100644 --- a/include/nir_c_runtime/checkpoint.h +++ b/include/nir_c_runtime/checkpoint.h @@ -1,12 +1,16 @@ -/* Torch-free checkpoint save/load plus the reverse bridge (#29, #30). +/* Torch-free checkpoint save/load plus the reverse bridge (#29, #30), + * with optimizer-state resume (#43). * * Checkpoints are version-stamped JSON (mirroring the - * `spikeforge-nir-graph` v1 envelope discipline): + * `spikeforge-nir-graph` v1 envelope discipline), currently version 2: * - * {"format": "spikeforge-c-checkpoint", "version": 1, + * {"format": "spikeforge-c-checkpoint", "version": 2, * "topology": "conv_net", "seed": 7, "step": 100, * "nodes": [{"name": "fc", "kind": "Affine", - * "shape": [4, 8], "weight": [...], "bias": [...]}]} + * "shape": [4, 8], "weight": [...], "bias": [...]}, ...], + * "optimizers": {"fc": {"kind": "adam", "step": 10, + * "m": [...], "v": [...], "step_bias": 10, + * "m_bias": [...], "v_bias": [...]}, ...}} * * Weights are flat row-major float arrays with an explicit shape; * floats print with %.9g, which round-trips float32 exactly. Unknown @@ -14,6 +18,24 @@ * kind/shape mismatches, and non-finite weights -- never a silent * reinterpretation. * + * The `optimizers` section (new in v2) is optional and keyed by node + * name. Each entry carries one node's optimizer state: `kind` is + * "sgd" (with `velocity` plus optional `velocity_bias`) or "adam" + * (with `step`/`m`/`v` plus optional `step_bias`/`m_bias`/`v_bias` + * for the bias). Bias arrays are omitted when the node has no bias + * or the bias is untrained. Hyperparameters (lr, momentum, betas, + * eps, weight decay) are NOT persisted: re-init optimizers with the + * same hyperparameters before load for a bit-identical resume. + * + * Version 1 files (no `optimizers` section) still load: weights + * restore normally and caller optimizers keep their current state + * (fresh when newly init). The plain dumps/save/loads/load entry + * points persist weights only, ignoring any `optimizers` section on + * load; the with_opt variants persist both. Every optimizer entry + * in the file must match a caller entry, else load fails rather + * than silently restarting that node's optimizer; caller entries + * absent from the file keep their state. + * * The JSON writer here is new code (the existing json module stays * parse-only, owned by another worker). The weight setter mutates * the graph's public param buffers in place, which is also the @@ -32,9 +54,11 @@ #include #include "nir_c_runtime/graph.h" +#include "nir_c_runtime/train.h" #define NIR_CHECKPOINT_FORMAT "spikeforge-c-checkpoint" -#define NIR_CHECKPOINT_VERSION 1 +#define NIR_CHECKPOINT_VERSION 2 +#define NIR_CHECKPOINT_VERSION_1 1 /* Growing JSON text builder. */ typedef struct nir_json_writer nir_json_writer_t; @@ -68,7 +92,8 @@ typedef struct { void nir_checkpoint_meta_free(nir_checkpoint_meta_t *meta); /* Serialize `graph`'s weights plus `meta` (NULL for weights only) - * to a malloc'd JSON document. */ + * to a malloc'd JSON document. Weights only: any optimizer state + * is omitted (see the with_opt variants). */ char *nir_checkpoint_dumps( const nir_graph_t *graph, const nir_checkpoint_meta_t *meta, char **error @@ -83,7 +108,8 @@ int nir_checkpoint_save( /* Load weights from a checkpoint document into an already-loaded * graph: every weighted node must be present with matching kind * and shape. Fills `meta` (NULL to skip; caller frees with - * `nir_checkpoint_meta_free`). */ + * `nir_checkpoint_meta_free`). Weights only: any `optimizers` + * section in the document is ignored. */ int nir_checkpoint_loads( const char *text, nir_graph_t *graph, nir_checkpoint_meta_t *meta, char **error @@ -95,6 +121,67 @@ int nir_checkpoint_load( char **error ); +/* One node's optimizer state for checkpoint save/load. + * + * `node` borrows the graph node name. `kind` selects which + * pointers are live: SGD uses `sgd_weight` (required) and + * `sgd_bias` (NULL when the node has no bias or the bias is + * untrained); Adam uses `adam_weight`/`adam_bias` the same way. + * The inactive kind's pointers must be NULL. Weight/bias + * optimizers are separate `train.h` instances stepped per buffer; + * Adam weight/bias steps may differ and persist separately. */ +typedef enum { + NIR_CHECKPOINT_OPT_SGD, + NIR_CHECKPOINT_OPT_ADAM +} nir_checkpoint_opt_kind_t; + +typedef struct { + const char *node; + nir_checkpoint_opt_kind_t kind; + nir_sgd_t *sgd_weight; + nir_sgd_t *sgd_bias; + nir_adam_t *adam_weight; + nir_adam_t *adam_bias; +} nir_checkpoint_opt_t; + +/* Serialize weights plus `opt_count` optimizer states. NULL/0 + * writes weights only (same document as `nir_checkpoint_dumps`). + * Fails on unknown/unweighted/duplicate nodes, buffer counts that + * do not match the graph, mixed-kind entries, and non-finite + * state. */ +char *nir_checkpoint_dumps_with_opt( + const nir_graph_t *graph, const nir_checkpoint_meta_t *meta, + const nir_checkpoint_opt_t *opts, size_t opt_count, + char **error +); + +/* Write the same document to `path`. */ +int nir_checkpoint_save_with_opt( + const char *path, const nir_graph_t *graph, + const nir_checkpoint_meta_t *meta, + const nir_checkpoint_opt_t *opts, size_t opt_count, + char **error +); + +/* Load weights plus optimizer state: every `optimizers` entry in + * the document must match a caller entry (by node name and kind), + * else load fails instead of silently restarting that optimizer. + * Caller entries absent from the document (including all of them + * for a v1 file) keep their current state. NULL/0 restores weights + * only, ignoring any `optimizers` section. */ +int nir_checkpoint_loads_with_opt( + const char *text, nir_graph_t *graph, nir_checkpoint_meta_t *meta, + nir_checkpoint_opt_t *opts, size_t opt_count, + char **error +); + +/* Load weights plus optimizer state from the file at `path`. */ +int nir_checkpoint_load_with_opt( + const char *path, nir_graph_t *graph, nir_checkpoint_meta_t *meta, + nir_checkpoint_opt_t *opts, size_t opt_count, + char **error +); + /* Point `*weight`/`*bias` at node `name`'s buffers with element * counts (`*bias` NULL, `*bias_count` 0 when the node has no bias). * Returns 0 for unknown nodes and nodes without weights. */ diff --git a/src/checkpoint.c b/src/checkpoint.c index 9a887b3..9d2818e 100644 --- a/src/checkpoint.c +++ b/src/checkpoint.c @@ -6,6 +6,7 @@ #include #include "nir_c_runtime/json.h" +#include "nir_c_runtime/train.h" struct nir_json_writer { char *text; @@ -466,8 +467,192 @@ static int write_node(nir_json_writer_t *w, const nir_node_t *node) { nir_json_end(w); } -char *nir_checkpoint_dumps( +/* Validate caller optimizer entries against the graph: known + * weighted nodes, no duplicates, live-kind weight set with + * matching buffer counts, inactive-kind pointers NULL. */ +static int validate_opt_entries( + const nir_graph_t *graph, const nir_checkpoint_opt_t *opts, + size_t opt_count, char **error +) { + if (opt_count == 0) { + return 1; + } + if (opts == NULL) { + set_error(error, "checkpoint optimizer entry needs a node"); + return 0; + } + for (size_t i = 0; i < opt_count; i++) { + const nir_checkpoint_opt_t *o = &opts[i]; + if (o->node == NULL) { + set_error(error, "checkpoint optimizer entry needs a node"); + return 0; + } + if (o->kind != NIR_CHECKPOINT_OPT_SGD && + o->kind != NIR_CHECKPOINT_OPT_ADAM) { + set_error(error, "checkpoint optimizer entry is invalid"); + return 0; + } + if (o->kind == NIR_CHECKPOINT_OPT_SGD) { + if (o->sgd_weight == NULL) { + set_error( + error, "checkpoint optimizer entry needs state" + ); + return 0; + } + if (o->adam_weight != NULL || o->adam_bias != NULL) { + set_error( + error, "checkpoint optimizer entry mixes kinds" + ); + return 0; + } + } else { + if (o->adam_weight == NULL) { + set_error( + error, "checkpoint optimizer entry needs state" + ); + return 0; + } + if (o->sgd_weight != NULL || o->sgd_bias != NULL) { + set_error( + error, "checkpoint optimizer entry mixes kinds" + ); + return 0; + } + } + size_t index = nir_graph_find(graph, o->node); + if (index == (size_t)-1) { + set_error( + error, "checkpoint optimizer names an unknown node" + ); + return 0; + } + const float *w = NULL; + const float *b = NULL; + size_t wn = 0; + size_t bn = 0; + if (!weight_shapes(&graph->nodes[index], &w, &wn, &b, &bn)) { + set_error( + error, "checkpoint optimizer targets an unweighted node" + ); + return 0; + } + if (o->kind == NIR_CHECKPOINT_OPT_SGD) { + if (o->sgd_weight->velocity == NULL || + o->sgd_weight->count != wn) { + set_error( + error, + "checkpoint optimizer state does not match graph" + ); + return 0; + } + if (o->sgd_bias != NULL && + (b == NULL || o->sgd_bias->velocity == NULL || + o->sgd_bias->count != bn)) { + set_error( + error, + "checkpoint optimizer state does not match graph" + ); + return 0; + } + } else { + if (o->adam_weight->m == NULL || o->adam_weight->v == NULL || + o->adam_weight->count != wn) { + set_error( + error, + "checkpoint optimizer state does not match graph" + ); + return 0; + } + if (o->adam_bias != NULL && + (b == NULL || o->adam_bias->m == NULL || + o->adam_bias->v == NULL || + o->adam_bias->count != bn)) { + set_error( + error, + "checkpoint optimizer state does not match graph" + ); + return 0; + } + } + for (size_t j = 0; j < i; j++) { + if (strcmp(opts[j].node, o->node) == 0) { + set_error( + error, "checkpoint optimizer entry is duplicated" + ); + return 0; + } + } + } + return 1; +} + +static int write_sgd_entry( + nir_json_writer_t *w, const nir_checkpoint_opt_t *o +) { + return nir_json_key(w, "kind") && nir_json_string(w, "sgd") && + nir_json_key(w, "velocity") && + write_floats( + w, o->sgd_weight->velocity, o->sgd_weight->count + ) && + (o->sgd_bias == NULL || + (nir_json_key(w, "velocity_bias") && + write_floats( + w, o->sgd_bias->velocity, o->sgd_bias->count + ))); +} + +static int write_adam_entry( + nir_json_writer_t *w, const nir_checkpoint_opt_t *o +) { + int ok = nir_json_key(w, "kind") && nir_json_string(w, "adam") && + nir_json_key(w, "step") && + nir_json_int(w, (long long)o->adam_weight->step) && + nir_json_key(w, "m") && + write_floats(w, o->adam_weight->m, o->adam_weight->count) && + nir_json_key(w, "v") && + write_floats(w, o->adam_weight->v, o->adam_weight->count); + if (ok && o->adam_bias != NULL) { + ok = nir_json_key(w, "step_bias") && + nir_json_int(w, (long long)o->adam_bias->step) && + nir_json_key(w, "m_bias") && + write_floats( + w, o->adam_bias->m, o->adam_bias->count + ) && + nir_json_key(w, "v_bias") && + write_floats( + w, o->adam_bias->v, o->adam_bias->count + ); + } + return ok; +} + +static int write_optimizers( + nir_json_writer_t *w, const nir_checkpoint_opt_t *opts, + size_t opt_count +) { + if (!nir_json_key(w, "optimizers") || !nir_json_begin_object(w)) { + return 0; + } + for (size_t i = 0; i < opt_count; i++) { + const nir_checkpoint_opt_t *o = &opts[i]; + int ok = nir_json_key(w, o->node) && + nir_json_begin_object(w); + if (ok) { + ok = o->kind == NIR_CHECKPOINT_OPT_SGD + ? write_sgd_entry(w, o) + : write_adam_entry(w, o); + } + ok = ok && nir_json_end(w); + if (!ok) { + return 0; + } + } + return nir_json_end(w); +} + +static char *dumps_impl( const nir_graph_t *graph, const nir_checkpoint_meta_t *meta, + const nir_checkpoint_opt_t *opts, size_t opt_count, char **error ) { if (graph == NULL) { @@ -501,7 +686,13 @@ char *nir_checkpoint_dumps( ok = write_node(w, node); } } - ok = ok && nir_json_end(w) && nir_json_end(w); + ok = ok && nir_json_end(w); + int opt_failed = 0; + if (ok && opts != NULL && opt_count > 0) { + ok = write_optimizers(w, opts, opt_count); + opt_failed = !ok; + } + ok = ok && nir_json_end(w); char *text = NULL; if (ok && nir_json_writer_data(w) != NULL) { size_t length = nir_json_writer_length(w) + 1; @@ -512,23 +703,35 @@ char *nir_checkpoint_dumps( } nir_json_writer_free(w); if (text == NULL) { - set_error(error, "checkpoint weights must be finite"); + set_error( + error, opt_failed + ? "checkpoint optimizer state must be finite" + : "checkpoint weights must be finite" + ); } return text; } -int nir_checkpoint_save( - const char *path, const nir_graph_t *graph, - const nir_checkpoint_meta_t *meta, char **error +char *nir_checkpoint_dumps( + const nir_graph_t *graph, const nir_checkpoint_meta_t *meta, + char **error ) { - if (path == NULL) { - set_error(error, "checkpoint save needs a path"); - return 0; - } - char *text = nir_checkpoint_dumps(graph, meta, error); - if (text == NULL) { - return 0; + return dumps_impl(graph, meta, NULL, 0, error); +} + +char *nir_checkpoint_dumps_with_opt( + const nir_graph_t *graph, const nir_checkpoint_meta_t *meta, + const nir_checkpoint_opt_t *opts, size_t opt_count, + char **error +) { + if (graph != NULL && + !validate_opt_entries(graph, opts, opt_count, error)) { + return NULL; } + return dumps_impl(graph, meta, opts, opt_count, error); +} + +static int save_text(const char *path, char *text, char **error) { FILE *f = fopen(path, "wb"); int ok = 0; if (f != NULL) { @@ -543,6 +746,40 @@ int nir_checkpoint_save( return ok; } +int nir_checkpoint_save( + const char *path, const nir_graph_t *graph, + const nir_checkpoint_meta_t *meta, char **error +) { + if (path == NULL) { + set_error(error, "checkpoint save needs a path"); + return 0; + } + char *text = nir_checkpoint_dumps(graph, meta, error); + if (text == NULL) { + return 0; + } + return save_text(path, text, error); +} + +int nir_checkpoint_save_with_opt( + const char *path, const nir_graph_t *graph, + const nir_checkpoint_meta_t *meta, + const nir_checkpoint_opt_t *opts, size_t opt_count, + char **error +) { + if (path == NULL) { + set_error(error, "checkpoint save needs a path"); + return 0; + } + char *text = nir_checkpoint_dumps_with_opt( + graph, meta, opts, opt_count, error + ); + if (text == NULL) { + return 0; + } + return save_text(path, text, error); +} + /* Read a flat numeric array into `out[0..n)`; every element must be * a finite number. */ static int read_floats( @@ -687,9 +924,259 @@ static char *read_whole_file(const char *path) { return text; } -int nir_checkpoint_loads( - const char *text, nir_graph_t *graph, nir_checkpoint_meta_t *meta, +/* True when every element of an already length-checked array is a + * finite JSON number. */ +static int floats_finite(const json_value_t *array) { + size_t n = json_array_length(array); + for (size_t i = 0; i < n; i++) { + const json_value_t *item = json_array_get(array, i); + if (item == NULL || item->type != JSON_NUMBER || + !isfinite(item->as.number)) { + return 0; + } + } + return 1; +} + +/* Validate an Adam step counter: a non-negative integer within the + * exactly representable double range (2^53-1 covers any realistic + * training run). */ +static int read_opt_step(const json_value_t *v, unsigned long *out) { + if (v == NULL || v->type != JSON_NUMBER) { + return 0; + } + double d = v->as.number; + if (!isfinite(d) || d < 0.0 || d != floor(d) || + d > 9007199254740991.0) { + return 0; + } + *out = (unsigned long)d; + return 1; +} + +static nir_checkpoint_opt_t *find_opt( + nir_checkpoint_opt_t *opts, size_t opt_count, const char *node +) { + for (size_t i = 0; i < opt_count; i++) { + if (opts[i].node != NULL && strcmp(opts[i].node, node) == 0) { + return &opts[i]; + } + } + return NULL; +} + +/* Check one required optimizer buffer: present array, exact + * length, all finite. Length mismatches report a shape error; + * missing/non-array/non-finite report a flat-array error. */ +static int check_opt_buffer( + const json_value_t *entry, const char *key, size_t want, char **error +) { + const json_value_t *arr = json_object_get(entry, key); + if (arr == NULL || arr->type != JSON_ARRAY) { + set_error( + error, "checkpoint optimizer state is not a flat array" + ); + return 0; + } + if (json_array_length(arr) != want) { + set_error( + error, "checkpoint optimizer shape does not match graph" + ); + return 0; + } + if (!floats_finite(arr)) { + set_error( + error, "checkpoint optimizer state is not a flat array" + ); + return 0; + } + return 1; +} + +/* Validate one file optimizer entry and restore it into its + * matching caller entry (already located). Validates every field + * before writing any, so a failing entry leaves that node's + * optimizer untouched. */ +static int load_one_opt( + const char *name, const json_value_t *entry, + nir_checkpoint_opt_t *dst, size_t wn, size_t bn, int has_bias, + char **error +) { + (void)name; + if (entry == NULL || entry->type != JSON_OBJECT) { + set_error(error, "checkpoint optimizer entry is invalid"); + return 0; + } + const char *kind = json_as_string(json_object_get(entry, "kind")); + int is_sgd = kind != NULL && strcmp(kind, "sgd") == 0; + int is_adam = kind != NULL && strcmp(kind, "adam") == 0; + if (!is_sgd && !is_adam) { + set_error(error, "checkpoint optimizer has an unknown kind"); + return 0; + } + if ((is_sgd && dst->kind != NIR_CHECKPOINT_OPT_SGD) || + (is_adam && dst->kind != NIR_CHECKPOINT_OPT_ADAM)) { + set_error( + error, "checkpoint optimizer kind does not match entry" + ); + return 0; + } + if (is_sgd) { + if (!check_opt_buffer(entry, "velocity", wn, error)) { + return 0; + } + const json_value_t *vb = json_object_get(entry, "velocity_bias"); + int want_bias = vb != NULL && vb->type != JSON_NULL; + if (want_bias) { + if (!has_bias) { + set_error( + error, + "checkpoint optimizer shape does not match graph" + ); + return 0; + } + if (!check_opt_buffer( + entry, "velocity_bias", bn, error + )) { + return 0; + } + } + if (!read_floats( + json_object_get(entry, "velocity"), + dst->sgd_weight->velocity, wn + )) { + set_error( + error, "checkpoint optimizer state is not a flat array" + ); + return 0; + } + if (want_bias && dst->sgd_bias != NULL && + !read_floats(vb, dst->sgd_bias->velocity, bn)) { + set_error( + error, "checkpoint optimizer state is not a flat array" + ); + return 0; + } + return 1; + } + unsigned long step = 0; + if (!read_opt_step(json_object_get(entry, "step"), &step)) { + set_error(error, "checkpoint optimizer step is invalid"); + return 0; + } + if (!check_opt_buffer(entry, "m", wn, error) || + !check_opt_buffer(entry, "v", wn, error)) { + return 0; + } + const json_value_t *mb = json_object_get(entry, "m_bias"); + const json_value_t *vb = json_object_get(entry, "v_bias"); + int has_mb = mb != NULL && mb->type != JSON_NULL; + int has_vb = vb != NULL && vb->type != JSON_NULL; + if (has_mb != has_vb) { + set_error( + error, "checkpoint optimizer bias state is incomplete" + ); + return 0; + } + unsigned long step_bias = step; + if (has_mb) { + if (!has_bias) { + set_error( + error, + "checkpoint optimizer shape does not match graph" + ); + return 0; + } + const json_value_t *sb = json_object_get(entry, "step_bias"); + if (sb != NULL && sb->type != JSON_NULL && + !read_opt_step(sb, &step_bias)) { + set_error(error, "checkpoint optimizer step is invalid"); + return 0; + } + if (!check_opt_buffer(entry, "m_bias", bn, error) || + !check_opt_buffer(entry, "v_bias", bn, error)) { + return 0; + } + } + if (!read_floats( + json_object_get(entry, "m"), dst->adam_weight->m, wn + ) || + !read_floats( + json_object_get(entry, "v"), dst->adam_weight->v, wn + )) { + set_error( + error, "checkpoint optimizer state is not a flat array" + ); + return 0; + } + dst->adam_weight->step = step; + if (has_mb && dst->adam_bias != NULL) { + if (!read_floats(mb, dst->adam_bias->m, bn) || + !read_floats(vb, dst->adam_bias->v, bn)) { + set_error( + error, "checkpoint optimizer state is not a flat array" + ); + return 0; + } + dst->adam_bias->step = step_bias; + } + return 1; +} + +static int load_optimizers( + const json_value_t *root, nir_graph_t *graph, + nir_checkpoint_opt_t *opts, size_t opt_count, char **error +) { + const json_value_t *opt_root = json_object_get(root, "optimizers"); + if (opt_root == NULL || opt_root->type == JSON_NULL) { + return 1; + } + if (opt_root->type != JSON_OBJECT) { + set_error( + error, "checkpoint optimizers section is not an object" + ); + return 0; + } + for (size_t i = 0; i < opt_root->as.object.count; i++) { + const char *name = opt_root->as.object.keys[i]; + const json_value_t *entry = opt_root->as.object.values[i]; + size_t index = nir_graph_find(graph, name); + if (index == (size_t)-1) { + set_error( + error, "checkpoint optimizer names an unknown node" + ); + return 0; + } + const float *w = NULL; + const float *b = NULL; + size_t wn = 0; + size_t bn = 0; + if (!weight_shapes(&graph->nodes[index], &w, &wn, &b, &bn)) { + set_error( + error, "checkpoint optimizer targets an unweighted node" + ); + return 0; + } + nir_checkpoint_opt_t *dst = find_opt(opts, opt_count, name); + if (dst == NULL) { + set_error( + error, "checkpoint optimizer has no matching entry" + ); + return 0; + } + if (!load_one_opt( + name, entry, dst, wn, bn, b != NULL, error + )) { + return 0; + } + } + return 1; +} + +static int loads_impl( + const char *text, nir_graph_t *graph, nir_checkpoint_meta_t *meta, + nir_checkpoint_opt_t *opts, size_t opt_count, char **error ) { if (text == NULL || graph == NULL) { set_error(error, "checkpoint load needs text and a graph"); @@ -708,7 +1195,7 @@ int nir_checkpoint_loads( if (!ok) { set_error(error, "checkpoint has an unknown format marker"); } else if (version == NULL || version->type != JSON_NUMBER || - version->as.number != (double)NIR_CHECKPOINT_VERSION) { + (version->as.number != 1.0 && version->as.number != 2.0)) { set_error(error, "checkpoint has an unsupported version"); ok = 0; } @@ -745,6 +1232,9 @@ int nir_checkpoint_loads( ok = 0; } } + if (ok && opts != NULL && opt_count > 0) { + ok = load_optimizers(root, graph, opts, opt_count, error); + } if (ok && meta != NULL) { memset(meta, 0, sizeof(*meta)); const char *topology = json_as_string( @@ -771,6 +1261,29 @@ int nir_checkpoint_loads( return ok; } +int nir_checkpoint_loads( + const char *text, nir_graph_t *graph, nir_checkpoint_meta_t *meta, + char **error +) { + return loads_impl(text, graph, meta, NULL, 0, error); +} + +int nir_checkpoint_loads_with_opt( + const char *text, nir_graph_t *graph, nir_checkpoint_meta_t *meta, + nir_checkpoint_opt_t *opts, size_t opt_count, + char **error +) { + if (graph != NULL && opt_count > 0 && + !validate_opt_entries(graph, opts, opt_count, error)) { + return 0; + } + /* NULL/0 restores weights only, ignoring any file optimizers. */ + if (opts == NULL || opt_count == 0) { + return loads_impl(text, graph, meta, NULL, 0, error); + } + return loads_impl(text, graph, meta, opts, opt_count, error); +} + int nir_checkpoint_load( const char *path, nir_graph_t *graph, nir_checkpoint_meta_t *meta, char **error @@ -789,6 +1302,27 @@ int nir_checkpoint_load( return ok; } +int nir_checkpoint_load_with_opt( + const char *path, nir_graph_t *graph, nir_checkpoint_meta_t *meta, + nir_checkpoint_opt_t *opts, size_t opt_count, + char **error +) { + if (path == NULL) { + set_error(error, "checkpoint load needs a path"); + return 0; + } + char *text = read_whole_file(path); + if (text == NULL) { + set_error(error, "cannot read checkpoint file"); + return 0; + } + int ok = nir_checkpoint_loads_with_opt( + text, graph, meta, opts, opt_count, error + ); + free(text); + return ok; +} + /* Emit `weight` (flat row-major, `wn` elements) as a nested array * with per-level sizes `dims[0..ndim)`. */ static int write_nested( diff --git a/tests/test_checkpoint.c b/tests/test_checkpoint.c index 3e7038e..0e91943 100644 --- a/tests/test_checkpoint.c +++ b/tests/test_checkpoint.c @@ -1,4 +1,5 @@ -/* Checkpoint round-trips, weight accessors, summary export (#29). */ +/* Checkpoint round-trips, weight accessors, summary export (#29), + * plus optimizer-state resume (#43). */ #include #include #include @@ -7,6 +8,7 @@ #include "nir_c_runtime/graph.h" #include "nir_c_runtime/json.h" #include "nir_c_runtime/runtime.h" +#include "nir_c_runtime/train.h" #include "test_util.h" static void check_writer(void) { @@ -108,6 +110,641 @@ static nir_tensor_t run_readout(nir_graph_t *graph) { return readout; } +static void fill_grad(float *g, size_t n, float base) { + for (size_t i = 0; i < n; i++) { + g[i] = base + 0.001f * (float)i; + } +} + +static float *copy_buf(const float *src, size_t n) { + float *dst = (float *)malloc(n * sizeof(float)); + CHECK(dst != NULL); + memcpy(dst, src, n * sizeof(float)); + return dst; +} + +static char *copy_text(const char *src) { + size_t n = strlen(src) + 1; + char *dst = (char *)malloc(n); + CHECK(dst != NULL); + memcpy(dst, src, n); + return dst; +} + +/* SGD resume: K steps, save, N more steps vs fresh+load+N steps + * are bit-identical in weights and velocities. */ +static void check_opt_sgd_resume(void) { + char *error = NULL; + nir_graph_t a = load_fixture("fixtures/lif_linear_graph.json"); + const float *w1 = NULL; + const float *b1 = NULL; + const float *w2 = NULL; + const float *b2 = NULL; + size_t w1n = 0; + size_t b1n = 0; + size_t w2n = 0; + size_t b2n = 0; + CHECK(nir_graph_get_weights(&a, "fc1", &w1, &w1n, &b1, &b1n)); + CHECK(nir_graph_get_weights(&a, "fc2", &w2, &w2n, &b2, &b2n)); + CHECK(b1 == NULL && b1n == 0); + nir_sgd_t s1w; + nir_sgd_t s2w; + nir_sgd_t s2b; + CHECK(nir_sgd_init(&s1w, w1n, 0.1f, 0.9f, 0.01f, &error)); + CHECK(nir_sgd_init(&s2w, w2n, 0.1f, 0.9f, 0.01f, &error)); + CHECK(nir_sgd_init(&s2b, b2n, 0.1f, 0.9f, 0.01f, &error)); + float *g1w = copy_buf(w1, w1n); + float *g2w = copy_buf(w2, w2n); + float *g2b = copy_buf(b2, b2n); + fill_grad(g1w, w1n, 0.1f); + fill_grad(g2w, w2n, 0.1f); + fill_grad(g2b, b2n, 0.2f); + for (int t = 0; t < 3; t++) { + CHECK(nir_sgd_step(&s1w, (float *)w1, g1w, w1n)); + CHECK(nir_sgd_step(&s2w, (float *)w2, g2w, w2n)); + CHECK(nir_sgd_step(&s2b, (float *)b2, g2b, b2n)); + } + nir_checkpoint_opt_t opts[2]; + memset(opts, 0, sizeof(opts)); + opts[0].node = "fc1"; + opts[0].kind = NIR_CHECKPOINT_OPT_SGD; + opts[0].sgd_weight = &s1w; + opts[1].node = "fc2"; + opts[1].kind = NIR_CHECKPOINT_OPT_SGD; + opts[1].sgd_weight = &s2w; + opts[1].sgd_bias = &s2b; + char *dump = nir_checkpoint_dumps_with_opt( + &a, NULL, opts, 2, &error + ); + CHECK(error == NULL && dump != NULL); + CHECK(strstr(dump, "\"optimizers\"") != NULL); + CHECK(strstr(dump, "\"velocity\"") != NULL); + CHECK(strstr(dump, "\"velocity_bias\"") != NULL); + for (int t = 0; t < 3; t++) { + CHECK(nir_sgd_step(&s1w, (float *)w1, g1w, w1n)); + CHECK(nir_sgd_step(&s2w, (float *)w2, g2w, w2n)); + CHECK(nir_sgd_step(&s2b, (float *)b2, g2b, b2n)); + } + float *a_w1 = copy_buf(w1, w1n); + float *a_w2 = copy_buf(w2, w2n); + float *a_b2 = copy_buf(b2, b2n); + float *a_v1 = copy_buf(s1w.velocity, w1n); + float *a_v2 = copy_buf(s2w.velocity, w2n); + float *a_vb = copy_buf(s2b.velocity, b2n); + + nir_graph_t b = load_fixture("fixtures/lif_linear_graph.json"); + nir_sgd_t r1w; + nir_sgd_t r2w; + nir_sgd_t r2b; + CHECK(nir_sgd_init(&r1w, w1n, 0.1f, 0.9f, 0.01f, &error)); + CHECK(nir_sgd_init(&r2w, w2n, 0.1f, 0.9f, 0.01f, &error)); + CHECK(nir_sgd_init(&r2b, b2n, 0.1f, 0.9f, 0.01f, &error)); + nir_checkpoint_opt_t ropts[2]; + memset(ropts, 0, sizeof(ropts)); + ropts[0].node = "fc1"; + ropts[0].kind = NIR_CHECKPOINT_OPT_SGD; + ropts[0].sgd_weight = &r1w; + ropts[1].node = "fc2"; + ropts[1].kind = NIR_CHECKPOINT_OPT_SGD; + ropts[1].sgd_weight = &r2w; + ropts[1].sgd_bias = &r2b; + CHECK(nir_checkpoint_loads_with_opt( + dump, &b, NULL, ropts, 2, &error + )); + CHECK(error == NULL); + const float *bw1 = NULL; + const float *bb1 = NULL; + const float *bw2 = NULL; + const float *bb2 = NULL; + size_t bw1n = 0; + size_t bb1n = 0; + size_t bw2n = 0; + size_t bb2n = 0; + CHECK(nir_graph_get_weights(&b, "fc1", &bw1, &bw1n, &bb1, &bb1n)); + CHECK(nir_graph_get_weights(&b, "fc2", &bw2, &bw2n, &bb2, &bb2n)); + for (int t = 0; t < 3; t++) { + CHECK(nir_sgd_step(&r1w, (float *)bw1, g1w, w1n)); + CHECK(nir_sgd_step(&r2w, (float *)bw2, g2w, w2n)); + CHECK(nir_sgd_step(&r2b, (float *)bb2, g2b, b2n)); + } + CHECK(memcmp(bw1, a_w1, w1n * sizeof(float)) == 0); + CHECK(memcmp(bw2, a_w2, w2n * sizeof(float)) == 0); + CHECK(memcmp(bb2, a_b2, b2n * sizeof(float)) == 0); + CHECK(memcmp(r1w.velocity, a_v1, w1n * sizeof(float)) == 0); + CHECK(memcmp(r2w.velocity, a_v2, w2n * sizeof(float)) == 0); + CHECK(memcmp(r2b.velocity, a_vb, b2n * sizeof(float)) == 0); + + /* File round-trip carries the same optimizer state. */ + const char *tmp = getenv("TMPDIR"); + char path[512]; + sprintf(path, "%s/nir-opt-test.json", tmp != NULL ? tmp : "/tmp"); + nir_graph_t c = load_fixture("fixtures/lif_linear_graph.json"); + nir_sgd_t f1w; + nir_sgd_t f2w; + nir_sgd_t f2b; + CHECK(nir_sgd_init(&f1w, w1n, 0.1f, 0.9f, 0.01f, &error)); + CHECK(nir_sgd_init(&f2w, w2n, 0.1f, 0.9f, 0.01f, &error)); + CHECK(nir_sgd_init(&f2b, b2n, 0.1f, 0.9f, 0.01f, &error)); + nir_checkpoint_opt_t fopts[2]; + memset(fopts, 0, sizeof(fopts)); + fopts[0].node = "fc1"; + fopts[0].kind = NIR_CHECKPOINT_OPT_SGD; + fopts[0].sgd_weight = &f1w; + fopts[1].node = "fc2"; + fopts[1].kind = NIR_CHECKPOINT_OPT_SGD; + fopts[1].sgd_weight = &f2w; + fopts[1].sgd_bias = &f2b; + CHECK(nir_checkpoint_save_with_opt( + path, &a, NULL, opts, 2, &error + )); + CHECK(error == NULL); + CHECK(nir_checkpoint_load_with_opt( + path, &c, NULL, fopts, 2, &error + )); + CHECK(error == NULL); + const float *cw1 = NULL; + const float *cb1 = NULL; + const float *cw2 = NULL; + const float *cb2 = NULL; + size_t cw1n = 0; + size_t cb1n = 0; + size_t cw2n = 0; + size_t cb2n = 0; + CHECK(nir_graph_get_weights(&c, "fc1", &cw1, &cw1n, &cb1, &cb1n)); + CHECK(nir_graph_get_weights(&c, "fc2", &cw2, &cw2n, &cb2, &cb2n)); + CHECK(memcmp(cw1, a_w1, w1n * sizeof(float)) == 0); + CHECK(memcmp(f1w.velocity, a_v1, w1n * sizeof(float)) == 0); + + free(g1w); + free(g2w); + free(g2b); + free(a_w1); + free(a_w2); + free(a_b2); + free(a_v1); + free(a_v2); + free(a_vb); + free(dump); + nir_sgd_free(&s1w); + nir_sgd_free(&s2w); + nir_sgd_free(&s2b); + nir_sgd_free(&r1w); + nir_sgd_free(&r2w); + nir_sgd_free(&r2b); + nir_sgd_free(&f1w); + nir_sgd_free(&f2w); + nir_sgd_free(&f2b); + nir_graph_free(&a); + nir_graph_free(&b); + nir_graph_free(&c); +} + +/* Adam resume: weights, m/v, and step counters are bit-identical. + * Bias correction makes the step counter load-bearing: a dropped + * step would diverge on the very next update. */ +static void check_opt_adam_resume(void) { + char *error = NULL; + nir_graph_t a = load_fixture("fixtures/lif_linear_graph.json"); + const float *w1 = NULL; + const float *b1 = NULL; + const float *w2 = NULL; + const float *b2 = NULL; + size_t w1n = 0; + size_t b1n = 0; + size_t w2n = 0; + size_t b2n = 0; + CHECK(nir_graph_get_weights(&a, "fc1", &w1, &w1n, &b1, &b1n)); + CHECK(nir_graph_get_weights(&a, "fc2", &w2, &w2n, &b2, &b2n)); + nir_adam_t s1w; + nir_adam_t s2w; + nir_adam_t s2b; + CHECK(nir_adam_init( + &s1w, w1n, 0.01f, 0.9f, 0.999f, 1e-8f, 0.0f, &error + )); + CHECK(nir_adam_init( + &s2w, w2n, 0.01f, 0.9f, 0.999f, 1e-8f, 0.0f, &error + )); + CHECK(nir_adam_init( + &s2b, b2n, 0.01f, 0.9f, 0.999f, 1e-8f, 0.0f, &error + )); + float *g1w = copy_buf(w1, w1n); + float *g2w = copy_buf(w2, w2n); + float *g2b = copy_buf(b2, b2n); + fill_grad(g1w, w1n, 0.1f); + fill_grad(g2w, w2n, 0.1f); + fill_grad(g2b, b2n, 0.2f); + for (int t = 0; t < 4; t++) { + CHECK(nir_adam_step(&s1w, (float *)w1, g1w, w1n)); + CHECK(nir_adam_step(&s2w, (float *)w2, g2w, w2n)); + CHECK(nir_adam_step(&s2b, (float *)b2, g2b, b2n)); + } + nir_checkpoint_opt_t opts[2]; + memset(opts, 0, sizeof(opts)); + opts[0].node = "fc1"; + opts[0].kind = NIR_CHECKPOINT_OPT_ADAM; + opts[0].adam_weight = &s1w; + opts[1].node = "fc2"; + opts[1].kind = NIR_CHECKPOINT_OPT_ADAM; + opts[1].adam_weight = &s2w; + opts[1].adam_bias = &s2b; + char *dump = nir_checkpoint_dumps_with_opt( + &a, NULL, opts, 2, &error + ); + CHECK(error == NULL && dump != NULL); + CHECK(strstr(dump, "\"kind\":\"adam\"") != NULL); + CHECK(strstr(dump, "\"step\":4") != NULL); + for (int t = 0; t < 4; t++) { + CHECK(nir_adam_step(&s1w, (float *)w1, g1w, w1n)); + CHECK(nir_adam_step(&s2w, (float *)w2, g2w, w2n)); + CHECK(nir_adam_step(&s2b, (float *)b2, g2b, b2n)); + } + CHECK(s1w.step == 8 && s2w.step == 8 && s2b.step == 8); + float *a_w1 = copy_buf(w1, w1n); + float *a_w2 = copy_buf(w2, w2n); + float *a_b2 = copy_buf(b2, b2n); + float *a_m1 = copy_buf(s1w.m, w1n); + float *a_v1 = copy_buf(s1w.v, w1n); + + nir_graph_t b = load_fixture("fixtures/lif_linear_graph.json"); + nir_adam_t r1w; + nir_adam_t r2w; + nir_adam_t r2b; + CHECK(nir_adam_init( + &r1w, w1n, 0.01f, 0.9f, 0.999f, 1e-8f, 0.0f, &error + )); + CHECK(nir_adam_init( + &r2w, w2n, 0.01f, 0.9f, 0.999f, 1e-8f, 0.0f, &error + )); + CHECK(nir_adam_init( + &r2b, b2n, 0.01f, 0.9f, 0.999f, 1e-8f, 0.0f, &error + )); + nir_checkpoint_opt_t ropts[2]; + memset(ropts, 0, sizeof(ropts)); + ropts[0].node = "fc1"; + ropts[0].kind = NIR_CHECKPOINT_OPT_ADAM; + ropts[0].adam_weight = &r1w; + ropts[1].node = "fc2"; + ropts[1].kind = NIR_CHECKPOINT_OPT_ADAM; + ropts[1].adam_weight = &r2w; + ropts[1].adam_bias = &r2b; + CHECK(nir_checkpoint_loads_with_opt( + dump, &b, NULL, ropts, 2, &error + )); + CHECK(error == NULL); + CHECK(r1w.step == 4 && r2w.step == 4 && r2b.step == 4); + const float *bw1 = NULL; + const float *bb1 = NULL; + const float *bw2 = NULL; + const float *bb2 = NULL; + size_t bw1n = 0; + size_t bb1n = 0; + size_t bw2n = 0; + size_t bb2n = 0; + CHECK(nir_graph_get_weights(&b, "fc1", &bw1, &bw1n, &bb1, &bb1n)); + CHECK(nir_graph_get_weights(&b, "fc2", &bw2, &bw2n, &bb2, &bb2n)); + for (int t = 0; t < 4; t++) { + CHECK(nir_adam_step(&r1w, (float *)bw1, g1w, w1n)); + CHECK(nir_adam_step(&r2w, (float *)bw2, g2w, w2n)); + CHECK(nir_adam_step(&r2b, (float *)bb2, g2b, b2n)); + } + CHECK(r1w.step == 8 && r2w.step == 8 && r2b.step == 8); + CHECK(memcmp(bw1, a_w1, w1n * sizeof(float)) == 0); + CHECK(memcmp(bw2, a_w2, w2n * sizeof(float)) == 0); + CHECK(memcmp(bb2, a_b2, b2n * sizeof(float)) == 0); + CHECK(memcmp(r1w.m, a_m1, w1n * sizeof(float)) == 0); + CHECK(memcmp(r1w.v, a_v1, w1n * sizeof(float)) == 0); + + free(g1w); + free(g2w); + free(g2b); + free(a_w1); + free(a_w2); + free(a_b2); + free(a_m1); + free(a_v1); + free(dump); + nir_adam_free(&s1w); + nir_adam_free(&s2w); + nir_adam_free(&s2b); + nir_adam_free(&r1w); + nir_adam_free(&r2w); + nir_adam_free(&r2b); + nir_graph_free(&a); + nir_graph_free(&b); +} + +/* A v1 checkpoint (no optimizers section) still loads; caller + * optimizers keep their fresh state. */ +static void check_opt_v1_compat(void) { + char *error = NULL; + nir_graph_t a = load_fixture("fixtures/lif_linear_graph.json"); + char *dump = nir_checkpoint_dumps(&a, NULL, &error); + CHECK(error == NULL && dump != NULL); + char *ver = strstr(dump, "\"version\":2"); + CHECK(ver != NULL); + ver[strlen("\"version\":")] = '1'; + CHECK(strstr(dump, "\"optimizers\"") == NULL); + + nir_graph_t b = load_fixture("fixtures/lif_linear_graph.json"); + const float *w1 = NULL; + const float *b1 = NULL; + size_t w1n = 0; + size_t b1n = 0; + CHECK(nir_graph_get_weights(&b, "fc1", &w1, &w1n, &b1, &b1n)); + nir_adam_t opt; + CHECK(nir_adam_init( + &opt, w1n, 0.01f, 0.9f, 0.999f, 1e-8f, 0.0f, &error + )); + nir_checkpoint_opt_t opts[1]; + memset(opts, 0, sizeof(opts)); + opts[0].node = "fc1"; + opts[0].kind = NIR_CHECKPOINT_OPT_ADAM; + opts[0].adam_weight = &opt; + CHECK(nir_checkpoint_loads_with_opt( + dump, &b, NULL, opts, 1, &error + )); + CHECK(error == NULL); + CHECK(opt.step == 0); + for (size_t i = 0; i < w1n; i++) { + CHECK(opt.m[i] == 0.0f && opt.v[i] == 0.0f); + } + /* Plain loads accepts v1 too. */ + nir_graph_t c = load_fixture("fixtures/lif_linear_graph.json"); + CHECK(nir_checkpoint_loads(dump, &c, NULL, &error)); + CHECK(error == NULL); + + free(dump); + nir_adam_free(&opt); + nir_graph_free(&a); + nir_graph_free(&b); + nir_graph_free(&c); +} + +/* Plain loads restores weights from a v2-with-optimizers file, + * ignoring the optimizer section. */ +static void check_opt_plain_ignores(void) { + char *error = NULL; + nir_graph_t a = load_fixture("fixtures/lif_linear_graph.json"); + const float *w1 = NULL; + const float *b1 = NULL; + size_t w1n = 0; + size_t b1n = 0; + CHECK(nir_graph_get_weights(&a, "fc1", &w1, &w1n, &b1, &b1n)); + nir_sgd_t opt; + CHECK(nir_sgd_init(&opt, w1n, 0.1f, 0.9f, 0.0f, &error)); + float *grad = copy_buf(w1, w1n); + fill_grad(grad, w1n, 0.1f); + CHECK(nir_sgd_step(&opt, (float *)w1, grad, w1n)); + nir_checkpoint_opt_t opts[1]; + memset(opts, 0, sizeof(opts)); + opts[0].node = "fc1"; + opts[0].kind = NIR_CHECKPOINT_OPT_SGD; + opts[0].sgd_weight = &opt; + char *dump = nir_checkpoint_dumps_with_opt( + &a, NULL, opts, 1, &error + ); + CHECK(error == NULL && dump != NULL); + CHECK(strstr(dump, "\"optimizers\"") != NULL); + + nir_graph_t b = load_fixture("fixtures/lif_linear_graph.json"); + CHECK(nir_checkpoint_loads(dump, &b, NULL, &error)); + CHECK(error == NULL); + const float *bw1 = NULL; + const float *bb1 = NULL; + size_t bw1n = 0; + size_t bb1n = 0; + CHECK(nir_graph_get_weights(&b, "fc1", &bw1, &bw1n, &bb1, &bb1n)); + CHECK(memcmp(bw1, w1, w1n * sizeof(float)) == 0); + + free(grad); + free(dump); + nir_sgd_free(&opt); + nir_graph_free(&a); + nir_graph_free(&b); +} + +static void check_opt_rejects(void) { + char *error = NULL; + nir_graph_t g = load_fixture("fixtures/lif_linear_graph.json"); + const float *w1 = NULL; + const float *b1 = NULL; + const float *w2 = NULL; + const float *b2 = NULL; + size_t w1n = 0; + size_t b1n = 0; + size_t w2n = 0; + size_t b2n = 0; + CHECK(nir_graph_get_weights(&g, "fc1", &w1, &w1n, &b1, &b1n)); + CHECK(nir_graph_get_weights(&g, "fc2", &w2, &w2n, &b2, &b2n)); + nir_sgd_t sw; + nir_adam_t aw; + nir_adam_t aw2; + nir_adam_t ab2; + CHECK(nir_sgd_init(&sw, w1n, 0.1f, 0.9f, 0.0f, &error)); + CHECK(nir_adam_init( + &aw, w1n, 0.01f, 0.9f, 0.999f, 1e-8f, 0.0f, &error + )); + CHECK(nir_adam_init( + &aw2, w2n, 0.01f, 0.9f, 0.999f, 1e-8f, 0.0f, &error + )); + CHECK(nir_adam_init( + &ab2, b2n, 0.01f, 0.9f, 0.999f, 1e-8f, 0.0f, &error + )); + + /* Save rejects: unknown node, unweighted node, duplicates, + * mixed kinds, and count mismatches. */ + nir_checkpoint_opt_t bad; + memset(&bad, 0, sizeof(bad)); + bad.node = "nope"; + bad.kind = NIR_CHECKPOINT_OPT_SGD; + bad.sgd_weight = &sw; + CHECK(nir_checkpoint_dumps_with_opt( + &g, NULL, &bad, 1, &error + ) == NULL); + CHECK(error != NULL); + free(error); + error = NULL; + bad.node = "lif1"; + CHECK(nir_checkpoint_dumps_with_opt( + &g, NULL, &bad, 1, &error + ) == NULL); + CHECK(error != NULL); + free(error); + error = NULL; + nir_checkpoint_opt_t dup[2]; + memset(dup, 0, sizeof(dup)); + dup[0].node = "fc1"; + dup[0].kind = NIR_CHECKPOINT_OPT_SGD; + dup[0].sgd_weight = &sw; + dup[1].node = "fc1"; + dup[1].kind = NIR_CHECKPOINT_OPT_SGD; + dup[1].sgd_weight = &sw; + CHECK(nir_checkpoint_dumps_with_opt( + &g, NULL, dup, 2, &error + ) == NULL); + CHECK(error != NULL); + free(error); + error = NULL; + nir_checkpoint_opt_t mixed; + memset(&mixed, 0, sizeof(mixed)); + mixed.node = "fc1"; + mixed.kind = NIR_CHECKPOINT_OPT_SGD; + mixed.sgd_weight = &sw; + mixed.adam_weight = &aw; + CHECK(nir_checkpoint_dumps_with_opt( + &g, NULL, &mixed, 1, &error + ) == NULL); + CHECK(error != NULL); + free(error); + error = NULL; + nir_sgd_t wrong; + CHECK(nir_sgd_init(&wrong, w1n + 1, 0.1f, 0.0f, 0.0f, &error)); + nir_checkpoint_opt_t mismatch; + memset(&mismatch, 0, sizeof(mismatch)); + mismatch.node = "fc1"; + mismatch.kind = NIR_CHECKPOINT_OPT_SGD; + mismatch.sgd_weight = &wrong; + CHECK(nir_checkpoint_dumps_with_opt( + &g, NULL, &mismatch, 1, &error + ) == NULL); + CHECK(error != NULL); + free(error); + error = NULL; + nir_sgd_free(&wrong); + + /* A valid Adam dump for load-reject hacks. */ + nir_checkpoint_opt_t good[2]; + memset(good, 0, sizeof(good)); + good[0].node = "fc1"; + good[0].kind = NIR_CHECKPOINT_OPT_ADAM; + good[0].adam_weight = &aw; + good[1].node = "fc2"; + good[1].kind = NIR_CHECKPOINT_OPT_ADAM; + good[1].adam_weight = &aw2; + good[1].adam_bias = &ab2; + float *gw1 = copy_buf(w1, w1n); + fill_grad(gw1, w1n, 0.1f); + CHECK(nir_adam_step(&aw, (float *)w1, gw1, w1n)); + free(gw1); + char *dump = nir_checkpoint_dumps_with_opt( + &g, NULL, good, 2, &error + ); + CHECK(error == NULL && dump != NULL); + + /* Load rejects a kind mismatch (file adam vs caller sgd). */ + nir_graph_t h = load_fixture("fixtures/lif_linear_graph.json"); + nir_sgd_t rs; + CHECK(nir_sgd_init(&rs, w1n, 0.1f, 0.9f, 0.0f, &error)); + nir_checkpoint_opt_t as_sgd[1]; + memset(as_sgd, 0, sizeof(as_sgd)); + as_sgd[0].node = "fc1"; + as_sgd[0].kind = NIR_CHECKPOINT_OPT_SGD; + as_sgd[0].sgd_weight = &rs; + /* Caller omits fc2 as well: file fc1 hits the kind check + * first only if iteration meets fc1 first; to force the kind + * path, pass both nodes with fc1 mistyped. */ + nir_checkpoint_opt_t kind_mix[2]; + memset(kind_mix, 0, sizeof(kind_mix)); + kind_mix[0].node = "fc1"; + kind_mix[0].kind = NIR_CHECKPOINT_OPT_SGD; + kind_mix[0].sgd_weight = &rs; + kind_mix[1].node = "fc2"; + kind_mix[1].kind = NIR_CHECKPOINT_OPT_ADAM; + kind_mix[1].adam_weight = &aw2; + kind_mix[1].adam_bias = &ab2; + CHECK(!nir_checkpoint_loads_with_opt( + dump, &h, NULL, kind_mix, 2, &error + )); + CHECK(error != NULL); + free(error); + error = NULL; + + /* Load rejects file entries with no caller entry. */ + CHECK(!nir_checkpoint_loads_with_opt( + dump, &h, NULL, as_sgd, 1, &error + )); + CHECK(error != NULL); + free(error); + error = NULL; + nir_sgd_free(&rs); + + /* Load rejects an unknown optimizer node. */ + char *hacked = copy_text(dump); + char *fc2key = strstr(hacked, "\"fc2\":{\"kind\""); + CHECK(fc2key != NULL); + memcpy(fc2key + 1, "nope", 4); + nir_graph_t u = load_fixture("fixtures/lif_linear_graph.json"); + nir_adam_t u1; + nir_adam_t u2; + nir_adam_t ub2; + CHECK(nir_adam_init( + &u1, w1n, 0.01f, 0.9f, 0.999f, 1e-8f, 0.0f, &error + )); + CHECK(nir_adam_init( + &u2, w2n, 0.01f, 0.9f, 0.999f, 1e-8f, 0.0f, &error + )); + CHECK(nir_adam_init( + &ub2, b2n, 0.01f, 0.9f, 0.999f, 1e-8f, 0.0f, &error + )); + nir_checkpoint_opt_t uopts[2]; + memset(uopts, 0, sizeof(uopts)); + uopts[0].node = "fc1"; + uopts[0].kind = NIR_CHECKPOINT_OPT_ADAM; + uopts[0].adam_weight = &u1; + uopts[1].node = "fc2"; + uopts[1].kind = NIR_CHECKPOINT_OPT_ADAM; + uopts[1].adam_weight = &u2; + uopts[1].adam_bias = &ub2; + CHECK(!nir_checkpoint_loads_with_opt( + hacked, &u, NULL, uopts, 2, &error + )); + CHECK(error != NULL); + free(error); + error = NULL; + free(hacked); + + /* Load rejects a truncated optimizer buffer. */ + hacked = copy_text(dump); + char *vel = strstr(hacked, "\"m\":["); + CHECK(vel != NULL); + char *end = strchr(vel, ']'); + CHECK(end != NULL); + memmove(vel + 5, end, strlen(end) + 1); + memcpy(vel + 5, "0.0]", 4); + CHECK(!nir_checkpoint_loads_with_opt( + hacked, &u, NULL, uopts, 2, &error + )); + CHECK(error != NULL); + free(error); + error = NULL; + free(hacked); + + /* Load rejects an invalid Adam step. */ + hacked = copy_text(dump); + char *step = strstr(hacked, "\"step\":1"); + CHECK(step != NULL); + step[strlen("\"step\":")] = '-'; + CHECK(!nir_checkpoint_loads_with_opt( + hacked, &u, NULL, uopts, 2, &error + )); + CHECK(error != NULL); + free(error); + error = NULL; + free(hacked); + + free(dump); + nir_sgd_free(&sw); + nir_adam_free(&aw); + nir_adam_free(&aw2); + nir_adam_free(&ab2); + nir_adam_free(&u1); + nir_adam_free(&u2); + nir_adam_free(&ub2); + nir_graph_free(&g); + nir_graph_free(&h); + nir_graph_free(&u); +} + int main(void) { check_writer(); @@ -169,7 +806,7 @@ int main(void) { CHECK(error == NULL); CHECK(dump != NULL); CHECK(strstr(dump, "\"format\":\"spikeforge-c-checkpoint\"") != NULL); - CHECK(strstr(dump, "\"version\":1") != NULL); + CHECK(strstr(dump, "\"version\":2") != NULL); CHECK(strstr(dump, "\"topology\":\"lif_linear\"") != NULL); CHECK(nir_graph_set_weights(&graph, "fc2", zeros, wn, zeros, bn, &error)); @@ -205,21 +842,21 @@ int main(void) { /* Rejects: bad version, bad format, missing node, bad shape. */ char hacked[65536]; strcpy(hacked, dump); - char *ver = strstr(hacked, "\"version\":1"); + char *ver = strstr(hacked, "\"version\":2"); CHECK(ver != NULL); ver[strlen("\"version\":")] = '9'; CHECK(!nir_checkpoint_loads(hacked, &graph, NULL, &error)); CHECK(error != NULL); free(error); error = NULL; - CHECK(!nir_checkpoint_loads("{\"format\":\"nope\",\"version\":1," + CHECK(!nir_checkpoint_loads("{\"format\":\"nope\",\"version\":2," "\"nodes\":[]}", &graph, NULL, &error)); CHECK(error != NULL); free(error); error = NULL; CHECK(!nir_checkpoint_loads("{\"format\":" "\"spikeforge-c-checkpoint\",\"version\":1,\"nodes\":[]}", - &graph, NULL, &error)); /* drops every weighted node */ + &graph, NULL, &error)); /* v1 version ok, drops every node */ CHECK(error != NULL); free(error); error = NULL; @@ -263,6 +900,12 @@ int main(void) { free(summary); free(dump); + check_opt_sgd_resume(); + check_opt_adam_resume(); + check_opt_v1_compat(); + check_opt_plain_ignores(); + check_opt_rejects(); + nir_graph_free(&graph); return 0; } diff --git a/tools/patch_nir_weights.py b/tools/patch_nir_weights.py index f418418..fba0bd0 100644 --- a/tools/patch_nir_weights.py +++ b/tools/patch_nir_weights.py @@ -43,10 +43,12 @@ def patch(summary: dict, checkpoint: dict) -> dict: """Return a copy of `summary` with checkpoint weights applied.""" if checkpoint.get("format") != "spikeforge-c-checkpoint": raise ValueError("not a spikeforge-c-checkpoint document") - if checkpoint.get("version") != 1: + if checkpoint.get("version") not in (1, 2): raise ValueError( f"unsupported checkpoint version {checkpoint.get('version')!r}" ) + # v2 optimizers are intentionally ignored: the reverse bridge + # carries weights only. nodes = {node["name"]: node for node in summary["nodes"]} patched = 0 for entry in checkpoint["nodes"]: