@simulatte/doppler 0.1.4 → 0.1.6

package/src/training/distillation/dataset.js

@@ -0,0 +1,337 @@
+import { mkdir, readFile, writeFile } from 'node:fs/promises';
+import { dirname, join, resolve } from 'node:path';
+
+import { parseJsonl } from '../datasets/jsonl.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 asOptionalString(value) {
+  if (value === undefined || value === null) return null;
+  const trimmed = String(value).trim();
+  return trimmed || null;
+}
+
+function asOptionalStringArray(value) {
+  if (value === undefined || value === null) return null;
+  const input = Array.isArray(value)
+    ? value
+    : (typeof value === 'string' ? value.split(',') : null);
+  if (!Array.isArray(input)) {
+    return null;
+  }
+  const normalized = input
+    .map((entry) => asOptionalString(entry))
+    .filter(Boolean);
+  return normalized.length > 0 ? normalized : null;
+}
+
+function normalizeLangCode(value) {
+  const normalized = asOptionalString(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;
+}
+
+export function normalizeDistillationPair(value, srcLang = null, tgtLang = null) {
+  const pair = asOptionalString(value);
+  if (!pair) {
+    if (!srcLang || !tgtLang) return null;
+    return `${srcLang}->${tgtLang}`;
+  }
+  const normalized = pair.toLowerCase().replace(/_/g, '-').replace(/\s+/g, '');
+  const separator = normalized.includes('->') ? '->' : '-';
+  const parts = normalized.split(separator).filter(Boolean);
+  if (parts.length !== 2) return null;
+  const source = normalizeLangCode(parts[0]) || parts[0];
+  const target = normalizeLangCode(parts[1]) || parts[1];
+  return `${source}->${target}`;
+}
+
+function resolveStringCandidate(record, keys) {
+  for (const key of keys) {
+    const value = asOptionalString(record?.[key]);
+    if (value) return value;
+  }
+  return null;
+}
+
+function resolveStableRowId(record, index, canonical) {
+  const explicit = asOptionalString(record?.row_id ?? record?.rowId);
+  if (explicit) return explicit;
+  return sha256Hex(stableJson({
+    index,
+    src_lang: canonical.src_lang,
+    tgt_lang: canonical.tgt_lang,
+    pair: canonical.pair,
+    source: canonical.source,
+    target_pos: canonical.target_pos,
+    target_neg: canonical.target_neg,
+  }));
+}
+
+export function normalizeTranslationPairRow(record, index, options = {}) {
+  if (!record || typeof record !== 'object' || Array.isArray(record)) {
+    return null;
+  }
+  const source = resolveStringCandidate(record, ['source', 'query', 'prompt']);
+  const targetPos = resolveStringCandidate(record, ['target_pos', 'target', 'pos', 'completion']);
+  const targetNeg = resolveStringCandidate(record, ['target_neg', 'neg']);
+  if (!source || !targetPos) {
+    return null;
+  }
+  const srcLang = normalizeLangCode(record?.src_lang ?? record?.source_lang);
+  const tgtLang = normalizeLangCode(record?.tgt_lang ?? record?.target_lang ?? record?.lang);
+  const pair = normalizeDistillationPair(record?.pair, srcLang, tgtLang);
+  const strict = options.strictPairContract === true;
+  if (strict) {
+    if (!srcLang || !tgtLang) {
+      throw new Error('strictPairContract requires src_lang and tgt_lang on each row.');
+    }
+    if (!pair) {
+      throw new Error('strictPairContract requires pair on each row.');
+    }
+    if (pair !== `${srcLang}->${tgtLang}`) {
+      throw new Error(`row pair "${record?.pair}" does not match src/tgt "${srcLang}->${tgtLang}".`);
+    }
+  }
+  const canonical = {
+    src_lang: srcLang,
+    tgt_lang: tgtLang,
+    pair: pair || null,
+    source,
+    target_pos: targetPos,
+    target_neg: targetNeg,
+  };
+  canonical.row_id = resolveStableRowId(record, index, canonical);
+  return canonical;
+}
+
+function normalizeFilterSet(value, normalizer) {
+  const entries = asOptionalStringArray(value);
+  if (!entries) return null;
+  const normalized = entries
+    .map((entry) => normalizer(entry))
+    .filter(Boolean);
+  return normalized.length > 0 ? new Set(normalized) : null;
+}
+
+function applyDatasetFilters(rows, options = {}) {
+  const sourceLangs = normalizeFilterSet(options.sourceLangs, normalizeLangCode);
+  const targetLangs = normalizeFilterSet(options.targetLangs, normalizeLangCode);
+  const pairs = normalizeFilterSet(options.pairAllowlist, normalizeDistillationPair);
+  return rows.filter((row) => {
+    if (sourceLangs && (!row.src_lang || !sourceLangs.has(row.src_lang))) {
+      return false;
+    }
+    if (targetLangs && (!row.tgt_lang || !targetLangs.has(row.tgt_lang))) {
+      return false;
+    }
+    if (pairs && (!row.pair || !pairs.has(row.pair))) {
+      return false;
+    }
+    return true;
+  });
+}
+
+function normalizeSubsetSpec(value) {
+  if (!value || typeof value !== 'object' || Array.isArray(value)) {
+    return null;
+  }
+  const sizeRaw = Number(value.size ?? value.count ?? value.rowCount ?? 0);
+  const size = Number.isInteger(sizeRaw) && sizeRaw > 0 ? sizeRaw : null;
+  const seedRaw = Number(value.seed ?? 1337);
+  const seed = Number.isInteger(seedRaw) ? seedRaw : 1337;
+  const balanceBy = asOptionalString(value.balanceBy ?? value.allocationMode ?? value.stratifyBy);
+  return {
+    id: asOptionalString(value.id ?? value.name) || null,
+    size,
+    seed,
+    balanceBy,
+    parentSubsetManifest: asOptionalString(
+      value.parentSubsetManifest
+      ?? value.parentSubset
+      ?? value.parentManifest
+    ),
+  };
+}
+
+function buildDeterministicRank(seed, rowId) {
+  return sha256Hex(`${seed}:${rowId}`);
+}
+
+async function resolveParentRowIds(parentSubsetManifest) {
+  if (!parentSubsetManifest) return null;
+  const absolutePath = resolve(parentSubsetManifest);
+  const raw = await readFile(absolutePath, 'utf8');
+  const parsed = JSON.parse(raw);
+  const rowIds = Array.isArray(parsed?.rowIds)
+    ? parsed.rowIds
+    : [];
+  return rowIds.length > 0 ? new Set(rowIds.map((entry) => String(entry))) : new Set();
+}
+
+function buildPairBalancedSubset(rows, size, seed) {
+  const byPair = new Map();
+  for (const row of rows) {
+    const key = row.pair || 'unknown';
+    const bucket = byPair.get(key) || [];
+    bucket.push(row);
+    byPair.set(key, bucket);
+  }
+  for (const bucket of byPair.values()) {
+    bucket.sort((left, right) => {
+      const leftRank = buildDeterministicRank(seed, left.row_id);
+      const rightRank = buildDeterministicRank(seed, right.row_id);
+      return leftRank.localeCompare(rightRank);
+    });
+  }
+  const pairKeys = [...byPair.keys()].sort((left, right) => left.localeCompare(right));
+  const selected = [];
+  let cursor = 0;
+  while (selected.length < size) {
+    let progressed = false;
+    for (const pairKey of pairKeys) {
+      const bucket = byPair.get(pairKey);
+      if (!bucket || cursor >= bucket.length) continue;
+      selected.push(bucket[cursor]);
+      progressed = true;
+      if (selected.length >= size) break;
+    }
+    if (!progressed) break;
+    cursor += 1;
+  }
+  return selected;
+}
+
+function selectSubsetRows(rows, subsetSpec) {
+  if (!subsetSpec || !subsetSpec.size || subsetSpec.size >= rows.length) {
+    return rows.slice();
+  }
+  if (subsetSpec.balanceBy === 'pair' || subsetSpec.balanceBy === 'pair_balance') {
+    return buildPairBalancedSubset(rows, subsetSpec.size, subsetSpec.seed);
+  }
+  return rows
+    .slice()
+    .sort((left, right) => {
+      const leftRank = buildDeterministicRank(subsetSpec.seed, left.row_id);
+      const rightRank = buildDeterministicRank(subsetSpec.seed, right.row_id);
+      return leftRank.localeCompare(rightRank);
+    })
+    .slice(0, subsetSpec.size);
+}
+
+function computeDirectionCounts(rows) {
+  const counts = {};
+  for (const row of rows) {
+    const key = row.pair || 'unknown';
+    counts[key] = (counts[key] || 0) + 1;
+  }
+  return counts;
+}
+
+export async function loadCanonicalTranslationDataset(datasetPath, options = {}) {
+  const absolutePath = resolve(String(datasetPath));
+  const raw = await readFile(absolutePath, 'utf8');
+  const parsed = absolutePath.endsWith('.json')
+    ? JSON.parse(raw)
+    : parseJsonl(raw);
+  if (!Array.isArray(parsed)) {
+    throw new Error(`Distillation dataset "${absolutePath}" must be a JSON array or JSONL file.`);
+  }
+  const normalizedRows = [];
+  for (let index = 0; index < parsed.length; index += 1) {
+    const row = normalizeTranslationPairRow(parsed[index], index, options);
+    if (row) {
+      normalizedRows.push(row);
+    }
+  }
+  const filteredRows = applyDatasetFilters(normalizedRows, options);
+  if (filteredRows.length === 0) {
+    throw new Error(`Distillation dataset "${absolutePath}" has no usable rows after contract checks and filters.`);
+  }
+  const rowIds = filteredRows.map((row) => row.row_id);
+  return {
+    absolutePath,
+    raw,
+    rows: filteredRows,
+    rowCount: filteredRows.length,
+    directionCounts: computeDirectionCounts(filteredRows),
+    datasetHash: sha256Hex(raw),
+    canonicalHash: sha256Hex(stableJson(filteredRows)),
+    rowIdsHash: sha256Hex(rowIds.join('\n')),
+  };
+}
+
+export async function buildFrozenSubset(options) {
+  const dataset = await loadCanonicalTranslationDataset(options.datasetPath, {
+    strictPairContract: options.strictPairContract === true,
+    sourceLangs: options.sourceLangs,
+    targetLangs: options.targetLangs,
+    pairAllowlist: options.pairAllowlist,
+  });
+  const subsetSpec = normalizeSubsetSpec(options.subsetSpec);
+  const parentRowIds = subsetSpec?.parentSubsetManifest
+    ? await resolveParentRowIds(subsetSpec.parentSubsetManifest)
+    : null;
+  const scopedRows = parentRowIds
+    ? dataset.rows.filter((row) => parentRowIds.has(row.row_id))
+    : dataset.rows;
+  const subsetRows = selectSubsetRows(scopedRows, subsetSpec);
+  const outputDir = resolve(String(options.outputDir));
+  const subsetJsonlPath = join(outputDir, 'subset.jsonl');
+  const rowIdsPath = join(outputDir, 'row_ids.txt');
+  const manifestPath = join(outputDir, 'subset_manifest.json');
+  const serializedRows = `${subsetRows.map((row) => JSON.stringify(row)).join('\n')}\n`;
+  const rowIdsText = `${subsetRows.map((row) => row.row_id).join('\n')}\n`;
+  const manifest = {
+    artifactType: 'subset_manifest',
+    schemaVersion: 1,
+    generatedAt: new Date().toISOString(),
+    datasetPath: dataset.absolutePath,
+    datasetHash: dataset.datasetHash,
+    canonicalHash: dataset.canonicalHash,
+    rowIdsHash: sha256Hex(rowIdsText),
+    universeRowCount: dataset.rowCount,
+    subsetRowCount: subsetRows.length,
+    directionCounts: computeDirectionCounts(subsetRows),
+    subsetSpec,
+    parentSubsetManifest: subsetSpec?.parentSubsetManifest || null,
+    rowIds: subsetRows.map((row) => row.row_id),
+    output: {
+      subsetJsonlPath,
+      rowIdsPath,
+    },
+  };
+  await mkdir(outputDir, { recursive: true });
+  await writeFile(subsetJsonlPath, serializedRows, 'utf8');
+  await writeFile(rowIdsPath, rowIdsText, 'utf8');
+  await mkdir(dirname(manifestPath), { recursive: true });
+  await writeFile(manifestPath, `${JSON.stringify(manifest, null, 2)}\n`, 'utf8');
+  return {
+    dataset,
+    subsetRows,
+    subsetJsonlPath,
+    rowIdsPath,
+    manifestPath,
+    manifest,
+  };
+}

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),
+  };
+}