@simulatte/doppler 0.1.6 → 0.1.7

package/src/training/distillation/suite-data.d.ts

@@ -0,0 +1,45 @@
+export interface DistillDataScope {
+  sourceLangs: string[] | null;
+  targetLangs: string[] | null;
+  pairAllowlist: string[] | null;
+  sourceLangSet: Set<string> | null;
+  targetLangSet: Set<string> | null;
+  pairAllowlistSet: Set<string> | null;
+  strictPairContract: boolean;
+}
+
+export interface DistillSample {
+  index?: number;
+  direction?: string | null;
+  sourceLang?: string | null;
+  targetLang?: string | null;
+  source?: string | null;
+  targetPos?: string | null;
+  targetNeg?: string | null;
+}
+
+export declare function normalizeOptionalString(value: unknown): string | null;
+
+export declare function normalizeDistillDatasetPath(value: unknown): string | null;
+
+export declare function resolveDistillDataScope(
+  options?: Record<string, unknown>,
+  trainingConfig?: Record<string, unknown> | null
+): DistillDataScope;
+
+export declare function encodeDistillRow(
+  record: Record<string, unknown> | null | undefined,
+  index: number,
+  scope?: DistillDataScope | null
+): DistillSample | null;
+
+export declare function summarizeDirectionCounts(
+  samples: Array<Record<string, unknown> | null | undefined>
+): Record<string, number>;
+
+export declare function buildDistillPrompt(sample: Record<string, unknown> | null | undefined): string;
+
+export declare function buildDistillCandidatePrompt(
+  sample: Record<string, unknown> | null | undefined,
+  candidate: unknown
+): string;

package/src/training/distillation/suite-data.js

