@simulatte/doppler 0.1.8 → 0.2.0

package/src/training/tensor-factory.d.ts

@@ -1,9 +0,0 @@
-import type { Tensor, TensorDType } from '../gpu/tensor.js';
-
-export declare function createUploadedTensor(
-  data: ArrayBufferView,
-  dtype: TensorDType,
-  shape: number[],
-  label: string,
-  usage?: number | undefined
-): Tensor;

package/src/training/tensor-factory.js

@@ -1,13 +0,0 @@
-import { acquireBuffer, releaseBuffer, uploadData } from '../memory/buffer-pool.js';
-import { createTensor } from '../gpu/tensor.js';
-
-export function createUploadedTensor(data, dtype, shape, label, usage = undefined) {
-  const buffer = acquireBuffer(data.byteLength, usage, label);
-  try {
-    uploadData(buffer, data);
-    return createTensor(buffer, dtype, [...shape], label);
-  } catch (error) {
-    releaseBuffer(buffer);
-    throw error;
-  }
-}

package/src/training/trainer.d.ts

@@ -1,89 +0,0 @@
-import type { BackwardRegistrySchema } from '../config/schema/backward-registry.schema.js';
-import type { TrainingConfigSchema } from '../config/training-defaults.d.ts';
-import type { Tensor } from '../gpu/tensor.js';
-import type { TrainingObjective } from './objectives/base.js';
-
-export interface TrainingBatch {
-  input: Tensor;
-  targets: Tensor;
-}
-
-export interface TrainingOptimizer {
-  step: (
-    params: Tensor[],
-    grads: Map<Tensor, Tensor>,
-    config: TrainingConfigSchema,
-    context?: {
-      trainableGroups?: string[];
-      frozenGroups?: string[];
-      allGroups?: string[];
-    }
-  ) => Promise<OptimizerMetrics>;
-}
-
-export interface TrainStepOptions {
-  registry?: BackwardRegistrySchema;
-  crossEntropyLoss?: (
-    logits: Tensor,
-    targets: Tensor,
-    config: TrainingConfigSchema,
-    tape: unknown
-  ) => Promise<Tensor>;
-  clipGradients?: (
-    grads: Map<Tensor, Tensor>,
-    config: TrainingConfigSchema
-  ) => Promise<ClipMetrics>;
-  optimizer?: TrainingOptimizer;
-  lossScale?: number;
-  applyClip?: boolean;
-  applyOptimizer?: boolean;
-  trainingObjective?: TrainingObjective;
-  stepIndex?: number | null;
-  epochIndex?: number | null;
-  batchIndex?: number | null;
-  stage1ArtifactContext?: Record<string, unknown> | null;
-  stageAArtifactContext?: Record<string, unknown> | null;
-}
-
-export interface ClipMetrics {
-  clippedGrads: Map<Tensor, Tensor>;
-  gradient_norm_unclipped: number;
-  gradient_norm_clipped: number;
-  clipped_event_count: number;
-  total_param_count: number;
-}
-
-export interface OptimizerMetrics {
-  optimizer_ms: number;
-  effective_lr?: number | null;
-  scheduler_index?: number | null;
-  scheduler_phase?: string | null;
-}
-
-export interface TrainStepResult {
-  loss: Tensor;
-  grads: Map<Tensor, Tensor>;
-  forward_ms?: number;
-  backward_ms?: number;
-  clipMetrics?: ClipMetrics;
-  optimizerMetrics?: OptimizerMetrics;
-  objectiveName?: string;
-  objectiveMetrics?: Record<string, number | string | boolean | null | undefined>;
-  paramGroupMetrics?: {
-    trainableGroups: string[];
-    frozenGroups: string[];
-    allGroups: string[];
-    trainableParamCount: number;
-  };
-}
-
-export declare function trainStep(
-  model: {
-    forward: (input: Tensor, tape: unknown) => Promise<Tensor>;
-    loraParams?: () => Tensor[];
-    paramGroups?: () => Record<string, Tensor[]>;
-  },
-  batch: TrainingBatch,
-  config: TrainingConfigSchema,
-  options?: TrainStepOptions
-): Promise<TrainStepResult>;

