@simulatte/doppler 0.1.6 → 0.1.7

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

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

package/src/training/checkpoint-watch.js

@@ -55,6 +55,36 @@ async function readProcessedManifest(manifestPath) {
   }
 }
 
+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));
@@ -65,6 +95,7 @@ export async function watchFinalizedCheckpoints(options) {
   const onCheckpoint = typeof options.onCheckpoint === 'function'
     ? options.onCheckpoint
     : null;
+  const signal = options.signal ?? null;
   if (!onCheckpoint) {
     throw new Error('watchFinalizedCheckpoints requires onCheckpoint(markerPath).');
   }
@@ -72,6 +103,9 @@ export async function watchFinalizedCheckpoints(options) {
   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)
@@ -92,15 +126,14 @@ export async function watchFinalizedCheckpoints(options) {
     if (!sawNewMarker) {
       idlePolls += 1;
       if (stopWhenIdle && idlePolls > 0) {
-        return {
-          ok: true,
-          processedCount: processed.size,
-          manifestPath,
-        };
+        return createWatchResult(processed, manifestPath);
       }
     } else {
       idlePolls = 0;
     }
-    await new Promise((resolvePromise) => setTimeout(resolvePromise, pollIntervalMs));
+    const shouldContinue = await waitForPollInterval(pollIntervalMs, signal);
+    if (!shouldContinue) {
+      return createWatchResult(processed, manifestPath, true);
+    }
   }
 }

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
@@ -194,13 +207,25 @@ export async function saveCheckpoint(key, payload, options = {}) {
 
   return new Promise((resolve, reject) => {
     const tx = browserStore.db.transaction(browserStore.storeName, 'readwrite');
-    tx.oncomplete = () => resolve({
-      key,
-      path: null,
-      metadata: data.metadata,
-      data,
-    });
-    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);
   });
@@ -213,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;
+  }
 }

package/src/training/distillation/checkpoint-watch.js

@@ -14,6 +14,7 @@ export async function watchDistillationCheckpoints(options) {
     manifestPath,
     pollIntervalMs: options.pollIntervalMs || 2000,
     stopWhenIdle: options.stopWhenIdle === true,
+    signal: options.signal ?? null,
     onCheckpoint: async (markerPath) => {
       const { marker } = await readDistillCheckpointMarker(markerPath);
       const reports = await evaluateDistillationCheckpoint({

package/src/training/distillation/student-fixture.d.ts

@@ -0,0 +1,22 @@
+export interface DistillStudentFixture {
+  config: Record<string, unknown>;
+  model: {
+    forward: (input: unknown, tape: unknown) => Promise<unknown>;
+    forwardDistill?: (
+      batch: unknown,
+      tape: unknown,
+      options?: Record<string, unknown>
+    ) => Promise<{ logits: unknown }>;
+    cleanupDistillStep?: () => void;
+    loraParams?: () => unknown[];
+    paramGroups?: () => Record<string, unknown[]>;
+  };
+  outputDim?: number;
+  embeddingDim?: number;
+  cleanup(): void;
+}
+
+export declare function createDistillStudentRuntimeModelFixture(
+  overrides?: Record<string, unknown>,
+  options?: Record<string, unknown>
+): Promise<DistillStudentFixture>;