@simulatte/doppler 0.1.5 → 0.1.7

package/src/tooling/source-runtime-materializer.js

@@ -0,0 +1,93 @@
+import path from 'node:path';
+
+import {
+  DIRECT_SOURCE_PATH_ARTIFACT_RELATIVE,
+  DIRECT_SOURCE_RUNTIME_MODE,
+  DIRECT_SOURCE_RUNTIME_SCHEMA,
+  DIRECT_SOURCE_RUNTIME_SCHEMA_VERSION,
+  getSourceRuntimeMetadata,
+} from './source-runtime-bundle.js';
+
+function cloneJsonValue(value) {
+  if (typeof structuredClone === 'function') {
+    return structuredClone(value);
+  }
+  return JSON.parse(JSON.stringify(value));
+}
+
+function toRelativeArtifactPath(value, artifactDir, label) {
+  const raw = String(value || '').trim();
+  if (!raw) {
+    throw new Error(`${label} path is required.`);
+  }
+  const resolvedArtifactDir = path.resolve(artifactDir);
+  const resolvedTarget = path.resolve(raw);
+  const relativePath = path.relative(resolvedArtifactDir, resolvedTarget).replace(/\\/g, '/');
+  if (!relativePath || relativePath.startsWith('../') || relativePath === '..') {
+    throw new Error(
+      `${label} "${raw}" must live inside artifactDir "${resolvedArtifactDir}" for a persisted direct-source manifest.`
+    );
+  }
+  return relativePath;
+}
+
+export function materializeSourceRuntimeManifest(manifest, artifactDir) {
+  const sourceRuntime = getSourceRuntimeMetadata(manifest);
+  if (!sourceRuntime) {
+    throw new Error('materializeSourceRuntimeManifest requires manifest.metadata.sourceRuntime.');
+  }
+  const resolvedArtifactDir = String(artifactDir || '').trim();
+  if (!resolvedArtifactDir) {
+    throw new Error('materializeSourceRuntimeManifest requires artifactDir.');
+  }
+
+  const nextManifest = cloneJsonValue(manifest);
+  if (!nextManifest.metadata || typeof nextManifest.metadata !== 'object') {
+    nextManifest.metadata = {};
+  }
+  const sourceMetadata = nextManifest.metadata.sourceRuntime && typeof nextManifest.metadata.sourceRuntime === 'object'
+    ? cloneJsonValue(nextManifest.metadata.sourceRuntime)
+    : {};
+
+  sourceMetadata.mode = DIRECT_SOURCE_RUNTIME_MODE;
+  sourceMetadata.schema = DIRECT_SOURCE_RUNTIME_SCHEMA;
+  sourceMetadata.schemaVersion = DIRECT_SOURCE_RUNTIME_SCHEMA_VERSION;
+  sourceMetadata.hashAlgorithm = sourceRuntime.hashAlgorithm;
+  sourceMetadata.pathSemantics = DIRECT_SOURCE_PATH_ARTIFACT_RELATIVE;
+  sourceMetadata.sourceFiles = sourceRuntime.sourceFiles.map((entry) => ({
+    index: entry.index,
+    filename: entry.filename ?? null,
+    path: toRelativeArtifactPath(
+      entry.path,
+      resolvedArtifactDir,
+      `source runtime source file ${entry.index}`
+    ),
+    size: entry.size,
+    hash: entry.hash,
+    hashAlgorithm: entry.hashAlgorithm,
+  }));
+  sourceMetadata.auxiliaryFiles = sourceRuntime.auxiliaryFiles.map((entry) => ({
+    path: toRelativeArtifactPath(
+      entry.path,
+      resolvedArtifactDir,
+      `source runtime auxiliary file ${entry.kind}`
+    ),
+    size: entry.size,
+    hash: entry.hash,
+    hashAlgorithm: entry.hashAlgorithm,
+    kind: entry.kind,
+  }));
+  sourceMetadata.tokenizer = {
+    jsonPath: sourceRuntime.tokenizer.jsonPath
+      ? toRelativeArtifactPath(sourceRuntime.tokenizer.jsonPath, resolvedArtifactDir, 'source runtime tokenizer json')
+      : null,
+    configPath: sourceRuntime.tokenizer.configPath
+      ? toRelativeArtifactPath(sourceRuntime.tokenizer.configPath, resolvedArtifactDir, 'source runtime tokenizer config')
+      : null,
+    modelPath: sourceRuntime.tokenizer.modelPath
+      ? toRelativeArtifactPath(sourceRuntime.tokenizer.modelPath, resolvedArtifactDir, 'source runtime tokenizer model')
+      : null,
+  };
+  nextManifest.metadata.sourceRuntime = sourceMetadata;
+  return nextManifest;
+}

