@simulatte/doppler 0.1.5 → 0.1.7

package/src/training/distillation/eval.d.ts

@@ -0,0 +1,34 @@
+import type { LoadedTrainingWorkload } from '../workloads.js';
+
+export declare function evaluateDistillationModel(options: {
+  loadedWorkload: LoadedTrainingWorkload;
+  layout?: Record<string, string> | null;
+  stageId: string;
+  checkpointId: string;
+  checkpointStep: number | null;
+  checkpointPath?: string | null;
+  distillRuntime: Record<string, unknown>;
+  model: Record<string, unknown>;
+  evalDatasetId?: string | null;
+  configHash?: string | null;
+  parentArtifacts?: Array<Record<string, unknown>>;
+}): Promise<Record<string, unknown>[]>;
+
+export declare function evaluateDistillationCheckpoint(options: {
+  loadedWorkload: LoadedTrainingWorkload;
+  checkpointPath: string;
+  checkpointId?: string | null;
+  checkpointStep?: number | null;
+  stageId?: string | null;
+  layout?: Record<string, string> | null;
+  datasetPath?: string | null;
+  stageAArtifact?: string | null;
+  stageAArtifactHash?: string | null;
+  evalDatasetId?: string | null;
+  parentArtifacts?: Array<Record<string, unknown>>;
+}): Promise<Record<string, unknown>[]>;
+
+export declare function readDistillCheckpointMarker(markerPath: string): Promise<{
+  absolutePath: string;
+  marker: Record<string, unknown>;
+}>;

package/src/training/distillation/eval.js