package/src/training/trainer.js

@@ -1,299 +0,0 @@
-import { AutogradTape } from './autograd.js';
-import { loadBackwardRegistry } from '../config/backward-registry-loader.js';
-import { runScale } from '../gpu/kernels/index.js';
-import { releaseBuffer } from '../memory/buffer-pool.js';
-import { createCrossEntropyObjective } from './objectives/cross_entropy.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);
-  return createUploadedTensor(gradData, 'f32', loss.shape, 'loss_grad_output');
-}
-
-function normalizeLossResult(value) {
-  if (value && typeof value === 'object' && value.loss) {
-    return value;
-  }
-  return { loss: value, components: null };
-}
-
-function resolveTrainingObjective(options) {
-  if (options.trainingObjective) {
-    return options.trainingObjective;
-  }
-  return createCrossEntropyObjective({ crossEntropyLoss: options.crossEntropyLoss });
-}
-
-function resolveModelParamGroups(model) {
-  if (model && typeof model.paramGroups === 'function') {
-    const groups = model.paramGroups();
-    if (!groups || typeof groups !== 'object') {
-      throw new Error('model.paramGroups() must return an object of tensor arrays.');
-    }
-    return groups;
-  }
-  if (model && typeof model.loraParams === 'function') {
-    return { lora: model.loraParams() };
-  }
-  return {};
-}
-
-function selectTrainableParamGroups(paramGroups, freezeMap) {
-  const trainableGroups = {};
-  const frozenGroups = [];
-  for (const [groupName, params] of Object.entries(paramGroups)) {
-    const normalizedParams = Array.isArray(params) ? params.filter(Boolean) : [];
-    if (freezeMap?.[groupName] === true) {
-      frozenGroups.push(groupName);
-      continue;
-    }
-    trainableGroups[groupName] = normalizedParams;
-  }
-  return { trainableGroups, frozenGroups };
-}
-
-function flattenUniqueParams(paramGroups) {
-  const unique = new Set();
-  const output = [];
-  for (const params of Object.values(paramGroups)) {
-    for (const tensor of params) {
-      if (!tensor || unique.has(tensor)) continue;
-      unique.add(tensor);
-      output.push(tensor);
-    }
-  }
-  return output;
-}
-
-function isTensorLike(value) {
-  return !!value
-    && typeof value === 'object'
-    && Array.isArray(value.shape)
-    && value.buffer != null;
-}
-
-function normalizeBackwardSeedPayload(payload, loss, lossScale) {
-  const collectSeedTensors = (value) => {
-    const seedTensors = [];
-    if (value instanceof Map) {
-      for (const grad of value.values()) {
-        if (isTensorLike(grad)) {
-          seedTensors.push(grad);
-        }
-      }
-      return seedTensors;
-    }
-    if (Array.isArray(value)) {
-      for (const seed of value) {
-        if (isTensorLike(seed?.grad)) {
-          seedTensors.push(seed.grad);
-        }
-      }
-      return seedTensors;
-    }
-    if (value && typeof value === 'object' && Array.isArray(value.seeds)) {
-      for (const seed of value.seeds) {
-        if (isTensorLike(seed?.grad)) {
-          seedTensors.push(seed.grad);
-        }
-      }
-      return seedTensors;
-    }
-    if (isTensorLike(value)) {
-      seedTensors.push(value);
-    }
-    return seedTensors;
-  };
-
-  if (payload instanceof Map) {
-    return { seedArg: payload, seedTensors: collectSeedTensors(payload) };
-  }
-  if (Array.isArray(payload)) {
-    return { seedArg: payload, seedTensors: collectSeedTensors(payload) };
-  }
-  if (payload && typeof payload === 'object' && Array.isArray(payload.seeds)) {
-    return { seedArg: payload, seedTensors: collectSeedTensors(payload) };
-  }
-  if (isTensorLike(payload)) {
-    return { seedArg: payload, seedTensors: [payload] };
-  }
-  const fallback = createLossGradient(loss, lossScale);
-  return { seedArg: fallback, seedTensors: [fallback] };
-}
-
-function filterGradientsForParams(grads, params) {
-  const filtered = new Map();
-  for (const param of params) {
-    const grad = grads.get(param);
-    if (grad) {
-      filtered.set(param, grad);
-    }
-  }
-  return filtered;
-}
-
-export async function trainStep(
-  model,
-  batch,
-  config,
-  options = {}
-) {
-  const {
-    registry = loadBackwardRegistry(),
-    crossEntropyLoss,
-    clipGradients,
-    optimizer,
-    lossScale = 1,
-    applyClip = true,
-    applyOptimizer = true,
-    trainingObjective = null,
-  } = options;
-
-  if (!crossEntropyLoss || !clipGradients || !optimizer) {
-    throw new Error('trainStep requires crossEntropyLoss, clipGradients, and optimizer');
-  }
-
-  const objective = resolveTrainingObjective({
-    trainingObjective,
-    crossEntropyLoss,
-  });
-  const tape = new AutogradTape(registry);
-  let preparedBatch = batch;
-  if (typeof objective.prepareBatch === 'function') {
-    const nextBatch = await objective.prepareBatch({
-      model,
-      batch,
-      config,
-      tape,
-      options,
-      lossScale,
-    });
-    if (nextBatch && typeof nextBatch === 'object') {
-      preparedBatch = nextBatch;
-    }
-  }
-
-  const t0 = globalThis.performance.now();
-  const forwardState = await objective.forward({
-    model,
-    batch: preparedBatch,
-    config,
-    tape,
-    options,
-    lossScale,
-  });
-  const lossResult = normalizeLossResult(await objective.computeLoss({
-    model,
-    batch: preparedBatch,
-    config,
-    tape,
-    options,
-    lossScale,
-    forwardState,
-  }));
-  const loss = lossResult.loss;
-  const gradOutput = typeof objective.backwardTargets === 'function'
-    ? await objective.backwardTargets({
-      model,
-      batch: preparedBatch,
-      config,
-      tape,
-      options,
-      lossScale,
-      loss,
-      lossResult,
-      forwardState,
-    })
-    : createLossGradient(loss, lossScale);
-  const forward_ms = globalThis.performance.now() - t0;
-
-  const t1 = globalThis.performance.now();
-  const backwardSeed = normalizeBackwardSeedPayload(gradOutput, loss, lossScale);
-  const grads = await tape.backward(backwardSeed.seedArg);
-  const backward_ms = globalThis.performance.now() - t1;
-  for (const seedTensor of backwardSeed.seedTensors) {
-    if (seedTensor?.buffer) {
-      releaseBuffer(seedTensor.buffer);
-    }
-  }
-  let processed = grads;
-  if (lossScale !== 1) {
-    const invScale = 1 / lossScale;
-    const unscaled = new Map();
-    for (const [param, grad] of grads.entries()) {
-      const scaled = await runScale(grad, invScale, { inplace: true });
-      unscaled.set(param, scaled);
-    }
-    processed = unscaled;
-  }
-  const paramGroups = resolveModelParamGroups(model);
-  const freezeMap = config.training?.ul?.freeze
-    ?? config.training?.distill?.freeze
-    ?? {};
-  const { trainableGroups, frozenGroups } = selectTrainableParamGroups(paramGroups, freezeMap);
-  const trainableParams = flattenUniqueParams(trainableGroups);
-  processed = filterGradientsForParams(processed, trainableParams);
-
-  let clipMetrics = null;
-  if (applyClip) {
-    clipMetrics = await clipGradients(processed, config);
-    processed = clipMetrics.clippedGrads;
-  }
-
-  let optimizerMetrics = null;
-  let paramGroupMetrics = null;
-  if (applyOptimizer) {
-    optimizerMetrics = await optimizer.step(trainableParams, processed, config, {
-      trainableGroups: Object.keys(trainableGroups),
-      frozenGroups,
-      allGroups: Object.keys(paramGroups),
-    });
-    paramGroupMetrics = {
-      trainableGroups: Object.keys(trainableGroups),
-      frozenGroups,
-      allGroups: Object.keys(paramGroups),
-      trainableParamCount: trainableParams.length,
-    };
-  }
-
-  let objectiveMetrics = lossResult.components || null;
-  if (typeof objective.metrics === 'function') {
-    objectiveMetrics = await objective.metrics({
-      model,
-      batch: preparedBatch,
-      config,
-      tape,
-      options,
-      lossScale,
-      loss,
-      lossResult,
-      forwardState,
-    });
-  }
-
-  if (typeof objective.cleanup === 'function') {
-    await objective.cleanup({
-      model,
-      batch,
-      preparedBatch,
-      config,
-      tape,
-      options,
-      lossScale,
-    });
-  }
-
-  return {
-    loss,
-    grads: processed,
-    forward_ms,
-    backward_ms,
-    clipMetrics,
-    optimizerMetrics,
-    objectiveName: objective.name,
-    objectiveMetrics,
-    paramGroupMetrics,
-  };
-}