package/src/training/attention-backward.js

@@ -1,6 +1,7 @@
-import { acquireBuffer, uploadData, readBuffer } from '../memory/buffer-pool.js';
+import { acquireBuffer, uploadData, readBuffer, releaseBuffer } from '../memory/buffer-pool.js';
 import { createTensor, tensorBytes } from '../gpu/tensor.js';
 import { f16ToF32Array } from '../inference/kv-cache/types.js';
+import { createUploadedTensor } from './tensor-factory.js';
 
 function toFloat32(buffer, dtype) {
   if (dtype === 'f16') {
@@ -67,9 +68,7 @@ export async function buildAttentionSoftmaxCache(q, k, options) {
   const kData = toFloat32(kBuf, k.dtype);
   const sData = computeSoftmax(qData, kData, options);
   const { seqLen, numHeads } = options;
-  const outBuf = acquireBuffer(tensorBytes([numHeads, seqLen, seqLen], 'f32'), undefined, 'attn_softmax_cache');
-  uploadData(outBuf, sData);
-  return createTensor(outBuf, 'f32', [numHeads, seqLen, seqLen], 'attn_softmax_cache');
+  return createUploadedTensor(sData, 'f32', [numHeads, seqLen, seqLen], 'attn_softmax_cache');
 }
 
 export async function attentionBackwardCpu(
@@ -201,17 +200,33 @@ export async function attentionBackwardCpu(
     }
   }
 
-  const qBufOut = acquireBuffer(tensorBytes([seqLen, numHeads, headDim], 'f32'), undefined, 'attn_backward_q');
-  const kBufOut = acquireBuffer(tensorBytes([seqLen, numHeads, headDim], 'f32'), undefined, 'attn_backward_k');
-  const vBufOut = acquireBuffer(tensorBytes([seqLen, numHeads, headDim], 'f32'), undefined, 'attn_backward_v');
-
-  uploadData(qBufOut, dQ);
-  uploadData(kBufOut, dK);
-  uploadData(vBufOut, dV);
-
-  return {
-    gradQ: createTensor(qBufOut, 'f32', [seqLen, numHeads, headDim], 'attn_grad_q'),
-    gradK: createTensor(kBufOut, 'f32', [seqLen, numHeads, headDim], 'attn_grad_k'),
-    gradV: createTensor(vBufOut, 'f32', [seqLen, numHeads, headDim], 'attn_grad_v'),
-  };
+  let qBufOut = null;
+  let kBufOut = null;
+  let vBufOut = null;
+  try {
+    qBufOut = acquireBuffer(tensorBytes([seqLen, numHeads, headDim], 'f32'), undefined, 'attn_backward_q');
+    kBufOut = acquireBuffer(tensorBytes([seqLen, numHeads, headDim], 'f32'), undefined, 'attn_backward_k');
+    vBufOut = acquireBuffer(tensorBytes([seqLen, numHeads, headDim], 'f32'), undefined, 'attn_backward_v');
+
+    uploadData(qBufOut, dQ);
+    uploadData(kBufOut, dK);
+    uploadData(vBufOut, dV);
+
+    return {
+      gradQ: createTensor(qBufOut, 'f32', [seqLen, numHeads, headDim], 'attn_grad_q'),
+      gradK: createTensor(kBufOut, 'f32', [seqLen, numHeads, headDim], 'attn_grad_k'),
+      gradV: createTensor(vBufOut, 'f32', [seqLen, numHeads, headDim], 'attn_grad_v'),
+    };
+  } catch (error) {
+    if (qBufOut) {
+      releaseBuffer(qBufOut);
+    }
+    if (kBufOut) {
+      releaseBuffer(kBufOut);
+    }
+    if (vBufOut) {
+      releaseBuffer(vBufOut);
+    }
+    throw error;
+  }
 }

package/src/training/autograd.js

@@ -6,6 +6,7 @@ import { acquireBuffer, readBuffer, releaseBuffer, uploadData } from '../memory/
 import { createTensor } from '../gpu/tensor.js';
 import { attentionBackwardCpu } from './attention-backward.js';
 import { f16ToF32Array, f32ToF16Array } from '../inference/kv-cache/types.js';
+import { createUploadedTensor } from './tensor-factory.js';
 
 export const OpType = {
   EMBED: 'embed',
@@ -35,6 +36,7 @@ export class AutogradTape {
   constructor(registry) {
     this.registry = registry;
     this.records = [];
+    this.retainedBuffers = new Set();
   }
 
   watch(tensor) {
@@ -43,6 +45,13 @@ export class AutogradTape {
 
   async record(op, fn, inputs, options = {}) {
     const output = await fn(...inputs);
+    if (Array.isArray(options.retainBuffers)) {
+      for (const buffer of options.retainBuffers) {
+        if (buffer) {
+          this.retainedBuffers.add(buffer);
+        }
+      }
+    }
     this.records.push({ op, inputs, output, options });
     return output;
   }
@@ -50,31 +59,40 @@ export class AutogradTape {
   async backward(gradOutput) {
     const grads = new Map();
     const seeds = this.normalizeBackwardSeeds(gradOutput);
-    for (const seed of seeds) {
-      await this.accumulateGrad(grads, seed.tensor, seed.grad);
-    }
-
-    for (let i = this.records.length - 1; i >= 0; i -= 1) {
-      const record = this.records[i];
-      const entry = this.registry.ops[record.op];
-      if (!entry) {
-        continue;
+    try {
+      for (const seed of seeds) {
+        await this.accumulateGrad(grads, seed.tensor, seed.grad);
       }
 
-      const gradOut = grads.get(record.output);
-      if (!gradOut) {
-        continue;
-      }
+      for (let i = this.records.length - 1; i >= 0; i -= 1) {
+        const record = this.records[i];
+        const entry = this.registry.ops[record.op];
+        if (!entry) {
+          continue;
+        }
+
+        const gradOut = grads.get(record.output);
+        if (!gradOut) {
+          continue;
+        }
 
-      const gradsOut = await this.runBackward(entry.backward, record, gradOut);
-      for (const { input, grad } of gradsOut) {
-        if (input && grad) {
-          await this.accumulateGrad(grads, input, grad);
+        const gradsOut = await this.runBackward(entry.backward, record, gradOut);
+        for (const { input, grad } of gradsOut) {
+          if (input && grad) {
+            await this.accumulateGrad(grads, input, grad);
+          }
         }
       }
-    }
 
-    return grads;
+      return grads;
+    } finally {
+      for (const buffer of this.retainedBuffers) {
+        try {
+          releaseBuffer(buffer);
+        } catch {}
+      }
+      this.retainedBuffers.clear();
+    }
   }
 
   isTensorLike(value) {
@@ -245,9 +263,7 @@ export class AutogradTape {
     expanded.set(gradRow.subarray(0, copyCount), rowOffset);
     const dtype = gradOut.dtype === 'f16' ? 'f16' : 'f32';
     const payload = dtype === 'f16' ? f32ToF16Array(expanded) : expanded;
-    const outBuffer = acquireBuffer(payload.byteLength, undefined, 'row_slice_backward_output');
-    uploadData(outBuffer, payload);
-    return createTensor(outBuffer, dtype, [rows, cols], 'row_slice_backward_output');
+    return createUploadedTensor(payload, dtype, [rows, cols], 'row_slice_backward_output');
   }
 
   resolveSiluRowsplitGate(gateValue, activation) {
@@ -305,9 +321,7 @@ export class AutogradTape {
 
     const dtype = gradOut.dtype === 'f16' ? 'f16' : 'f32';
     const payload = dtype === 'f16' ? f32ToF16Array(output) : output;
-    const outBuffer = acquireBuffer(payload.byteLength, undefined, 'silu_rowsplit_backward_output');
-    uploadData(outBuffer, payload);
-    return createTensor(outBuffer, dtype, [numTokens, dim * 2], 'silu_rowsplit_backward_output');
+    return createUploadedTensor(payload, dtype, [numTokens, dim * 2], 'silu_rowsplit_backward_output');
   }
 
   async accumulateLargeGradF32(existing, grad, size, shape) {
@@ -317,35 +331,49 @@ export class AutogradTape {
     }
     const bytesPerElement = 4;
     const outputBuffer = acquireBuffer(size * bytesPerElement, undefined, 'grad_accum_large_output');
-
-    for (let offset = 0; offset < size; offset += MAX_RESIDUAL_ELEMENTS_PER_DISPATCH) {
-      const chunkElements = Math.min(MAX_RESIDUAL_ELEMENTS_PER_DISPATCH, size - offset);
-      const chunkBytes = chunkElements * bytesPerElement;
-      const chunkOffsetBytes = offset * bytesPerElement;
-
-      const aChunkBuffer = acquireBuffer(chunkBytes, undefined, 'grad_accum_large_a_chunk');
-      const bChunkBuffer = acquireBuffer(chunkBytes, undefined, 'grad_accum_large_b_chunk');
-      const copyIn = device.createCommandEncoder();
-      copyIn.copyBufferToBuffer(existing.buffer, chunkOffsetBytes, aChunkBuffer, 0, chunkBytes);
-      copyIn.copyBufferToBuffer(grad.buffer, chunkOffsetBytes, bChunkBuffer, 0, chunkBytes);
-      device.queue.submit([copyIn.finish()]);
-
-      const aChunk = createTensor(aChunkBuffer, 'f32', [chunkElements], 'grad_accum_large_a_tensor');
-      const bChunk = createTensor(bChunkBuffer, 'f32', [chunkElements], 'grad_accum_large_b_tensor');
-      const summedChunk = await runResidualAdd(aChunk, bChunk, chunkElements);
-
-      const copyOut = device.createCommandEncoder();
-      copyOut.copyBufferToBuffer(summedChunk.buffer, 0, outputBuffer, chunkOffsetBytes, chunkBytes);
-      device.queue.submit([copyOut.finish()]);
-
-      releaseBuffer(aChunkBuffer);
-      releaseBuffer(bChunkBuffer);
-      if (summedChunk?.buffer && summedChunk.buffer !== outputBuffer) {
-        releaseBuffer(summedChunk.buffer);
+    try {
+      for (let offset = 0; offset < size; offset += MAX_RESIDUAL_ELEMENTS_PER_DISPATCH) {
+        const chunkElements = Math.min(MAX_RESIDUAL_ELEMENTS_PER_DISPATCH, size - offset);
+        const chunkBytes = chunkElements * bytesPerElement;
+        const chunkOffsetBytes = offset * bytesPerElement;
+
+        let aChunkBuffer = null;
+        let bChunkBuffer = null;
+        let summedChunkBuffer = null;
+        try {
+          aChunkBuffer = acquireBuffer(chunkBytes, undefined, 'grad_accum_large_a_chunk');
+          bChunkBuffer = acquireBuffer(chunkBytes, undefined, 'grad_accum_large_b_chunk');
+          const copyIn = device.createCommandEncoder();
+          copyIn.copyBufferToBuffer(existing.buffer, chunkOffsetBytes, aChunkBuffer, 0, chunkBytes);
+          copyIn.copyBufferToBuffer(grad.buffer, chunkOffsetBytes, bChunkBuffer, 0, chunkBytes);
+          device.queue.submit([copyIn.finish()]);
+
+          const aChunk = createTensor(aChunkBuffer, 'f32', [chunkElements], 'grad_accum_large_a_tensor');
+          const bChunk = createTensor(bChunkBuffer, 'f32', [chunkElements], 'grad_accum_large_b_tensor');
+          const summedChunk = await runResidualAdd(aChunk, bChunk, chunkElements);
+          summedChunkBuffer = summedChunk?.buffer ?? null;
+
+          const copyOut = device.createCommandEncoder();
+          copyOut.copyBufferToBuffer(summedChunk.buffer, 0, outputBuffer, chunkOffsetBytes, chunkBytes);
+          device.queue.submit([copyOut.finish()]);
+        } finally {
+          if (aChunkBuffer) {
+            releaseBuffer(aChunkBuffer);
+          }
+          if (bChunkBuffer) {
+            releaseBuffer(bChunkBuffer);
+          }
+          if (summedChunkBuffer && summedChunkBuffer !== outputBuffer) {
+            releaseBuffer(summedChunkBuffer);
+          }
+        }
       }
-    }
 
-    return createTensor(outputBuffer, 'f32', [...shape], 'grad_accum_large_output');
+      return createTensor(outputBuffer, 'f32', [...shape], 'grad_accum_large_output');
+    } catch (error) {
+      releaseBuffer(outputBuffer);
+      throw error;
+    }
   }
 
 

package/src/training/checkpoint-watch.d.ts

@@ -0,0 +1,8 @@
+export declare function watchFinalizedCheckpoints(options: {
+  checkpointsDir: string;
+  manifestPath: string;
+  pollIntervalMs?: number | null;
+  stopWhenIdle?: boolean;
+  signal?: AbortSignal | null;
+  onCheckpoint: (markerPath: string) => Promise<void> | void;
+}): Promise<{ ok: true; processedCount: number; manifestPath: string; aborted?: boolean }>;

package/src/training/checkpoint-watch.js

@@ -0,0 +1,139 @@
+import { readdir, readFile } from 'node:fs/promises';
+import { join, resolve } from 'node:path';
+
+import { writeJsonArtifact } from './operator-artifacts.js';
+
+async function listCheckpointMarkers(checkpointsDir) {
+  const absoluteDir = resolve(String(checkpointsDir));
+  const entries = await readdir(absoluteDir, { withFileTypes: true });
+  const markers = [];
+  for (const entry of entries) {
+    if (!entry.isDirectory()) {
+      continue;
+    }
+    const entryPath = join(absoluteDir, entry.name);
+    const markerPath = join(entryPath, 'checkpoint.complete.json');
+    try {
+      await readFile(markerPath, 'utf8');
+      markers.push(markerPath);
+      continue;
+    } catch (error) {
+      if (error?.code !== 'ENOENT') {
+        throw error;
+      }
+    }
+    markers.push(...await listCheckpointMarkers(entryPath));
+  }
+  return markers.sort((left, right) => left.localeCompare(right));
+}
+
+async function ensureDirectoryExists(directoryPath) {
+  try {
+    const entries = await readdir(directoryPath, { withFileTypes: true });
+    return Array.isArray(entries);
+  } catch (error) {
+    if (error?.code === 'ENOENT') {
+      return false;
+    }
+    throw error;
+  }
+}
+
+async function readProcessedManifest(manifestPath) {
+  try {
+    const raw = await readFile(manifestPath, 'utf8');
+    const parsed = JSON.parse(raw);
+    const processed = Array.isArray(parsed?.processedCheckpointMarkers)
+      ? parsed.processedCheckpointMarkers.filter((entry) => typeof entry === 'string')
+      : [];
+    return new Set(processed);
+  } catch (error) {
+    if (error?.code === 'ENOENT') {
+      return new Set();
+    }
+    throw error;
+  }
+}
+
+function createWatchResult(processed, manifestPath, aborted = false) {
+  return {
+    ok: true,
+    processedCount: processed.size,
+    manifestPath,
+    aborted,
+  };
+}
+
+async function waitForPollInterval(pollIntervalMs, signal) {
+  if (!signal) {
+    await new Promise((resolvePromise) => setTimeout(resolvePromise, pollIntervalMs));
+    return true;
+  }
+  if (signal.aborted) {
+    return false;
+  }
+  return new Promise((resolvePromise) => {
+    const onAbort = () => {
+      clearTimeout(timer);
+      resolvePromise(false);
+    };
+    const timer = setTimeout(() => {
+      signal.removeEventListener('abort', onAbort);
+      resolvePromise(true);
+    }, pollIntervalMs);
+    signal.addEventListener('abort', onAbort, { once: true });
+  });
+}
+
+export async function watchFinalizedCheckpoints(options) {
+  const checkpointsDir = resolve(String(options.checkpointsDir));
+  const manifestPath = resolve(String(options.manifestPath));
+  const pollIntervalMs = Number.isFinite(options.pollIntervalMs)
+    ? Math.max(100, Math.floor(options.pollIntervalMs))
+    : 2000;
+  const stopWhenIdle = options.stopWhenIdle === true;
+  const onCheckpoint = typeof options.onCheckpoint === 'function'
+    ? options.onCheckpoint
+    : null;
+  const signal = options.signal ?? null;
+  if (!onCheckpoint) {
+    throw new Error('watchFinalizedCheckpoints requires onCheckpoint(markerPath).');
+  }
+
+  const processed = await readProcessedManifest(manifestPath);
+  let idlePolls = 0;
+  for (;;) {
+    if (signal?.aborted) {
+      return createWatchResult(processed, manifestPath, true);
+    }
+    const checkpointsExist = await ensureDirectoryExists(checkpointsDir);
+    const markers = checkpointsExist
+      ? await listCheckpointMarkers(checkpointsDir)
+      : [];
+    let sawNewMarker = false;
+    for (const markerPath of markers) {
+      if (processed.has(markerPath)) continue;
+      sawNewMarker = true;
+      await onCheckpoint(markerPath);
+      processed.add(markerPath);
+      await writeJsonArtifact(manifestPath, {
+        artifactType: 'training_checkpoint_watch_manifest',
+        schemaVersion: 1,
+        generatedAt: new Date().toISOString(),
+        processedCheckpointMarkers: [...processed].sort((left, right) => left.localeCompare(right)),
+      });
+    }
+    if (!sawNewMarker) {
+      idlePolls += 1;
+      if (stopWhenIdle && idlePolls > 0) {
+        return createWatchResult(processed, manifestPath);
+      }
+    } else {
+      idlePolls = 0;
+    }
+    const shouldContinue = await waitForPollInterval(pollIntervalMs, signal);
+    if (!shouldContinue) {
+      return createWatchResult(processed, manifestPath, true);
+    }
+  }
+}

package/src/training/checkpoint.d.ts

@@ -23,7 +23,12 @@ export declare function saveCheckpoint(
   key: string,
   data: unknown,
   options?: CheckpointStoreOptions
-): Promise<void>;
+): Promise<{
+  key: string;
+  path: string | null;
+  metadata: Record<string, unknown>;
+  data: unknown;
+}>;
 
 export declare function loadCheckpoint(
   key: string,

package/src/training/checkpoint.js

@@ -31,6 +31,13 @@ function openCheckpointDB(options = {}) {
   });
 }
 
+function closeCheckpointDB(db) {
+  if (!db || typeof db.close !== 'function') {
+    return;
+  }
+  db.close();
+}
+
 async function resolveNodeCheckpointPath(key, options = {}) {
   const [{ resolve, join, dirname }, { mkdir }] = await Promise.all([
     import('node:path'),
@@ -140,9 +147,15 @@ export async function saveCheckpoint(key, payload, options = {}) {
   const useNodeStore = isNodeRuntime() && typeof indexedDB === 'undefined';
   const nodePath = useNodeStore ? await resolveNodeCheckpointPath(key, options) : null;
   const browserStore = useNodeStore ? null : await openCheckpointDB(options);
-  const previousData = useNodeStore
-    ? await readNodeCheckpointRecord(nodePath)
-    : await readCheckpointRecord(browserStore.db, browserStore.storeName, key);
+  let previousData;
+  try {
+    previousData = useNodeStore
+      ? await readNodeCheckpointRecord(nodePath)
+      : await readCheckpointRecord(browserStore.db, browserStore.storeName, key);
+  } catch (error) {
+    closeCheckpointDB(browserStore?.db);
+    throw error;
+  }
   const previousMetadata = previousData?.metadata || {};
   const previousLineage = previousMetadata.lineage || {};
   const previousCheckpointHash = options.priorCheckpointHash
@@ -184,13 +197,35 @@ export async function saveCheckpoint(key, payload, options = {}) {
 
   if (useNodeStore) {
     await writeNodeCheckpointRecord(nodePath, data);
-    return;
+    return {
+      key,
+      path: nodePath,
+      metadata: data.metadata,
+      data,
+    };
   }
 
   return new Promise((resolve, reject) => {
     const tx = browserStore.db.transaction(browserStore.storeName, 'readwrite');
-    tx.oncomplete = () => resolve();
-    tx.onerror = () => reject(tx.error);
+    tx.oncomplete = () => {
+      closeCheckpointDB(browserStore.db);
+      resolve({
+        key,
+        path: null,
+        metadata: data.metadata,
+        data,
+      });
+    };
+    tx.onerror = () => {
+      const error = tx.error;
+      closeCheckpointDB(browserStore.db);
+      reject(error);
+    };
+    tx.onabort = () => {
+      const error = tx.error ?? new Error('Checkpoint transaction aborted');
+      closeCheckpointDB(browserStore.db);
+      reject(error);
+    };
     const store = tx.objectStore(browserStore.storeName);
     store.put(data, key);
   });
@@ -203,7 +238,11 @@ export async function loadCheckpoint(key, options = {}) {
     ? await readNodeCheckpointRecord(nodePath)
     : await (async () => {
       const { db, storeName } = await openCheckpointDB(options);
-      return readCheckpointRecord(db, storeName, key);
+      try {
+        return await readCheckpointRecord(db, storeName, key);
+      } finally {
+        closeCheckpointDB(db);
+      }
     })();
 
   if (!data || !data.metadata || !options.expectedMetadata) {

package/src/training/clip.js

@@ -12,7 +12,8 @@ async function readGradData(grad) {
 }
 
 export async function clipGradients(grads, config) {
-  const maxNorm = config?.training?.gradient?.maxNorm;
+  const maxNorm = config?.training?.gradientClipping?.maxNorm
+    ?? config?.training?.gradient?.maxNorm;
   let sumSq = 0;
   let totalParamCount = 0;
 

package/src/training/datasets/token-batch.js

@@ -1,5 +1,5 @@
 
-import { acquireBuffer, uploadData } from '../../memory/buffer-pool.js';
+import { acquireBuffer, uploadData, releaseBuffer } from '../../memory/buffer-pool.js';
 import { createTensor } from '../../gpu/tensor.js';
 
 function flattenTokenBatch(samples, key) {
@@ -27,14 +27,26 @@ export function buildTokenBatch(samples) {
 }
 
 export function createTokenBatchTensors(batch) {
-  const inputBuf = acquireBuffer(batch.inputFlat.byteLength, undefined, 'train_input_tokens');
-  uploadData(inputBuf, batch.inputFlat);
+  let inputBuf = null;
+  let targetBuf = null;
+  try {
+    inputBuf = acquireBuffer(batch.inputFlat.byteLength, undefined, 'train_input_tokens');
+    uploadData(inputBuf, batch.inputFlat);
 
-  const targetBuf = acquireBuffer(batch.targetFlat.byteLength, undefined, 'train_target_tokens');
-  uploadData(targetBuf, batch.targetFlat);
+    targetBuf = acquireBuffer(batch.targetFlat.byteLength, undefined, 'train_target_tokens');
+    uploadData(targetBuf, batch.targetFlat);
 
-  const input = createTensor(inputBuf, 'f32', [batch.inputFlat.length], 'train_input_tokens');
-  const targets = createTensor(targetBuf, 'f32', [batch.targetFlat.length], 'train_target_tokens');
+    const input = createTensor(inputBuf, 'f32', [batch.inputFlat.length], 'train_input_tokens');
+    const targets = createTensor(targetBuf, 'f32', [batch.targetFlat.length], 'train_target_tokens');
 
-  return { input, targets, offsets: batch.offsets };
+    return { input, targets, offsets: batch.offsets };
+  } catch (error) {
+    if (inputBuf) {
+      releaseBuffer(inputBuf);
+    }
+    if (targetBuf) {
+      releaseBuffer(targetBuf);
+    }
+    throw error;
+  }
 }