@@ -0,0 +1,310 @@
+import { readFile } from 'node:fs/promises';
+import { dirname, join, resolve } from 'node:path';
+
+import { loadBackwardRegistry } from '../../config/backward-registry-loader.js';
+import { f16ToF32Array } from '../../inference/kv-cache/types.js';
+import { readBuffer, releaseBuffer } from '../../memory/buffer-pool.js';
+import { AutogradTape } from '../autograd.js';
+import { loadCheckpoint } from '../checkpoint.js';
+import { computeEvalMetrics } from '../operator-eval.js';
+import {
+  buildDistillPrompt,
+  createDistillRuntimeContext,
+  createDistillStudentRuntimeModelFixture,
+  resolveDistillDataScope,
+} from '../suite.js';
+import { restoreTrainingCheckpointState } from '../runner.js';
+import { loadCanonicalTranslationDataset } from './dataset.js';
+import { buildDistillArtifactBase, writeDistillEvalReport } from './artifacts.js';
+import { buildDistillationTrainingConfigFromWorkload, resolveInternalDistillStage } from './runtime.js';
+
+function toFloat32Array(raw, dtype = 'f32') {
+  if (dtype === 'f16') {
+    return f16ToF32Array(new Uint16Array(raw));
+  }
+  return new Float32Array(raw);
+}
+
+function resolveEvalDatasets(workload, requestedEvalDatasetId = null) {
+  const evalDatasets = Array.isArray(workload.evalDatasets) ? workload.evalDatasets : [];
+  if (!requestedEvalDatasetId) {
+    return evalDatasets;
+  }
+  return evalDatasets.filter((entry) => entry.id === requestedEvalDatasetId);
+}
+
+function argmax(values) {
+  let bestIndex = 0;
+  let bestValue = Number.NEGATIVE_INFINITY;
+  for (let index = 0; index < values.length; index += 1) {
+    const value = Number.isFinite(values[index]) ? values[index] : Number.NEGATIVE_INFINITY;
+    if (value > bestValue) {
+      bestValue = value;
+      bestIndex = index;
+    }
+  }
+  return bestIndex;
+}
+
+function releaseTensorLike(value, released) {
+  if (!value || typeof value !== 'object') return;
+  const buffer = value.buffer;
+  if (!buffer || released.has(buffer)) return;
+  released.add(buffer);
+  releaseBuffer(buffer);
+}
+
+function disposeTapeOutputs(tape, protectedBuffers = new Set()) {
+  if (!tape || !Array.isArray(tape.records)) return;
+  const released = new Set();
+  for (const record of tape.records) {
+    const output = record?.output;
+    if (!output || typeof output !== 'object') continue;
+    if (output.buffer && !protectedBuffers.has(output.buffer)) {
+      releaseTensorLike(output, released);
+      continue;
+    }
+    if (Array.isArray(output)) {
+      for (const entry of output) {
+        if (entry?.buffer && !protectedBuffers.has(entry.buffer)) {
+          releaseTensorLike(entry, released);
+        }
+      }
+    }
+  }
+}
+
+function collectProtectedBuffers(model) {
+  const protectedBuffers = new Set();
+  const groups = typeof model?.paramGroups === 'function'
+    ? model.paramGroups()
+    : {};
+  for (const params of Object.values(groups || {})) {
+    for (const tensor of Array.isArray(params) ? params : []) {
+      if (tensor?.buffer) {
+        protectedBuffers.add(tensor.buffer);
+      }
+    }
+  }
+  return protectedBuffers;
+}
+
+async function readLogitsTensor(tensor) {
+  const raw = await readBuffer(tensor.buffer);
+  return toFloat32Array(raw, tensor.dtype || 'f32');
+}
+
+async function greedyDecodeFixture(model, tokenizer, prompt, decodePolicy = {}) {
+  const maxTokens = Number.isInteger(decodePolicy?.maxTokens) && decodePolicy.maxTokens > 0
+    ? decodePolicy.maxTokens
+    : null;
+  if (!maxTokens) {
+    throw new Error('Translation eval requires evalDatasets[].decodePolicy.maxTokens in the workload pack.');
+  }
+  const stopOnEos = decodePolicy?.stopOnEos !== false;
+  const eosToken = tokenizer?.getSpecialTokens?.()?.eos ?? null;
+  const protectedBuffers = collectProtectedBuffers(model);
+  const generated = [];
+  let currentPrompt = prompt;
+  for (let step = 0; step < maxTokens; step += 1) {
+    const tape = new AutogradTape(loadBackwardRegistry());
+    let logits = null;
+    try {
+      const result = await model.forwardDistill(
+        {
+          distill: {
+            prompts: [currentPrompt],
+          },
+        },
+        tape,
+        { phase: 'anchor' }
+      );
+      logits = result?.logits || result;
+      const values = await readLogitsTensor(logits);
+      const tokenId = argmax(values);
+      if (stopOnEos && eosToken != null && tokenId === eosToken) {
+        break;
+      }
+      generated.push(tokenId);
+      currentPrompt = `${prompt}${tokenizer.decode(generated, false, false)}`;
+    } finally {
+      if (logits?.buffer && !protectedBuffers.has(logits.buffer)) {
+        releaseBuffer(logits.buffer);
+      }
+      model.cleanupDistillStep?.();
+      disposeTapeOutputs(tape, protectedBuffers);
+    }
+  }
+  return tokenizer.decode(generated, true, true);
+}
+
+function flattenMetricSummary(metrics) {
+  return {
+    bleu: metrics?.bleu?.score ?? null,
+    chrf: metrics?.chrf?.score ?? null,
+    exact_match: metrics?.exactMatch?.score ?? null,
+    accuracy: metrics?.accuracy?.score ?? null,
+    primaryMetric: metrics?.primaryMetric || null,
+    primaryScore: metrics?.primaryScore ?? null,
+  };
+}
+
+export async function evaluateDistillationModel(options) {
+  const loadedWorkload = options.loadedWorkload;
+  const workload = loadedWorkload.workload;
+  const stageId = options.stageId;
+  const checkpointId = options.checkpointId;
+  const checkpointStep = options.checkpointStep;
+  const distillRuntime = options.distillRuntime;
+  const model = options.model;
+  if (distillRuntime.studentGraphMode !== 'transformer_full') {
+    throw new Error(
+      `Distillation eval requires studentGraphMode="transformer_full"; got "${distillRuntime.studentGraphMode}".`
+    );
+  }
+  const evalDatasets = resolveEvalDatasets(workload, options.evalDatasetId || null);
+  if (evalDatasets.length === 0) {
+    throw new Error(`No eval datasets resolved for workload "${workload.id}".`);
+  }
+  const reports = [];
+  for (const evalDataset of evalDatasets) {
+    if (evalDataset.evalKind !== 'translation') {
+      throw new Error(`Distillation eval currently supports translation eval only, got "${evalDataset.evalKind}".`);
+    }
+    const dataset = await loadCanonicalTranslationDataset(evalDataset.datasetPath, {
+      strictPairContract: workload.pipeline.strictPairContract === true,
+      sourceLangs: evalDataset.sourceLangs || workload.pipeline.sourceLangs,
+      targetLangs: evalDataset.targetLangs || workload.pipeline.targetLangs,
+      pairAllowlist: evalDataset.pairAllowlist || workload.pipeline.pairAllowlist,
+    });
+    const hypotheses = [];
+    const references = [];
+    const samples = [];
+    for (const row of dataset.rows) {
+      const prompt = buildDistillPrompt({
+        direction: row.pair || (
+          row.src_lang && row.tgt_lang
+            ? `${row.src_lang}->${row.tgt_lang}`
+            : 'unknown'
+        ),
+        source: row.source,
+      });
+      const hypothesis = await greedyDecodeFixture(
+        model,
+        distillRuntime.studentPipeline.tokenizer,
+        prompt,
+        evalDataset.decodePolicy || {}
+      );
+      hypotheses.push(hypothesis);
+      references.push(row.target_pos);
+      if (samples.length < 5) {
+        samples.push({
+          row_id: row.row_id,
+          source: row.source,
+          reference: row.target_pos,
+          hypothesis,
+        });
+      }
+    }
+    const computedMetrics = computeEvalMetrics('translation', hypotheses, references, {});
+    const flattened = flattenMetricSummary(computedMetrics);
+    const reportPayload = {
+      ...buildDistillArtifactBase(loadedWorkload, {
+        prefix: 'dst_eval',
+        artifactType: 'training_eval_report',
+        datasetPath: dataset.absolutePath,
+        datasetHash: dataset.canonicalHash,
+        stage: stageId,
+        checkpointStep,
+        parentArtifacts: options.parentArtifacts || [],
+        configHash: options.configHash || workload.configHash,
+      }),
+      checkpointId,
+      checkpointPath: options.checkpointPath || null,
+      evalDatasetId: evalDataset.id,
+      evalKind: evalDataset.evalKind,
+      metrics: computedMetrics,
+      ...flattened,
+      rowCount: dataset.rowCount,
+      sampleRows: samples,
+    };
+    const reportFile = options.layout
+      ? await writeDistillEvalReport(options.layout, reportPayload)
+      : null;
+    reports.push({
+      ...reportPayload,
+      reportPath: reportFile?.path || null,
+    });
+  }
+  return reports;
+}
+
+export async function evaluateDistillationCheckpoint(options) {
+  const loadedWorkload = options.loadedWorkload;
+  const workload = loadedWorkload.workload;
+  const stagePlan = workload.pipeline.stagePlan;
+  const requestedStageId = String(options.stageId || '').trim();
+  const stageEntry = stagePlan.find((entry) => entry.id === requestedStageId)
+    || stagePlan[0];
+  if (!stageEntry) {
+    throw new Error(`No stage entry resolved for workload "${workload.id}".`);
+  }
+  const internalStage = resolveInternalDistillStage(stageEntry);
+  const checkpointPath = resolve(String(options.checkpointPath));
+  const checkpointRecord = await loadCheckpoint(checkpointPath);
+  if (!checkpointRecord) {
+    throw new Error(`Checkpoint not found: ${checkpointPath}`);
+  }
+  const configBundle = buildDistillationTrainingConfigFromWorkload(loadedWorkload, stageEntry, {
+    datasetPath: options.datasetPath || workload.datasetPath,
+    stageAArtifact: options.stageAArtifact || null,
+    stageAArtifactHash: options.stageAArtifactHash || null,
+    artifactDir: dirname(dirname(checkpointPath)),
+  });
+  const distillRuntime = await createDistillRuntimeContext({
+    teacherModelId: workload.teacherModelId,
+    studentModelId: workload.studentModelId,
+    trainingStage: internalStage,
+    studentGraphMode: workload.pipeline.studentGraphMode,
+  }, configBundle.trainingConfig.training);
+  let fixture = null;
+  try {
+    fixture = await createDistillStudentRuntimeModelFixture({
+      training: configBundle.trainingConfig.training,
+    }, {
+      distillRuntime,
+      studentGraphMode: workload.pipeline.studentGraphMode,
+    });
+    await restoreTrainingCheckpointState(
+      fixture.model,
+      { getState: () => null, stepCount: 0 },
+      checkpointRecord,
+      fixture.config
+    );
+    return evaluateDistillationModel({
+      loadedWorkload,
+      layout: options.layout || null,
+      stageId: stageEntry.id,
+      checkpointId: options.checkpointId || 'checkpoint',
+      checkpointStep: options.checkpointStep || null,
+      checkpointPath,
+      distillRuntime,
+      model: fixture.model,
+      evalDatasetId: options.evalDatasetId || null,
+      configHash: configBundle.trainingConfigHash,
+      parentArtifacts: options.parentArtifacts || [],
+    });
+  } finally {
+    fixture?.cleanup?.();
+    await distillRuntime.cleanup();
+  }
+}
+
+export async function readDistillCheckpointMarker(markerPath) {
+  const absolutePath = resolve(String(markerPath));
+  const raw = await readFile(absolutePath, 'utf8');
+  return {
+    absolutePath,
+    marker: JSON.parse(raw),
+  };
+}