package/src/training/ul_dataset.d.ts

@@ -1,47 +0,0 @@
-import type { Tensor } from '../gpu/tensor.js';
-import type { UlTrainingConfigSchema } from '../config/schema/ul-training.schema.js';
-
-export interface UlLatentStats {
-  mean: number;
-  std: number;
-}
-
-export interface UlNoisyLatentBatchResult {
-  noisyTensor: Tensor;
-  cleanStats: UlLatentStats;
-  noiseStats: UlLatentStats;
-  noisyStats: UlLatentStats;
-  alpha: number;
-  sigma: number;
-  lambda0: number;
-  stepIndex: number;
-  cleanValues: Float32Array | null;
-  noiseValues: Float32Array | null;
-  noisyValues: Float32Array | null;
-  shape: number[];
-}
-
-export declare function resolveUlNoiseScale(lambda0: number): { alpha: number; sigma: number };
-export declare function resolveUlScheduledLambda(
-  ulConfig: UlTrainingConfigSchema,
-  stepIndex?: number
-): number;
-
-export declare function buildNoisyLatentsFromInputTensor(
-  inputTensor: Tensor,
-  ulConfig: UlTrainingConfigSchema,
-  options?: { seed?: number; stepIndex?: number; lambda0?: number; includeValues?: boolean }
-): Promise<UlNoisyLatentBatchResult>;
-
-export declare function applyUlStage1Batch(
-  batch: Record<string, unknown> & { input: Tensor },
-  ulConfig: UlTrainingConfigSchema,
-  options?: { seed?: number; stepIndex?: number; lambda0?: number; includeValues?: boolean }
-): Promise<Record<string, unknown>>;
-
-export declare function cleanupUlPreparedBatch(batch: Record<string, unknown>): void;
-
-export declare function computeLatentBitrateProxy(
-  stats: Record<string, unknown> | null | undefined,
-  ulConfig: UlTrainingConfigSchema
-): number;