@@ -0,0 +1,189 @@
+export function normalizeOptionalString(value) {
+  if (value === undefined || value === null) return null;
+  const trimmed = String(value).trim();
+  return trimmed || null;
+}
+
+export function normalizeDistillDatasetPath(value) {
+  return normalizeOptionalString(value);
+}
+
+function normalizeLangCode(value) {
+  const normalized = normalizeOptionalString(value);
+  if (!normalized) return null;
+  const compact = normalized.toLowerCase().replace(/_/g, '-');
+  if (compact.startsWith('en')) return 'en';
+  if (compact.startsWith('es')) return 'es';
+  return compact;
+}
+
+function normalizePairDirection(value) {
+  const pair = normalizeOptionalString(value);
+  if (!pair) return null;
+  const normalized = pair.toLowerCase().replace(/_/g, '-').replace(/\s+/g, '');
+  const parts = normalized.includes('->')
+    ? normalized.split('->').filter(Boolean)
+    : normalized.split('-').filter(Boolean);
+  if (parts.length !== 2) return null;
+  return `${normalizeLangCode(parts[0]) || parts[0]}->${normalizeLangCode(parts[1]) || parts[1]}`;
+}
+
+function normalizeOptionalStringArray(value) {
+  if (value === undefined || value === null) return null;
+  const list = Array.isArray(value)
+    ? value
+    : (typeof value === 'string' ? value.split(',') : null);
+  if (!Array.isArray(list)) return null;
+  const normalized = list
+    .map((entry) => normalizeOptionalString(entry))
+    .filter(Boolean);
+  return normalized.length > 0 ? normalized : null;
+}
+
+function normalizeDistillLanguageAllowlist(value) {
+  const list = normalizeOptionalStringArray(value);
+  if (!list) return null;
+  const normalized = list
+    .map((entry) => normalizeLangCode(entry))
+    .filter(Boolean);
+  if (normalized.length === 0) return null;
+  return [...new Set(normalized)];
+}
+
+function normalizeDistillPairAllowlist(value) {
+  const list = normalizeOptionalStringArray(value);
+  if (!list) return null;
+  const normalized = list
+    .map((entry) => normalizePairDirection(entry))
+    .filter(Boolean);
+  if (normalized.length === 0) return null;
+  return [...new Set(normalized)];
+}
+
+export function resolveDistillDataScope(options = {}, trainingConfig = null) {
+  const distillConfig = trainingConfig?.distill || {};
+  const sourceLangs = normalizeDistillLanguageAllowlist(
+    options.distillSourceLangs ?? distillConfig.sourceLangs ?? null
+  );
+  const targetLangs = normalizeDistillLanguageAllowlist(
+    options.distillTargetLangs ?? distillConfig.targetLangs ?? null
+  );
+  const pairAllowlist = normalizeDistillPairAllowlist(
+    options.distillPairAllowlist ?? distillConfig.pairAllowlist ?? null
+  );
+  const strictPairContract = (
+    options.strictPairContract === true
+    || distillConfig.strictPairContract === true
+  );
+  return {
+    sourceLangs,
+    targetLangs,
+    pairAllowlist,
+    sourceLangSet: sourceLangs ? new Set(sourceLangs) : null,
+    targetLangSet: targetLangs ? new Set(targetLangs) : null,
+    pairAllowlistSet: pairAllowlist ? new Set(pairAllowlist) : null,
+    strictPairContract,
+  };
+}
+
+function resolveDistillDirection(record) {
+  const pairDirection = normalizePairDirection(record?.pair);
+  if (pairDirection) return pairDirection;
+  const srcLang = normalizeLangCode(record?.src_lang);
+  const tgtLang = normalizeLangCode(record?.tgt_lang || record?.lang);
+  if (srcLang && tgtLang) {
+    return `${srcLang}->${tgtLang}`;
+  }
+  return null;
+}
+
+function resolveStringCandidate(record, keys) {
+  for (const key of keys) {
+    const value = normalizeOptionalString(record?.[key]);
+    if (value) return value;
+  }
+  return null;
+}
+
+export function encodeDistillRow(record, index, scope = null) {
+  if (!record || typeof record !== 'object') return null;
+  const source = resolveStringCandidate(record, ['source', 'query']);
+  const targetPos = resolveStringCandidate(record, ['target_pos', 'target', 'pos']);
+  const targetNeg = resolveStringCandidate(record, ['target_neg', 'neg']);
+  if (!source || !targetPos) return null;
+  const sourceLangRaw = normalizeLangCode(record?.src_lang);
+  const targetLangRaw = normalizeLangCode(record?.tgt_lang || record?.lang);
+  const pairDirection = normalizePairDirection(record?.pair);
+  const sourceTargetDirection = (
+    sourceLangRaw && targetLangRaw
+      ? `${sourceLangRaw}->${targetLangRaw}`
+      : null
+  );
+  if (scope?.strictPairContract === true) {
+    if (!sourceLangRaw || !targetLangRaw) {
+      throw new Error('strictPairContract requires src_lang and tgt_lang/lang on each row.');
+    }
+    if (!pairDirection) {
+      throw new Error('strictPairContract requires pair on each row.');
+    }
+    if (pairDirection !== sourceTargetDirection) {
+      throw new Error(`pair "${record?.pair}" does not match src/tgt "${sourceLangRaw}-${targetLangRaw}".`);
+    }
+  }
+  const direction = pairDirection || sourceTargetDirection || resolveDistillDirection(record) || 'unknown';
+  const [directionSourceLang, directionTargetLang] = String(direction).split('->');
+  const sourceLang = sourceLangRaw || normalizeLangCode(directionSourceLang);
+  const targetLang = targetLangRaw || normalizeLangCode(directionTargetLang);
+  if (scope?.sourceLangSet && (!sourceLang || !scope.sourceLangSet.has(sourceLang))) {
+    return null;
+  }
+  if (scope?.targetLangSet && (!targetLang || !scope.targetLangSet.has(targetLang))) {
+    return null;
+  }
+  if (scope?.pairAllowlistSet && !scope.pairAllowlistSet.has(direction)) {
+    return null;
+  }
+
+  return {
+    index,
+    direction,
+    sourceLang: sourceLang || null,
+    targetLang: targetLang || null,
+    source,
+    targetPos,
+    targetNeg: targetNeg || null,
+  };
+}
+
+export function summarizeDirectionCounts(samples) {
+  const counts = {};
+  for (const sample of samples) {
+    const key = sample?.direction || 'unknown';
+    counts[key] = (counts[key] || 0) + 1;
+  }
+  return counts;
+}
+
+function resolveLanguageName(langCode) {
+  const normalized = normalizeLangCode(langCode);
+  if (normalized === 'en') return 'English';
+  if (normalized === 'es') return 'Spanish';
+  return normalized || 'target';
+}
+
+export function buildDistillPrompt(sample) {
+  const direction = String(sample?.direction || '').trim();
+  const [srcCodeRaw, tgtCodeRaw] = direction.split('->');
+  const srcCode = normalizeLangCode(srcCodeRaw) || srcCodeRaw || 'source';
+  const tgtCode = normalizeLangCode(tgtCodeRaw) || tgtCodeRaw || 'target';
+  const srcName = resolveLanguageName(srcCode);
+  const tgtName = resolveLanguageName(tgtCode);
+  const source = String(sample?.source || '').trim();
+  return `Translate from ${srcName} to ${tgtName}:\n${source}\nTranslation:`;
+}
+
+export function buildDistillCandidatePrompt(sample, candidate) {
+  const base = buildDistillPrompt(sample);
+  const text = String(candidate || '').trim();
+  return text ? `${base} ${text}` : base;
+}