package/src/training/distillation/index.d.ts

@@ -0,0 +1,29 @@
+export {
+  normalizeDistillationPair,
+  normalizeTranslationPairRow,
+  loadCanonicalTranslationDataset,
+  buildFrozenSubset,
+} from './dataset.js';
+export {
+  createDistillationRunArtifacts,
+  writeDistillStageManifest,
+  writeDistillCheckpointMetadata,
+  writeDistillCheckpointComplete,
+  writeDistillEvalReport,
+  writeDistillCompareReport,
+  writeDistillQualityGateReport,
+  buildDistillArtifactBase,
+} from './artifacts.js';
+export { appendDistillationScoreboardRow } from './scoreboard.js';
+export {
+  buildDistillationTrainingConfigFromWorkload,
+  resolveInternalDistillStage,
+} from './runtime.js';
+export {
+  evaluateDistillationModel,
+  evaluateDistillationCheckpoint,
+  readDistillCheckpointMarker,
+} from './eval.js';
+export { runDistillationStage, runDistillationStageA } from './stage-a.js';
+export { runDistillationStageB } from './stage-b.js';
+export { watchDistillationCheckpoints } from './checkpoint-watch.js';

package/src/training/distillation/index.js

@@ -0,0 +1,29 @@
+export {
+  normalizeDistillationPair,
+  normalizeTranslationPairRow,
+  loadCanonicalTranslationDataset,
+  buildFrozenSubset,
+} from './dataset.js';
+export {
+  createDistillationRunArtifacts,
+  writeDistillStageManifest,
+  writeDistillCheckpointMetadata,
+  writeDistillCheckpointComplete,
+  writeDistillEvalReport,
+  writeDistillCompareReport,
+  writeDistillQualityGateReport,
+  buildDistillArtifactBase,
+} from './artifacts.js';
+export { appendDistillationScoreboardRow } from './scoreboard.js';
+export {
+  buildDistillationTrainingConfigFromWorkload,
+  resolveInternalDistillStage,
+} from './runtime.js';
+export {
+  evaluateDistillationModel,
+  evaluateDistillationCheckpoint,
+  readDistillCheckpointMarker,
+} from './eval.js';
+export { runDistillationStage, runDistillationStageA } from './stage-a.js';
+export { runDistillationStageB } from './stage-b.js';
+export { watchDistillationCheckpoints } from './checkpoint-watch.js';