package/src/training/ul_dataset.js

@@ -1,151 +0,0 @@
-import { readBuffer, releaseBuffer } from '../memory/buffer-pool.js';
-import { resolveUlScheduledLambda } from './ul_schedule.js';
-import { createUploadedTensor } from './tensor-factory.js';
-
-function xorshift32(value) {
-  let x = value >>> 0;
-  x ^= x << 13;
-  x ^= x >>> 17;
-  x ^= x << 5;
-  return x >>> 0;
-}
-
-function uniformFromSeed(seed, index) {
-  const mixed = xorshift32((seed ^ ((index + 1) * 0x9e3779b9)) >>> 0);
-  return (mixed + 1) / 4294967297;
-}
-
-function gaussianFromSeed(seed, index) {
-  const u1 = Math.max(uniformFromSeed(seed, index * 2), 1e-7);
-  const u2 = uniformFromSeed(seed ^ 0x85ebca6b, index * 2 + 1);
-  const radius = Math.sqrt(-2 * Math.log(u1));
-  const theta = 2 * Math.PI * u2;
-  return radius * Math.cos(theta);
-}
-
-export function resolveUlNoiseScale(lambda0) {
-  const snr = Math.exp(lambda0);
-  const alpha = Math.sqrt(snr / (1 + snr));
-  const sigma = Math.sqrt(1 / (1 + snr));
-  return { alpha, sigma };
-}
-
-export { resolveUlScheduledLambda } from './ul_schedule.js';
-
-function summarizeArray(values) {
-  if (!values.length) {
-    return { mean: 0, std: 0 };
-  }
-  let sum = 0;
-  for (let i = 0; i < values.length; i += 1) {
-    sum += values[i];
-  }
-  const mean = sum / values.length;
-  let variance = 0;
-  for (let i = 0; i < values.length; i += 1) {
-    const delta = values[i] - mean;
-    variance += delta * delta;
-  }
-  variance /= values.length;
-  return { mean, std: Math.sqrt(Math.max(variance, 0)) };
-}
-
-export async function buildNoisyLatentsFromInputTensor(inputTensor, ulConfig, options = {}) {
-  if (!inputTensor || !inputTensor.buffer || !Array.isArray(inputTensor.shape)) {
-    throw new Error('UL dataset: input tensor is required.');
-  }
-  if (inputTensor.dtype !== 'f32') {
-    throw new Error(`UL dataset: expected f32 input tensor, got ${inputTensor.dtype}.`);
-  }
-
-  const rawInput = new Float32Array(await readBuffer(inputTensor.buffer));
-  const elementCount = inputTensor.shape.reduce((acc, value) => acc * Math.max(1, Math.floor(Number(value) || 1)), 1);
-  const inputData = rawInput.length === elementCount
-    ? rawInput
-    : rawInput.slice(0, elementCount);
-  const noisy = new Float32Array(inputData.length);
-  const noise = new Float32Array(inputData.length);
-  const stepIndex = Math.max(0, Math.floor(Number(options.stepIndex) || 0));
-  const lambda0 = Number.isFinite(options.lambda0)
-    ? options.lambda0
-    : resolveUlScheduledLambda(ulConfig, stepIndex);
-  const seedBase = Number.isFinite(options.seed) ? options.seed : ulConfig.seed;
-  const seed = (seedBase + (stepIndex * 104729)) >>> 0;
-  const { alpha, sigma } = resolveUlNoiseScale(lambda0);
-
-  for (let i = 0; i < inputData.length; i += 1) {
-    const n = gaussianFromSeed(seed, i);
-    noise[i] = n;
-    noisy[i] = alpha * inputData[i] + sigma * n;
-  }
-
-  const noisyTensor = createUploadedTensor(noisy, 'f32', inputTensor.shape, 'ul_noisy_latents');
-  const cleanStats = summarizeArray(inputData);
-  const noiseStats = summarizeArray(noise);
-  const noisyStats = summarizeArray(noisy);
-
-  return {
-    noisyTensor,
-    cleanStats,
-    noiseStats,
-    noisyStats,
-    alpha,
-    sigma,
-    lambda0,
-    stepIndex,
-    cleanValues: options.includeValues === true ? inputData : null,
-    noiseValues: options.includeValues === true ? noise : null,
-    noisyValues: options.includeValues === true ? noisy : null,
-    shape: [...inputTensor.shape],
-  };
-}
-
-export async function applyUlStage1Batch(batch, ulConfig, options = {}) {
-  const stepIndex = Math.max(0, Math.floor(Number(options.stepIndex) || 0));
-  const lambda = Number.isFinite(options.lambda0)
-    ? options.lambda0
-    : resolveUlScheduledLambda(ulConfig, stepIndex);
-  const prepared = await buildNoisyLatentsFromInputTensor(batch.input, ulConfig, {
-    ...options,
-    lambda0: lambda,
-    stepIndex,
-  });
-  return {
-    ...batch,
-    input: prepared.noisyTensor,
-    ul: {
-      alpha: prepared.alpha,
-      sigma: prepared.sigma,
-      clean: prepared.cleanStats,
-      noise: prepared.noiseStats,
-      noisy: prepared.noisyStats,
-      lambda: prepared.lambda0,
-      schedule: ulConfig.noiseSchedule || null,
-      stepIndex,
-      shape: prepared.shape,
-      values: options.includeValues === true
-        ? {
-          clean: prepared.cleanValues ? Array.from(prepared.cleanValues) : null,
-          noise: prepared.noiseValues ? Array.from(prepared.noiseValues) : null,
-          noisy: prepared.noisyValues ? Array.from(prepared.noisyValues) : null,
-        }
-        : null,
-    },
-    _ulTemporaryTensors: [prepared.noisyTensor],
-  };
-}
-
-export function cleanupUlPreparedBatch(batch) {
-  const tensors = Array.isArray(batch?._ulTemporaryTensors) ? batch._ulTemporaryTensors : [];
-  for (const tensor of tensors) {
-    if (tensor?.buffer) {
-      releaseBuffer(tensor.buffer);
-    }
-  }
-}
-
-export function computeLatentBitrateProxy(stats, ulConfig) {
-  const variance = Math.max((stats?.noisy?.std ?? 0) ** 2, 1e-8);
-  const lambda = Number.isFinite(ulConfig?.lambda0) ? ulConfig.lambda0 : 5;
-  return Math.log2(1 + variance * Math.exp(lambda));
-}

