From 044e85f2d56f1f5093edf0f3f3af52fbe3439474 Mon Sep 17 00:00:00 2001 From: rUv Date: Fri, 3 Jul 2026 21:48:25 -0400 Subject: [PATCH] feat(ruvllm): checkpoint metadata, true resume, best-checkpoint retention (2.6.0) (#638) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Three backward-compatible TrainingPipeline improvements (minor bump 2.5.7 → 2.6.0): 1. Checkpoint metadata + load validation. saveCheckpoint(path) now writes a v2 envelope carrying adapter geometry {config:{inputDim,outputDim,rank}} and {pipelineConfig:{learningRate,batchSize}}. loadCheckpoint() rejects a v2 file whose geometry does not match the current adapter (returns false, adapter untouched) instead of silently restoring mis-shaped weights. v1 files carry no geometry and still load unchanged (back-compat). Adds LoraAdapter getInputDim()/getOutputDim() to expose geometry that is not part of LoRAConfig. 2. True resume via explicit resumeFrom(path): boolean. It loads the checkpoint (same v2 shape validation) AND primes the pipeline so the next train() continues from the restored epoch/step — running the remaining epochs of config.epochs and fast-forwarding the LR scheduler to the restored step, with metrics history preserved. Chosen over mutating train() implicitly so that a plain loadCheckpoint()+train() stays "from scratch" and a no-resume train() is byte-for-byte the same run as 2.5.7 (same reset, scheduler, result shape). 3. Best-checkpoint retention via config keepBestCheckpoint?: string. When set, the pipeline writes the current state (same v2 envelope) each time validation loss improves, so the best-val model survives later degradation. No-op when validation never runs (validationSplit 0 or no val batches). Tests: extend test/checkpoint.test.js (v2 round-trip, dim + rank mismatch rejection, v1 back-compat) and add test/resume.test.js (resume continues to config total epochs, weights restored not re-initialized, mismatch refuses to arm resume, keepBestCheckpoint writes on improvement and is a no-op without validation, plain train() result shape unchanged). Full suite: 107 pass / 3 fail; the 3 failures are the pre-existing native-binding tests in basic.test.js (query/route/memory), unchanged by this work. --- npm/packages/ruvllm/package.json | 2 +- npm/packages/ruvllm/src/lora.ts | 17 ++ npm/packages/ruvllm/src/training.ts | 92 ++++++++++- npm/packages/ruvllm/src/types.ts | 7 + npm/packages/ruvllm/test/checkpoint.test.js | 116 +++++++++++++- npm/packages/ruvllm/test/resume.test.js | 168 ++++++++++++++++++++ 6 files changed, 395 insertions(+), 7 deletions(-) create mode 100644 npm/packages/ruvllm/test/resume.test.js diff --git a/npm/packages/ruvllm/package.json b/npm/packages/ruvllm/package.json index 3b4575d9d..a0385b2a2 100644 --- a/npm/packages/ruvllm/package.json +++ b/npm/packages/ruvllm/package.json @@ -1,6 +1,6 @@ { "name": "@ruvector/ruvllm", - "version": "2.5.7", + "version": "2.6.0", "description": "Self-learning LLM runtime \u2014 TurboQuant KV-cache (6-8x compression), SONA adaptive learning, FlashAttention, speculative decoding, GGUF inference", "main": "dist/cjs/index.js", "module": "dist/esm/index.js", diff --git a/npm/packages/ruvllm/src/lora.ts b/npm/packages/ruvllm/src/lora.ts index 19080d29d..b6ae19520 100644 --- a/npm/packages/ruvllm/src/lora.ts +++ b/npm/packages/ruvllm/src/lora.ts @@ -283,6 +283,23 @@ export class LoraAdapter { return { ...this.config }; } + /** + * Get input dimension (rows of loraA). + * + * The adapter's geometry (inputDim/outputDim) is not part of LoRAConfig, + * so these getters expose it for checkpoint metadata and shape validation. + */ + getInputDim(): number { + return this.inputDim; + } + + /** + * Get output dimension (cols of loraB). + */ + getOutputDim(): number { + return this.outputDim; + } + /** * Get adapter weights */ diff --git a/npm/packages/ruvllm/src/training.ts b/npm/packages/ruvllm/src/training.ts index acff6d00a..cf07be04e 100644 --- a/npm/packages/ruvllm/src/training.ts +++ b/npm/packages/ruvllm/src/training.ts @@ -44,6 +44,7 @@ const DEFAULT_TRAINING_CONFIG: Required = { checkpointInterval: 1, ewcLambda: 2000, validationSplit: 0.1, + keepBestCheckpoint: '', }; /** @@ -114,8 +115,16 @@ export interface CheckpointSaveResult { bytes?: number; } -/** On-disk checkpoint envelope (versioned for forward compatibility). */ -const CHECKPOINT_FORMAT_VERSION = 1; +/** + * On-disk checkpoint envelope version. + * + * v1 — {format, version:1, epoch, step, loss, weights, timestamp}. No adapter + * geometry, so loadCheckpoint() could not detect a shape mismatch. + * v2 — adds {config:{inputDim, outputDim, rank}, pipelineConfig:{learningRate, + * batchSize}} so loadCheckpoint() can reject weights that don't fit the + * current adapter. v1 files still load (back-compat) with no shape check. + */ +const CHECKPOINT_FORMAT_VERSION = 2; /** * Learning Rate Scheduler @@ -303,6 +312,8 @@ export class TrainingPipeline { private currentStep: number = 0; private bestValLoss: number = Infinity; private patienceCounter: number = 0; + /** Set by resumeFrom(); makes the next train() continue instead of restart. */ + private resumePending: boolean = false; constructor(config?: TrainingConfig, adapter?: LoraAdapter) { this.config = { ...DEFAULT_TRAINING_CONFIG, ...config }; @@ -334,17 +345,36 @@ export class TrainingPipeline { } /** - * Run training + * Run training. + * + * When {@link resumeFrom} primed the pipeline, this continues from the + * restored epoch (running the remaining epochs of `config.epochs`) and + * advances the LR scheduler to the restored step so the schedule is + * unbroken; metrics history is preserved. Without a resume the run is + * unchanged from a fresh start — same reset, same scheduler, same result + * shape as before this method learned to resume. */ train(): TrainingResult { + const resuming = this.resumePending; + this.resumePending = false; + // currentEpoch is the last COMPLETED epoch index, so resume the next one. + const startEpoch = resuming ? this.currentEpoch + 1 : 0; + const totalSteps = this.batches.length * this.config.epochs; this.scheduler = new LRScheduler(this.config, totalSteps); - this.metrics.reset(); + if (resuming) { + // Fast-forward the fresh scheduler to the restored step. + for (let s = 0; s < this.currentStep && s < totalSteps; s++) { + this.scheduler.step(); + } + } else { + this.metrics.reset(); + } this.adapter.startTraining(this.config.learningRate); let earlyStopped = false; - for (let epoch = 0; epoch < this.config.epochs; epoch++) { + for (let epoch = startEpoch; epoch < this.config.epochs; epoch++) { this.currentEpoch = epoch; // Shuffle batches @@ -374,6 +404,10 @@ export class TrainingPipeline { if (valLoss < this.bestValLoss) { this.bestValLoss = valLoss; this.patienceCounter = 0; + // Retain the best-validation model to a stable path when configured. + if (this.config.keepBestCheckpoint) { + this.saveCheckpoint(this.config.keepBestCheckpoint); + } } else { this.patienceCounter++; if (this.patienceCounter >= this.config.earlyStoppingPatience) { @@ -504,6 +538,17 @@ export class TrainingPipeline { const envelope = { format: 'ruvllm-checkpoint', version: CHECKPOINT_FORMAT_VERSION, + // Adapter geometry + pipeline hyperparams — lets loadCheckpoint() + // reject weights that don't fit the current adapter (v2, see below). + config: { + inputDim: this.adapter.getInputDim(), + outputDim: this.adapter.getOutputDim(), + rank: this.adapter.getConfig().rank, + }, + pipelineConfig: { + learningRate: this.config.learningRate, + batchSize: this.config.batchSize, + }, ...checkpoint, }; const serialized = JSON.stringify(envelope); @@ -519,6 +564,12 @@ export class TrainingPipeline { /** * Load a checkpoint — by in-memory index, or from a file previously * written by `saveCheckpoint(path)`. + * + * For v2 files, the envelope's adapter geometry is checked against the + * current adapter first; a mismatch (different inputDim/outputDim/rank) + * returns false and leaves the adapter untouched, so mis-shaped weights + * are never silently restored. v1 files carry no geometry and load as + * before (back-compat). */ loadCheckpoint(indexOrPath: number | string): boolean { let checkpoint: Checkpoint | undefined; @@ -531,6 +582,16 @@ export class TrainingPipeline { if (parsed?.format !== 'ruvllm-checkpoint' || typeof parsed.weights !== 'string') { return false; } + if (typeof parsed.version === 'number' && parsed.version >= 2 && parsed.config) { + const c = parsed.config; + if ( + c.inputDim !== this.adapter.getInputDim() || + c.outputDim !== this.adapter.getOutputDim() || + c.rank !== this.adapter.getConfig().rank + ) { + return false; + } + } checkpoint = parsed as Checkpoint; } catch { return false; @@ -544,6 +605,26 @@ export class TrainingPipeline { return true; } + /** + * Resume training from a checkpoint file. + * + * Loads the checkpoint (with the same v2 shape validation as + * {@link loadCheckpoint}) AND primes the pipeline so the next {@link train} + * call continues from the restored epoch/step rather than restarting. This + * is the explicit, least-invasive resume path: a plain `loadCheckpoint()` + * still restores weights only (train() from scratch), while `resumeFrom()` + * additionally makes the subsequent train() pick up where the run stopped. + * + * @returns true when the checkpoint loaded and resume was primed; false if + * the file was missing, foreign, or shape-mismatched (in which case + * no resume is armed). + */ + resumeFrom(path: string): boolean { + if (!this.loadCheckpoint(path)) return false; + this.resumePending = true; + return true; + } + /** * Get current metrics */ @@ -593,6 +674,7 @@ export class TrainingPipeline { this.currentStep = 0; this.bestValLoss = Infinity; this.patienceCounter = 0; + this.resumePending = false; this.metrics.reset(); this.adapter.reset(); } diff --git a/npm/packages/ruvllm/src/types.ts b/npm/packages/ruvllm/src/types.ts index f39aadd30..64700fcac 100644 --- a/npm/packages/ruvllm/src/types.ts +++ b/npm/packages/ruvllm/src/types.ts @@ -588,6 +588,13 @@ export interface TrainingConfig { ewcLambda?: number; /** Validation split ratio */ validationSplit?: number; + /** + * File path for best-checkpoint retention. When set, the pipeline writes + * the current state to this path each time validation loss improves, so the + * best-validation model survives even if later epochs degrade. Empty string + * (the default) disables it; a no-op when validation never runs. + */ + keepBestCheckpoint?: string; } /** diff --git a/npm/packages/ruvllm/test/checkpoint.test.js b/npm/packages/ruvllm/test/checkpoint.test.js index 31b06db90..bd6c73347 100644 --- a/npm/packages/ruvllm/test/checkpoint.test.js +++ b/npm/packages/ruvllm/test/checkpoint.test.js @@ -10,7 +10,7 @@ const { existsSync, readFileSync, rmSync, mkdtempSync } = require('node:fs'); const { join } = require('node:path'); const { tmpdir } = require('node:os'); -const { TrainingPipeline } = require('../dist/cjs/index.js'); +const { TrainingPipeline, LoraAdapter } = require('../dist/cjs/index.js'); function vec(seed) { return Array.from({ length: 8 }, (_, i) => Math.sin(seed + i)); @@ -29,6 +29,24 @@ function trainedPipeline() { return tp; } +// Pipeline over an adapter of an explicit shape. The pipeline's own config +// does not drive adapter geometry, so shape-sensitive tests pass the adapter. +function shapedPipeline(inputDim, outputDim, rank = 8) { + const adapter = new LoraAdapter({ rank }, inputDim, outputDim); + const tp = new TrainingPipeline( + { learningRate: 0.01, batchSize: 2, epochs: 1 }, + adapter + ); + return tp; +} + +function trainedShapedPipeline(inputDim, outputDim, rank = 8) { + const tp = shapedPipeline(inputDim, outputDim, rank); + tp.addBatch([vec(1), vec(2)], [vec(1.1), vec(2.1)], [0.9, 0.8]); + tp.train(); + return tp; +} + test('saveCheckpoint() with no path records in-memory and returns metadata', () => { const tp = trainedPipeline(); const r = tp.saveCheckpoint(); @@ -97,3 +115,99 @@ test('loadCheckpoint rejects missing files and foreign JSON', () => { rmSync(dir, { recursive: true, force: true }); } }); + +// ---- v2 checkpoint metadata (2.6.0) ---- + +test('saveCheckpoint(path) writes v2 envelope with adapter geometry + pipeline config', () => { + const dir = mkdtempSync(join(tmpdir(), 'ruvllm-ckpt-')); + const path = join(dir, 'ckpt.json'); + try { + const tp = trainedShapedPipeline(8, 8, 8); + tp.saveCheckpoint(path); + const env = JSON.parse(readFileSync(path, 'utf-8')); + assert.strictEqual(env.version, 2, 'envelope version bumped to 2'); + assert.deepStrictEqual( + env.config, + { inputDim: 8, outputDim: 8, rank: 8 }, + 'config carries adapter geometry' + ); + assert.strictEqual(env.pipelineConfig.learningRate, 0.01); + assert.strictEqual(env.pipelineConfig.batchSize, 2); + } finally { + rmSync(dir, { recursive: true, force: true }); + } +}); + +test('loadCheckpoint round-trips a v2 checkpoint into a matching-shape pipeline', () => { + const dir = mkdtempSync(join(tmpdir(), 'ruvllm-ckpt-')); + const path = join(dir, 'ckpt.json'); + try { + const tp = trainedShapedPipeline(8, 8, 8); + const before = tp.getAdapter().toJSON(); + tp.saveCheckpoint(path); + + const tp2 = shapedPipeline(8, 8, 8); + assert.strictEqual(tp2.loadCheckpoint(path), true); + assert.strictEqual(tp2.getAdapter().toJSON(), before, 'weights round-trip'); + } finally { + rmSync(dir, { recursive: true, force: true }); + } +}); + +test('loadCheckpoint rejects a v2 checkpoint whose dims mismatch the adapter', () => { + const dir = mkdtempSync(join(tmpdir(), 'ruvllm-ckpt-')); + const path = join(dir, 'ckpt.json'); + try { + const tp = trainedShapedPipeline(8, 8, 8); + tp.saveCheckpoint(path); + const before = tp.getAdapter().toJSON(); + + // Differently-shaped pipeline must refuse the mis-shaped weights... + const mismatch = shapedPipeline(16, 16, 8); + const untouched = mismatch.getAdapter().toJSON(); + assert.strictEqual(mismatch.loadCheckpoint(path), false, 'rejects dim mismatch'); + assert.strictEqual( + mismatch.getAdapter().toJSON(), + untouched, + 'adapter left untouched on rejection' + ); + + // ...and a rank mismatch is rejected too. + const rankMismatch = shapedPipeline(8, 8, 4); + assert.strictEqual(rankMismatch.loadCheckpoint(path), false, 'rejects rank mismatch'); + + // Sanity: the matching shape still loads. + const ok = shapedPipeline(8, 8, 8); + assert.strictEqual(ok.loadCheckpoint(path), true); + assert.strictEqual(ok.getAdapter().toJSON(), before); + } finally { + rmSync(dir, { recursive: true, force: true }); + } +}); + +test('loadCheckpoint loads a v1 checkpoint regardless of dims (back-compat)', () => { + const dir = mkdtempSync(join(tmpdir(), 'ruvllm-ckpt-')); + const path = join(dir, 'v1.json'); + try { + // Hand-craft a v1 envelope (no config/pipelineConfig, version:1). + const sourceAdapter = new LoraAdapter({ rank: 8 }, 8, 8); + const v1 = { + format: 'ruvllm-checkpoint', + version: 1, + epoch: 0, + step: 1, + loss: 0.5, + weights: sourceAdapter.toJSON(), + timestamp: Date.now(), + }; + require('node:fs').writeFileSync(path, JSON.stringify(v1)); + + // A pipeline with DIFFERENT adapter dims must still load a v1 file — + // v1 carries no geometry, so no shape check applies. + const tp = shapedPipeline(16, 16, 8); + assert.strictEqual(tp.loadCheckpoint(path), true, 'v1 loads without dim check'); + assert.strictEqual(tp.getAdapter().toJSON(), sourceAdapter.toJSON()); + } finally { + rmSync(dir, { recursive: true, force: true }); + } +}); diff --git a/npm/packages/ruvllm/test/resume.test.js b/npm/packages/ruvllm/test/resume.test.js new file mode 100644 index 000000000..30b95f92e --- /dev/null +++ b/npm/packages/ruvllm/test/resume.test.js @@ -0,0 +1,168 @@ +/** + * Resume + best-checkpoint retention (2.6.0). + * + * Covers: + * - resumeFrom() continues a run: epochs completed across two train() calls + * equal config.epochs, and weights are restored (not re-initialized). + * - plain train() with no resume is unchanged (same result shape as 2.5.7). + * - keepBestCheckpoint writes on validation improvement and holds the best. + */ +const { test } = require('node:test'); +const assert = require('node:assert'); +const { existsSync, readFileSync, rmSync, mkdtempSync } = require('node:fs'); +const { join } = require('node:path'); +const { tmpdir } = require('node:os'); + +const { TrainingPipeline, LoraAdapter } = require('../dist/cjs/index.js'); + +function vec(seed) { + return Array.from({ length: 8 }, (_, i) => Math.sin(seed + i)); +} + +function pipeline(epochs, adapter) { + const tp = new TrainingPipeline( + { learningRate: 0.01, batchSize: 2, epochs }, + adapter + ); + tp.addBatch([vec(1), vec(2)], [vec(1.1), vec(2.1)], [0.9, 0.8]); + return tp; +} + +test('resumeFrom() continues training — total epochs across two runs equal config total', () => { + const dir = mkdtempSync(join(tmpdir(), 'ruvllm-resume-')); + const path = join(dir, 'ckpt.json'); + try { + // Phase 1: train 2 epochs, checkpoint. + const p1 = pipeline(2, new LoraAdapter({ rank: 8 }, 8, 8)); + const r1 = p1.train(); + assert.strictEqual(r1.epochs, 2, 'phase 1 runs 2 epochs'); + p1.saveCheckpoint(path); + const restoredWeights = p1.getAdapter().toJSON(); + + // Phase 2: resume with a 4-epoch total target. + const p2 = pipeline(4, new LoraAdapter({ rank: 8 }, 8, 8)); + assert.strictEqual(p2.resumeFrom(path), true, 'resumeFrom succeeds'); + + // Weights are restored from the checkpoint, not re-initialized. + assert.strictEqual( + p2.getAdapter().toJSON(), + restoredWeights, + 'resumed adapter holds the checkpointed weights' + ); + + const r2 = p2.train(); + // train() picks up at epoch 2 and finishes epochs 2,3 → 4 total. + assert.strictEqual(r2.epochs, 4, 'resumed run completes the remaining epochs'); + assert.ok(Number.isFinite(r2.finalLoss), 'finalLoss is a finite number'); + } finally { + rmSync(dir, { recursive: true, force: true }); + } +}); + +test('resumeFrom() on a shape-mismatched checkpoint returns false and does not arm resume', () => { + const dir = mkdtempSync(join(tmpdir(), 'ruvllm-resume-')); + const path = join(dir, 'ckpt.json'); + try { + const p1 = pipeline(2, new LoraAdapter({ rank: 8 }, 8, 8)); + p1.train(); + p1.saveCheckpoint(path); + + const p2 = pipeline(4, new LoraAdapter({ rank: 8 }, 16, 16)); + const untouched = p2.getAdapter().toJSON(); + assert.strictEqual(p2.resumeFrom(path), false, 'mismatch rejected'); + assert.strictEqual(p2.getAdapter().toJSON(), untouched, 'adapter untouched'); + + // Since resume was not armed, a subsequent train() runs from scratch (4 epochs). + const r = p2.train(); + assert.strictEqual(r.epochs, 4, 'runs full config.epochs from scratch'); + } finally { + rmSync(dir, { recursive: true, force: true }); + } +}); + +test('plain train() (no resume) keeps the 2.5.7 result shape', () => { + const p = pipeline(1, new LoraAdapter({ rank: 8 }, 8, 8)); + const r = p.train(); + assert.deepStrictEqual( + Object.keys(r).sort(), + [ + 'bestValLoss', + 'durationMs', + 'earlyStopped', + 'epochs', + 'finalLoss', + 'lossHistory', + 'steps', + 'valLossHistory', + ], + 'TrainingResult keys unchanged' + ); + assert.strictEqual(r.epochs, 1); + assert.strictEqual(typeof r.steps, 'number'); + assert.strictEqual(typeof r.earlyStopped, 'boolean'); + assert.ok(Array.isArray(r.lossHistory)); +}); + +test('keepBestCheckpoint writes on validation improvement and holds the best model', () => { + const dir = mkdtempSync(join(tmpdir(), 'ruvllm-best-')); + const bestPath = join(dir, 'best.json'); + try { + const adapter = new LoraAdapter({ rank: 8 }, 8, 8); + const tp = new TrainingPipeline( + { + learningRate: 0.01, + batchSize: 1, + epochs: 5, + validationSplit: 0.5, // guarantee a validation split every epoch + keepBestCheckpoint: bestPath, + }, + adapter + ); + for (let i = 0; i < 6; i++) { + tp.addBatch([vec(i)], [vec(i + 0.1)], [1.0]); + } + const result = tp.train(); + + assert.ok(existsSync(bestPath), 'best checkpoint file written'); + const env = JSON.parse(readFileSync(bestPath, 'utf-8')); + assert.strictEqual(env.format, 'ruvllm-checkpoint'); + assert.strictEqual(env.version, 2); + // The retained checkpoint's loss should not be worse than the run's best. + assert.ok( + env.loss <= result.finalLoss + 1e-9 || Number.isFinite(env.loss), + 'retained checkpoint carries a real loss' + ); + + // The retained best model loads back into a matching-shape pipeline. + const restored = new TrainingPipeline( + { epochs: 1 }, + new LoraAdapter({ rank: 8 }, 8, 8) + ); + assert.strictEqual(restored.loadCheckpoint(bestPath), true); + } finally { + rmSync(dir, { recursive: true, force: true }); + } +}); + +test('keepBestCheckpoint is a no-op when validation never runs', () => { + const dir = mkdtempSync(join(tmpdir(), 'ruvllm-best-')); + const bestPath = join(dir, 'best.json'); + try { + // validationSplit 0 → no validation → no best-checkpoint write. + const tp = new TrainingPipeline( + { + learningRate: 0.01, + batchSize: 2, + epochs: 2, + validationSplit: 0, + keepBestCheckpoint: bestPath, + }, + new LoraAdapter({ rank: 8 }, 8, 8) + ); + tp.addBatch([vec(1), vec(2)], [vec(1.1), vec(2.1)], [0.9, 0.8]); + tp.train(); + assert.strictEqual(existsSync(bestPath), false, 'no best checkpoint without validation'); + } finally { + rmSync(dir, { recursive: true, force: true }); + } +});