package/src/training/distillation/runtime.d.ts

@@ -0,0 +1,20 @@
+import type { LoadedTrainingWorkload, DistillStagePlanEntry } from '../workloads.js';
+
+export declare function resolveInternalDistillStage(
+  stageEntry: DistillStagePlanEntry | Record<string, unknown>
+): 'stage_a' | 'stage_b';
+
+export declare function buildDistillationTrainingConfigFromWorkload(
+  loadedWorkload: LoadedTrainingWorkload,
+  stageEntry: DistillStagePlanEntry | Record<string, unknown>,
+  options?: {
+    datasetPath?: string | null;
+    artifactDir?: string | null;
+    stageAArtifact?: string | null;
+    stageAArtifactHash?: string | null;
+  }
+): {
+  internalStage: 'stage_a' | 'stage_b';
+  trainingConfig: Record<string, unknown>;
+  trainingConfigHash: string;
+};

package/src/training/distillation/runtime.js

@@ -0,0 +1,121 @@
+import { createTrainingConfig } from '../../config/training-defaults.js';
+import { sha256Hex } from '../../utils/sha256.js';
+
+function stableSortObject(value) {
+  if (Array.isArray(value)) {
+    return value.map((entry) => stableSortObject(entry));
+  }
+  if (!value || typeof value !== 'object') {
+    return value;
+  }
+  const sorted = {};
+  for (const key of Object.keys(value).sort()) {
+    sorted[key] = stableSortObject(value[key]);
+  }
+  return sorted;
+}
+
+function stableJson(value) {
+  return JSON.stringify(stableSortObject(value));
+}
+
+function normalizeStageLabel(value) {
+  const normalized = String(value || '').trim().toLowerCase().replace(/[\s-]+/g, '_');
+  return normalized;
+}
+
+export function resolveInternalDistillStage(stageEntry) {
+  const trainingStage = normalizeStageLabel(stageEntry?.trainingStage || stageEntry?.id || '');
+  const objective = normalizeStageLabel(stageEntry?.objective || '');
+  if (trainingStage === 'sft' || objective === 'sft') {
+    throw new Error(
+      'Distillation workload stage uses "sft", but the current JS distill runner only supports the KD-oriented stage_a contract. Use objective="kd" / trainingStage="stage_a" explicitly.'
+    );
+  }
+  if (
+    trainingStage === 'stage_b'
+    || trainingStage === 'post_sft_distill'
+    || trainingStage === 'post_sft_triplet'
+    || objective === 'triplet'
+  ) {
+    return 'stage_b';
+  }
+  if (
+    trainingStage === 'stage_a'
+    || trainingStage === 'kd'
+    || objective === 'kd'
+    || objective === 'cross_entropy'
+  ) {
+    return 'stage_a';
+  }
+  throw new Error(
+    `Unsupported distillation stage "${stageEntry?.trainingStage || stageEntry?.id || 'unknown'}".`
+  );
+}
+
+export function buildDistillationTrainingConfigFromWorkload(loadedWorkload, stageEntry, options = {}) {
+  const workload = loadedWorkload.workload;
+  if (workload.kind !== 'distill') {
+    throw new Error('buildDistillationTrainingConfigFromWorkload requires a distill workload.');
+  }
+  const internalStage = resolveInternalDistillStage(stageEntry);
+  const distillTraining = {
+    enabled: true,
+    stage: internalStage,
+    teacherModelId: workload.teacherModelId,
+    studentModelId: workload.studentModelId,
+    datasetId: workload.datasetId,
+    datasetPath: options.datasetPath || workload.datasetPath,
+    sourceLangs: workload.pipeline.sourceLangs,
+    targetLangs: workload.pipeline.targetLangs,
+    pairAllowlist: workload.pipeline.pairAllowlist,
+    strictPairContract: workload.pipeline.strictPairContract === true,
+    stageAArtifact: options.stageAArtifact || null,
+    stageAArtifactHash: options.stageAArtifactHash || null,
+    artifactDir: options.artifactDir || null,
+    temperature: workload.pipeline.temperature,
+    alphaKd: workload.pipeline.alphaKd,
+    alphaCe: workload.pipeline.alphaCe,
+    tripletMargin: workload.pipeline.tripletMargin,
+    studentGraphMode: workload.pipeline.studentGraphMode,
+  };
+  if (internalStage === 'stage_b') {
+    distillTraining.freeze = {
+      encoder: true,
+      prior: true,
+      decoder: true,
+      base: false,
+      lora: false,
+    };
+  }
+  const trainingConfig = createTrainingConfig({
+    training: {
+      enabled: true,
+      optimizer: {
+        type: workload.training.optimizer.type,
+        lr: workload.training.optimizer.lr,
+        beta1: workload.training.optimizer.beta1,
+        beta2: workload.training.optimizer.beta2,
+        eps: workload.training.optimizer.eps,
+        weightDecay: workload.training.optimizer.weightDecay,
+        scheduler: workload.training.optimizer.scheduler,
+      },
+      gradient: {
+        maxNorm: workload.training.gradientClipping.maxNorm,
+      },
+      precision: workload.training.precision,
+      distill: distillTraining,
+    },
+  });
+  return {
+    internalStage,
+    trainingConfig,
+    trainingConfigHash: sha256Hex(stableJson({
+      workloadConfigHash: workload.configHash,
+      stageEntry,
+      datasetPath: options.datasetPath || workload.datasetPath,
+      stageAArtifact: options.stageAArtifact || null,
+      stageAArtifactHash: options.stageAArtifactHash || null,
+    })),
+  };
+}