package/src/training/ul_schedule.d.ts

@@ -1,6 +0,0 @@
-import type { UlTrainingConfigSchema } from '../config/schema/ul-training.schema.js';
-
-export declare function resolveUlScheduledLambda(
-  ulConfig: UlTrainingConfigSchema,
-  stepIndex?: number
-): number;

package/src/training/ul_schedule.js

@@ -1,29 +0,0 @@
-function clamp01(value) {
-  if (!Number.isFinite(value)) return 0;
-  if (value <= 0) return 0;
-  if (value >= 1) return 1;
-  return value;
-}
-
-export function resolveUlScheduledLambda(ulConfig, stepIndex = 0) {
-  const fallback = Number.isFinite(ulConfig?.lambda0) ? ulConfig.lambda0 : 5;
-  const schedule = ulConfig?.noiseSchedule;
-  if (!schedule || typeof schedule !== 'object') {
-    return fallback;
-  }
-
-  const steps = Math.max(1, Math.floor(Number(schedule.steps) || 1));
-  const step = Math.max(0, Math.floor(Number(stepIndex) || 0));
-  const t = steps <= 1 ? 1 : clamp01(step / (steps - 1));
-  const minLogSNR = Number.isFinite(schedule.minLogSNR) ? schedule.minLogSNR : fallback;
-  const maxLogSNR = Number.isFinite(schedule.maxLogSNR) ? schedule.maxLogSNR : fallback;
-  const type = typeof schedule.type === 'string' ? schedule.type : 'log_snr_linear';
-
-  if (type === 'log_snr_cosine') {
-    const cosineT = 0.5 - 0.5 * Math.cos(Math.PI * t);
-    return maxLogSNR + ((minLogSNR - maxLogSNR) * cosineT);
-  }
-
-  // Default: linear interpolation from max -> min across schedule steps.
-  return maxLogSNR + ((minLogSNR - maxLogSNR) * t);
-}