package/src/training/lora-pipeline.js

@@ -3,7 +3,6 @@ import { join, resolve } from 'node:path';
 
 import { loadBackwardRegistry } from '../config/backward-registry-loader.js';
 import { acquireBuffer, readBuffer, releaseBuffer, uploadData } from '../memory/buffer-pool.js';
-import { createTensor } from '../gpu/tensor.js';
 import { runMatmul } from '../gpu/kernels/index.js';
 import { runResidualAdd } from '../gpu/kernels/residual.js';
 import { parseJsonl } from './datasets/jsonl.js';
@@ -27,6 +26,7 @@ import {
 } from './operator-artifacts.js';
 import { watchFinalizedCheckpoints } from './checkpoint-watch.js';
 import { loadLoRAFromManifest } from '../adapters/lora-loader.js';
+import { createUploadedTensor } from './tensor-factory.js';
 
 function stableSortObject(value) {
   if (Array.isArray(value)) {
@@ -48,16 +48,12 @@ function stableJson(value) {
 
 function makeTensorFromFloat32(values, shape, label) {
   const data = values instanceof Float32Array ? values : new Float32Array(values);
-  const buffer = acquireBuffer(data.byteLength, undefined, label);
-  uploadData(buffer, data);
-  return createTensor(buffer, 'f32', [...shape], label);
+  return createUploadedTensor(data, 'f32', shape, label);
 }
 
 function makeTensorFromUint32(values, shape, label) {
   const data = values instanceof Uint32Array ? values : new Uint32Array(values);
-  const buffer = acquireBuffer(data.byteLength, undefined, label);
-  uploadData(buffer, data);
-  return createTensor(buffer, 'u32', [...shape], label);
+  return createUploadedTensor(data, 'u32', shape, label);
 }
 
 function releaseTensor(tensor) {
@@ -709,6 +705,7 @@ export async function watchLoraCheckpoints(options) {
     manifestPath: join(options.runRoot, 'scoreboard', 'watch-manifest.json'),
     pollIntervalMs: options.pollIntervalMs || 2000,
     stopWhenIdle: options.stopWhenIdle === true,
+    signal: options.signal ?? null,
     onCheckpoint: async (markerPath) => {
       const raw = await readFile(markerPath, 'utf8');
       const marker = JSON.parse(raw);

package/src/training/lora.js

@@ -12,18 +12,32 @@ export class LoraAdapter {
     const aBytes = tensorBytes([inDim, rank], dtype);
     const bBytes = tensorBytes([rank, outDim], dtype);
 
-    this.A = createTensor(
-      acquireBuffer(aBytes, BufferUsage.STORAGE, 'lora_A'),
-      dtype,
-      [inDim, rank],
-      'lora_A'
-    );
-    this.B = createTensor(
-      acquireBuffer(bBytes, BufferUsage.STORAGE, 'lora_B'),
-      dtype,
-      [rank, outDim],
-      'lora_B'
-    );
+    let aBuffer = null;
+    let bBuffer = null;
+    try {
+      aBuffer = acquireBuffer(aBytes, BufferUsage.STORAGE, 'lora_A');
+      bBuffer = acquireBuffer(bBytes, BufferUsage.STORAGE, 'lora_B');
+      this.A = createTensor(
+        aBuffer,
+        dtype,
+        [inDim, rank],
+        'lora_A'
+      );
+      this.B = createTensor(
+        bBuffer,
+        dtype,
+        [rank, outDim],
+        'lora_B'
+      );
+    } catch (error) {
+      if (aBuffer) {
+        releaseBuffer(aBuffer);
+      }
+      if (bBuffer) {
+        releaseBuffer(bBuffer);
+      }
+      throw error;
+    }
     this.alpha = alpha;
     this.rank = rank;
   }

package/src/training/loss.js

@@ -1,6 +1,5 @@
 
 import { runSoftmax, runCrossEntropyLoss, castF16ToF32 } from '../gpu/kernels/index.js';
-import { releaseBuffer } from '../memory/buffer-pool.js';
 import { OpType } from './autograd.js';
 
 export async function crossEntropyLoss(logits, targets, config, tape) {
@@ -25,13 +24,13 @@ export async function crossEntropyLoss(logits, targets, config, tape) {
     OpType.SOFTMAX,
     (input) => runSoftmax(input, -1, { batchSize: numTokens, size: vocabSize }),
     [logitsF32],
-    { rows: numTokens, cols: vocabSize }
+    {
+      rows: numTokens,
+      cols: vocabSize,
+      retainBuffers: logitsF32 !== logits ? [logitsF32.buffer] : [],
+    }
   );
 
-  if (logitsF32 !== logits) {
-    releaseBuffer(logitsF32.buffer);
-  }
-
   return tape.record(
     OpType.CROSS_ENTROPY,
     (input, target) => runCrossEntropyLoss(input, target, { numTokens, vocabSize }),

package/src/training/objectives/cross_entropy.js

@@ -1,15 +1,12 @@
 import { crossEntropyLoss as defaultCrossEntropyLoss } from '../loss.js';
-import { acquireBuffer, uploadData } from '../../memory/buffer-pool.js';
-import { createTensor } from '../../gpu/tensor.js';
 import { createTrainingObjective } from './base.js';
+import { createUploadedTensor } from '../tensor-factory.js';
 
 function createLossGradient(loss, lossScale) {
   const lossElements = loss.shape.reduce((acc, value) => acc * value, 1);
   const gradData = new Float32Array(lossElements);
   gradData.fill(lossScale);
-  const gradBuf = acquireBuffer(gradData.byteLength, undefined, 'loss_grad_output');
-  uploadData(gradBuf, gradData);
-  return createTensor(gradBuf, 'f32', [...loss.shape], 'loss_grad_output');
+  return createUploadedTensor(gradData, 'f32', loss.shape, 'loss_grad_output');
 }
 
 export function createCrossEntropyObjective(options = {}) {

package/src/training/objectives/distill_kd.js

@@ -1,8 +1,8 @@
 import { crossEntropyLoss as defaultCrossEntropyLoss } from '../loss.js';
 import { createTrainingObjective } from './base.js';
-import { acquireBuffer, readBuffer, uploadData } from '../../memory/buffer-pool.js';
-import { createTensor } from '../../gpu/tensor.js';
+import { readBuffer } from '../../memory/buffer-pool.js';
 import { f16ToF32Array, f32ToF16Array } from '../../inference/kv-cache/types.js';
+import { createUploadedTensor } from '../tensor-factory.js';
 
 const EPS = 1e-8;
 
@@ -31,9 +31,7 @@ function createLossGradient(loss, lossScale) {
   const lossElements = loss.shape.reduce((acc, value) => acc * value, 1);
   const gradData = new Float32Array(lossElements);
   gradData.fill(lossScale);
-  const gradBuf = acquireBuffer(gradData.byteLength, undefined, 'distill_kd_loss_grad_output');
-  uploadData(gradBuf, gradData);
-  return createTensor(gradBuf, 'f32', [...loss.shape], 'distill_kd_loss_grad_output');
+  return createUploadedTensor(gradData, 'f32', loss.shape, 'distill_kd_loss_grad_output');
 }
 
 function createGradientTensor(values, shape, dtype, label) {
@@ -42,9 +40,7 @@ function createGradientTensor(values, shape, dtype, label) {
   const payload = tensorDtype === 'f16'
     ? f32ToF16Array(floatValues)
     : floatValues;
-  const gradBuf = acquireBuffer(payload.byteLength, undefined, label);
-  uploadData(gradBuf, payload);
-  return createTensor(gradBuf, tensorDtype, [...shape], label);
+  return createUploadedTensor(payload, tensorDtype, shape, label);
 }
 
 async function readLogitsRows(logitsTensor) {

package/src/training/objectives/distill_triplet.js

@@ -1,8 +1,8 @@
 import { crossEntropyLoss as defaultCrossEntropyLoss } from '../loss.js';
 import { createTrainingObjective } from './base.js';
-import { acquireBuffer, readBuffer, uploadData } from '../../memory/buffer-pool.js';
-import { createTensor } from '../../gpu/tensor.js';
+import { readBuffer } from '../../memory/buffer-pool.js';
 import { f16ToF32Array, f32ToF16Array } from '../../inference/kv-cache/types.js';
+import { createUploadedTensor } from '../tensor-factory.js';
 
 function toFinite(value, fallback) {
   const parsed = Number(value);
@@ -29,9 +29,7 @@ function createLossGradient(loss, lossScale) {
   const lossElements = loss.shape.reduce((acc, value) => acc * value, 1);
   const gradData = new Float32Array(lossElements);
   gradData.fill(lossScale);
-  const gradBuf = acquireBuffer(gradData.byteLength, undefined, 'distill_triplet_loss_grad_output');
-  uploadData(gradBuf, gradData);
-  return createTensor(gradBuf, 'f32', [...loss.shape], 'distill_triplet_loss_grad_output');
+  return createUploadedTensor(gradData, 'f32', loss.shape, 'distill_triplet_loss_grad_output');
 }
 
 function createGradientTensor(values, shape, dtype, label) {
@@ -40,9 +38,7 @@ function createGradientTensor(values, shape, dtype, label) {
   const payload = tensorDtype === 'f16'
     ? f32ToF16Array(floatValues)
     : floatValues;
-  const gradBuf = acquireBuffer(payload.byteLength, undefined, label);
-  uploadData(gradBuf, payload);
-  return createTensor(gradBuf, tensorDtype, [...shape], label);
+  return createUploadedTensor(payload, tensorDtype, shape, label);
 }
 
 async function readLogitsRows(logitsTensor) {

package/src/training/objectives/ul_stage2_base.js

@@ -1,7 +1,7 @@
 import { crossEntropyLoss as defaultCrossEntropyLoss } from '../loss.js';
 import { createTrainingObjective } from './base.js';
-import { acquireBuffer, uploadData, releaseBuffer } from '../../memory/buffer-pool.js';
-import { createTensor } from '../../gpu/tensor.js';
+import { releaseBuffer } from '../../memory/buffer-pool.js';
+import { createUploadedTensor } from '../tensor-factory.js';
 
 function sigmoid(value) {
   return 1 / (1 + Math.exp(-value));
@@ -9,17 +9,13 @@ function sigmoid(value) {
 
 function createF32Tensor(values, shape, label) {
   const data = values instanceof Float32Array ? values : new Float32Array(values);
-  const buffer = acquireBuffer(data.byteLength, undefined, label);
-  uploadData(buffer, data);
-  return createTensor(buffer, 'f32', [...shape], label);
+  return createUploadedTensor(data, 'f32', shape, label);
 }
 
 function createU32TokenTensor(values, shape, label) {
   const data = values instanceof Uint32Array ? values : new Uint32Array(values);
-  const buffer = acquireBuffer(data.byteLength, undefined, label);
-  uploadData(buffer, data);
   // Token targets are consumed as raw u32 bytes by loss kernels.
-  return createTensor(buffer, 'f32', [...shape], label);
+  return createUploadedTensor(data, 'f32', shape, label);
 }
 
 function releaseTensor(tensor) {

package/src/training/operator-command.js

@@ -316,6 +316,7 @@ async function runDistillCommand(request) {
         layout: runArtifacts.layout,
         pollIntervalMs: request.pollIntervalMs || null,
         stopWhenIdle: request.stopWhenIdle === true,
+        signal: request.signal ?? null,
       })),
     };
   }
@@ -378,6 +379,7 @@ async function runLoraCommand(request) {
         runRoot: resolve(String(request.runRoot)),
         pollIntervalMs: request.pollIntervalMs || null,
         stopWhenIdle: request.stopWhenIdle === true,
+        signal: request.signal ?? null,
       })),
     };
   }

package/src/training/optimizer.js

@@ -1,4 +1,4 @@
-import { acquireBuffer, BufferUsage } from '../memory/buffer-pool.js';
+import { acquireBuffer, releaseBuffer, BufferUsage } from '../memory/buffer-pool.js';
 import { createTensor, tensorBytes } from '../gpu/tensor.js';
 import { runAdam } from '../gpu/kernels/backward/adam.js';
 
@@ -72,12 +72,24 @@ export class AdamOptimizer {
     let entry = this.state.get(param);
     if (!entry) {
       const bytes = tensorBytes(param.shape, param.dtype);
-      const mBuf = acquireBuffer(bytes, BufferUsage.STORAGE, 'adam_m');
-      const vBuf = acquireBuffer(bytes, BufferUsage.STORAGE, 'adam_v');
-      entry = {
-        m: createTensor(mBuf, param.dtype, [...param.shape], 'adam_m'),
-        v: createTensor(vBuf, param.dtype, [...param.shape], 'adam_v'),
-      };
+      let mBuf = null;
+      let vBuf = null;
+      try {
+        mBuf = acquireBuffer(bytes, BufferUsage.STORAGE, 'adam_m');
+        vBuf = acquireBuffer(bytes, BufferUsage.STORAGE, 'adam_v');
+        entry = {
+          m: createTensor(mBuf, param.dtype, [...param.shape], 'adam_m'),
+          v: createTensor(vBuf, param.dtype, [...param.shape], 'adam_v'),
+        };
+      } catch (error) {
+        if (mBuf) {
+          releaseBuffer(mBuf);
+        }
+        if (vBuf) {
+          releaseBuffer(vBuf);
+        }
+        throw error;
+      }
       this.state.set(param, entry);
     }
     return entry;

package/src/training/runner.js

@@ -617,7 +617,6 @@ function buildExpectedCheckpointMetadata(metadata) {
     'configHash',
     'datasetHash',
     'tokenizerHash',
-    'optimizerHash',
     'runtimePresetId',
     'kernelPathId',
   ]) {
@@ -845,6 +844,8 @@ export class TrainingRunner {
   }
 
   async run(model, dataset, options = {}) {
+    this.lastCheckpoint = null;
+    this.lastArtifact = null;
     const {
       epochs = 1,
       batchSize = 1,