package/src/training/distillation/scoreboard.d.ts

@@ -0,0 +1,6 @@
+export declare function appendDistillationScoreboardRow(
+  layout: Record<string, string>,
+  stageId: string,
+  row: Record<string, unknown>,
+  options?: { selectionMetric?: string | null; selectionGoal?: string | null }
+): Promise<{ rowsPath: string; summaryPath: string }>;

package/src/training/distillation/scoreboard.js

@@ -0,0 +1,8 @@
+import { join } from 'node:path';
+
+import { appendScoreboardRow } from '../operator-scoreboard.js';
+
+export async function appendDistillationScoreboardRow(layout, stageId, row, options = {}) {
+  const scoreboardDir = join(layout.scoreboard, String(stageId || 'stage'));
+  return appendScoreboardRow(scoreboardDir, row, options);
+}

package/src/training/distillation/stage-a.d.ts

@@ -0,0 +1,45 @@
+import type { LoadedTrainingWorkload, DistillStagePlanEntry } from '../workloads.js';
+
+export declare function runDistillationStage(options: {
+  loadedWorkload: LoadedTrainingWorkload;
+  stageEntry: DistillStagePlanEntry;
+  layout: Record<string, string>;
+  datasetPath?: string | null;
+  stageAArtifact?: string | null;
+  stageAArtifactHash?: string | null;
+  legacyArtifactDir?: string | null;
+  timestamp?: string | Date | null;
+  parentArtifacts?: Array<Record<string, unknown>>;
+}): Promise<{
+  stageId: string;
+  trainingStage: 'stage_a' | 'stage_b';
+  metrics: Record<string, unknown>[];
+  checkpointArtifacts: Array<Record<string, unknown>>;
+  evalReports: Array<Record<string, unknown>>;
+  bestReport: Record<string, unknown> | null;
+  stageManifestPath: string;
+  legacyArtifact: Record<string, unknown> | null;
+  lastCheckpoint: Record<string, unknown> | null;
+}>;
+
+export declare function runDistillationStageA(options: {
+  loadedWorkload: LoadedTrainingWorkload;
+  stageEntry: DistillStagePlanEntry;
+  layout: Record<string, string>;
+  datasetPath?: string | null;
+  stageAArtifact?: string | null;
+  stageAArtifactHash?: string | null;
+  legacyArtifactDir?: string | null;
+  timestamp?: string | Date | null;
+  parentArtifacts?: Array<Record<string, unknown>>;
+}): Promise<{
+  stageId: string;
+  trainingStage: 'stage_a' | 'stage_b';
+  metrics: Record<string, unknown>[];
+  checkpointArtifacts: Array<Record<string, unknown>>;
+  evalReports: Array<Record<string, unknown>>;
+  bestReport: Record<string, unknown> | null;
+  stageManifestPath: string;
+  legacyArtifact: Record<string, unknown> | null;
+  lastCheckpoint: Record<string, unknown> | null;
+}>;