@mlx-node/trl 0.0.2 → 0.0.4

package/dist/trainers/sft-config.d.ts

@@ -0,0 +1,33 @@
+export declare class SFTConfigError extends Error {
+    constructor(message: string);
+}
+export interface SFTTrainerConfig {
+    modelName: string;
+    outputDir: string;
+    runName: string;
+    learningRate: number;
+    batchSize: number;
+    gradientAccumulationSteps: number;
+    numEpochs: number;
+    maxTrainSamples: number;
+    maxGradNorm: number;
+    weightDecay: number;
+    maxSeqLength: number;
+    completionOnly: boolean;
+    labelSmoothing: number;
+    loggingSteps: number;
+    saveSteps: number;
+    maxCheckpoints: number;
+    logJsonl: boolean;
+    tuiMode: boolean;
+    gradientCheckpointing: boolean;
+    seed: number;
+    resumeFromCheckpoint: string;
+}
+declare const DEFAULT_SFT_CONFIG: SFTTrainerConfig;
+export declare function getDefaultSFTConfig(): SFTTrainerConfig;
+export declare function mergeSFTConfig(base: SFTTrainerConfig, update: Partial<SFTTrainerConfig>): SFTTrainerConfig;
+export declare function loadSFTTomlConfig(filePath: string): SFTTrainerConfig;
+export declare function applySFTOverrides(config: SFTTrainerConfig, overrides: string[]): SFTTrainerConfig;
+export { DEFAULT_SFT_CONFIG };
+//# sourceMappingURL=sft-config.d.ts.map

package/dist/trainers/sft-config.d.ts.map

@@ -0,0 +1 @@
+{"version":3,"file":"sft-config.d.ts","sourceRoot":"","sources":["../../src/trainers/sft-config.ts"],"names":[],"mappings":"AAMA,qBAAa,cAAe,SAAQ,KAAK;gBAC3B,OAAO,EAAE,MAAM;CAI5B;AAED,MAAM,WAAW,gBAAgB;IAE/B,SAAS,EAAE,MAAM,CAAC;IAClB,SAAS,EAAE,MAAM,CAAC;IAClB,OAAO,EAAE,MAAM,CAAC;IAGhB,YAAY,EAAE,MAAM,CAAC;IACrB,SAAS,EAAE,MAAM,CAAC;IAClB,yBAAyB,EAAE,MAAM,CAAC;IAClC,SAAS,EAAE,MAAM,CAAC;IAClB,eAAe,EAAE,MAAM,CAAC;IACxB,WAAW,EAAE,MAAM,CAAC;IACpB,WAAW,EAAE,MAAM,CAAC;IAGpB,YAAY,EAAE,MAAM,CAAC;IACrB,cAAc,EAAE,OAAO,CAAC;IACxB,cAAc,EAAE,MAAM,CAAC;IAGvB,YAAY,EAAE,MAAM,CAAC;IACrB,SAAS,EAAE,MAAM,CAAC;IAClB,cAAc,EAAE,MAAM,CAAC;IACvB,QAAQ,EAAE,OAAO,CAAC;IAClB,OAAO,EAAE,OAAO,CAAC;IAGjB,qBAAqB,EAAE,OAAO,CAAC;IAG/B,IAAI,EAAE,MAAM,CAAC;IACb,oBAAoB,EAAE,MAAM,CAAC;CAC9B;AAED,QAAA,MAAM,kBAAkB,EAAE,gBA2BxB,CAAC;AA6GH,wBAAgB,mBAAmB,IAAI,gBAAgB,CAEtD;AAED,wBAAgB,cAAc,CAAC,IAAI,EAAE,gBAAgB,EAAE,MAAM,EAAE,OAAO,CAAC,gBAAgB,CAAC,GAAG,gBAAgB,CAa1G;AAED,wBAAgB,iBAAiB,CAAC,QAAQ,EAAE,MAAM,GAAG,gBAAgB,CAsBpE;AAED,wBAAgB,iBAAiB,CAAC,MAAM,EAAE,gBAAgB,EAAE,SAAS,EAAE,MAAM,EAAE,GAAG,gBAAgB,CAoBjG;AAED,OAAO,EAAE,kBAAkB,EAAE,CAAC"}

package/dist/trainers/sft-config.js

@@ -0,0 +1,193 @@
+import { readFileSync } from 'node:fs';
+import { resolve as resolvePath } from 'node:path';
+import { parse as parseToml } from '@std/toml';
+import { camelCase } from 'change-case';
+export class SFTConfigError extends Error {
+    constructor(message) {
+        super(message);
+        this.name = 'SFTConfigError';
+    }
+}
+const DEFAULT_SFT_CONFIG = Object.freeze({
+    modelName: 'Qwen/Qwen3-0.6B',
+    outputDir: 'outputs/sft',
+    runName: 'sft-run',
+    learningRate: 2e-5,
+    batchSize: 4,
+    gradientAccumulationSteps: 1,
+    numEpochs: 3,
+    maxTrainSamples: 0,
+    maxGradNorm: 1.0,
+    weightDecay: 0.01,
+    maxSeqLength: 2048,
+    completionOnly: false, // Changed to false for TRL parity
+    labelSmoothing: 0.0,
+    loggingSteps: 10,
+    saveSteps: 100,
+    maxCheckpoints: 3,
+    logJsonl: true,
+    tuiMode: false,
+    gradientCheckpointing: true,
+    seed: 42,
+    resumeFromCheckpoint: '',
+});
+const SFT_CONFIG_VALUE_TYPES = {
+    modelName: 'string',
+    outputDir: 'string',
+    runName: 'string',
+    learningRate: 'number',
+    batchSize: 'number',
+    gradientAccumulationSteps: 'number',
+    numEpochs: 'number',
+    maxTrainSamples: 'number',
+    maxGradNorm: 'number',
+    weightDecay: 'number',
+    maxSeqLength: 'number',
+    completionOnly: 'boolean',
+    labelSmoothing: 'number',
+    loggingSteps: 'number',
+    saveSteps: 'number',
+    maxCheckpoints: 'number',
+    logJsonl: 'boolean',
+    tuiMode: 'boolean',
+    seed: 'number',
+    gradientCheckpointing: 'boolean',
+    resumeFromCheckpoint: 'string',
+};
+const SFT_INTEGER_KEYS = new Set([
+    'batchSize',
+    'gradientAccumulationSteps',
+    'numEpochs',
+    'maxTrainSamples',
+    'maxSeqLength',
+    'loggingSteps',
+    'saveSteps',
+    'maxCheckpoints',
+    'seed',
+]);
+function cloneDefaults() {
+    return { ...DEFAULT_SFT_CONFIG };
+}
+function isConfigKey(value) {
+    return Object.prototype.hasOwnProperty.call(SFT_CONFIG_VALUE_TYPES, value);
+}
+function coerceBoolean(value, key) {
+    if (typeof value === 'boolean')
+        return value;
+    if (typeof value === 'string') {
+        const normalized = value.trim().toLowerCase();
+        if (['true', '1', 'yes', 'on'].includes(normalized))
+            return true;
+        if (['false', '0', 'no', 'off'].includes(normalized))
+            return false;
+    }
+    throw new SFTConfigError(`Invalid boolean for ${key}: ${String(value)}`);
+}
+function coerceNumber(value, key) {
+    if (typeof value === 'string' && value.trim() === '') {
+        throw new SFTConfigError(`Invalid number for ${key}: empty string`);
+    }
+    const parsed = typeof value === 'number' ? value : Number(value);
+    if (!Number.isFinite(parsed)) {
+        throw new SFTConfigError(`Invalid number for ${key}: ${String(value)}`);
+    }
+    if (SFT_INTEGER_KEYS.has(key) && !Number.isInteger(parsed)) {
+        throw new SFTConfigError(`Expected integer for ${key}, received ${parsed}`);
+    }
+    return parsed;
+}
+function coerceString(value, key) {
+    if (typeof value === 'string')
+        return value;
+    throw new SFTConfigError(`Invalid string for ${key}: ${String(value)}`);
+}
+function coerceValue(key, value) {
+    const expected = SFT_CONFIG_VALUE_TYPES[key];
+    if (expected === 'boolean') {
+        return coerceBoolean(value, key);
+    }
+    if (expected === 'number') {
+        return coerceNumber(value, key);
+    }
+    return coerceString(value, key);
+}
+function setConfigValue(config, key, value) {
+    config[key] = value;
+}
+function normalizeTomlRecord(record) {
+    const normalized = {};
+    for (const [rawKey, rawValue] of Object.entries(record)) {
+        // TOML uses snake_case, config uses camelCase
+        const key = camelCase(rawKey);
+        if (!isConfigKey(key)) {
+            continue;
+        }
+        setConfigValue(normalized, key, coerceValue(key, rawValue));
+    }
+    return normalized;
+}
+export function getDefaultSFTConfig() {
+    return cloneDefaults();
+}
+export function mergeSFTConfig(base, update) {
+    if (!update) {
+        return { ...base };
+    }
+    const result = { ...base };
+    for (const [key, value] of Object.entries(update)) {
+        if (value === undefined)
+            continue;
+        if (!isConfigKey(key)) {
+            throw new SFTConfigError(`Unknown configuration key: ${key}`);
+        }
+        setConfigValue(result, key, value);
+    }
+    return result;
+}
+export function loadSFTTomlConfig(filePath) {
+    const absolutePath = resolvePath(filePath);
+    let fileContents;
+    try {
+        fileContents = readFileSync(absolutePath, 'utf8');
+    }
+    catch (error) {
+        const message = error instanceof Error ? error.message : String(error);
+        throw new SFTConfigError(`Failed to read config at ${absolutePath}: ${message}`);
+    }
+    let parsedRaw;
+    try {
+        parsedRaw = parseToml(fileContents);
+    }
+    catch (error) {
+        const message = error instanceof Error ? error.message : String(error);
+        throw new SFTConfigError(`Failed to parse TOML at ${absolutePath}: ${message}`);
+    }
+    if (parsedRaw === null || typeof parsedRaw !== 'object' || Array.isArray(parsedRaw)) {
+        throw new SFTConfigError(`Expected table at ${absolutePath}`);
+    }
+    const parsed = parsedRaw;
+    const normalized = normalizeTomlRecord(parsed);
+    return mergeSFTConfig(getDefaultSFTConfig(), normalized);
+}
+export function applySFTOverrides(config, overrides) {
+    if (!overrides.length) {
+        return { ...config };
+    }
+    const accumulated = {};
+    for (const entry of overrides) {
+        const idx = entry.indexOf('=');
+        if (idx === -1) {
+            throw new SFTConfigError(`Invalid override "${entry}", expected key=value format`);
+        }
+        const rawKey = entry.slice(0, idx).trim();
+        const rawValue = entry.slice(idx + 1).trim();
+        // Accept both snake_case and camelCase overrides
+        const key = camelCase(rawKey);
+        if (!isConfigKey(key)) {
+            throw new SFTConfigError(`Unknown configuration key in override: ${rawKey}`);
+        }
+        setConfigValue(accumulated, key, coerceValue(key, rawValue));
+    }
+    return mergeSFTConfig(config, accumulated);
+}
+export { DEFAULT_SFT_CONFIG };

package/dist/trainers/sft-trainer.d.ts

@@ -0,0 +1,148 @@
+/**
+ * SFT (Supervised Fine-Tuning) Trainer
+ *
+ * This module provides a Rust-native SFT training engine for training
+ * models on fixed prompt-completion pairs using cross-entropy loss.
+ *
+ * ## Key Features
+ * - Training loop runs in Rust (eliminates FFI overhead)
+ * - Cross-entropy loss with completion masking (ignore_index=-100)
+ * - Label smoothing support
+ * - Gradient accumulation and clipping
+ * - High-level train() method for full training runs
+ * - Low-level trainStep() for custom training loops
+ *
+ * ## Usage
+ * ```typescript
+ * const trainer = await SFTTrainer.create({
+ *   modelPath: './model',
+ *   learningRate: 2e-5,
+ *   numEpochs: 3,
+ * });
+ * await trainer.train(dataset);
+ * ```
+ */
+import { Qwen3Tokenizer, type SftStepMetrics } from '@mlx-node/core';
+import { type TrainableModel } from '@mlx-node/lm';
+import { SFTDataset, type SFTBatch } from '../data/sft-dataset.js';
+import type { SFTTrainerConfig } from './sft-config.js';
+import { type TrainingLogger } from './training-logger.js';
+export { SftTrainingEngine } from '@mlx-node/core';
+export type { SftEngineConfig, SftStepMetrics, SftEpochMetrics } from '@mlx-node/core';
+/**
+ * Training state saved with checkpoints for resumption
+ */
+export interface SFTTrainingState {
+    step: number;
+    epoch: number;
+    timestamp: string;
+    trainerType: 'sft';
+}
+/**
+ * Training step result
+ */
+export interface SFTTrainStepResult {
+    /** Step metrics */
+    metrics: SftStepMetrics;
+    /** Current epoch */
+    epoch: number;
+}
+/**
+ * SFT Trainer - Rust-Native Training Engine
+ *
+ * Provides a TypeScript-friendly interface to the Rust SFT training engine.
+ */
+export declare class SFTTrainer {
+    private engine;
+    private model;
+    private tokenizer;
+    private config;
+    private currentEpoch;
+    private currentStep;
+    /** Original model path (for tokenizer files when saving checkpoints) */
+    private originalModelPath?;
+    private paused;
+    private stopRequested;
+    private stdinInterface?;
+    private logger;
+    private sampleDisplayMode;
+    private signalHandlersInstalled;
+    /**
+     * Create a new SFT trainer from a model
+     *
+     * @param model - Pre-loaded Qwen3 model
+     * @param tokenizer - Pre-loaded tokenizer
+     * @param config - Training configuration
+     * @param logger - Optional custom logger
+     */
+    constructor(model: TrainableModel, tokenizer: Qwen3Tokenizer, config?: Partial<SFTTrainerConfig>, logger?: TrainingLogger);
+    /**
+     * Setup OS signal handlers for graceful shutdown on interrupt.
+     * Saves an emergency checkpoint before exiting to prevent progress loss.
+     */
+    private setupSignalHandlers;
+    /**
+     * Setup stdin handler for TUI control commands
+     */
+    private setupStdinHandler;
+    /**
+     * Handle a command received from stdin
+     */
+    private handleStdinCommand;
+    /**
+     * Wait for resume if paused
+     */
+    private waitForResume;
+    /**
+     * Create a trainer by loading a model from disk
+     *
+     * @param config - Configuration including modelPath
+     * @returns Promise<SFTTrainer>
+     */
+    static create(config: Partial<SFTTrainerConfig>): Promise<SFTTrainer>;
+    /**
+     * Find the latest checkpoint in the output directory
+     */
+    static findLatestCheckpoint(outputDir?: string): string | null;
+    /**
+     * Run a single training step
+     *
+     * @param batch - Tokenized batch with input_ids and labels
+     * @returns Training step metrics
+     */
+    trainStep(batch: SFTBatch): Promise<SFTTrainStepResult>;
+    /**
+     * Run a full training loop over a dataset
+     *
+     * @param dataset - SFT dataset or path to JSONL file
+     */
+    train(dataset: SFTDataset | string): Promise<void>;
+    /**
+     * Save a checkpoint with model weights and training state
+     *
+     * @param name - Checkpoint name (default: "checkpoint-{step}")
+     * @returns Path to saved checkpoint
+     */
+    saveCheckpoint(name?: string): Promise<string>;
+    /**
+     * Remove old checkpoints, keeping only the most recent ones
+     */
+    private cleanupOldCheckpoints;
+    /**
+     * Get current training step
+     */
+    get step(): number;
+    /**
+     * Get current epoch
+     */
+    get epoch(): number;
+    /**
+     * Get the underlying model for inference
+     */
+    getModel(): TrainableModel;
+    /**
+     * Get the tokenizer
+     */
+    getTokenizer(): Qwen3Tokenizer;
+}
+//# sourceMappingURL=sft-trainer.d.ts.map

package/dist/trainers/sft-trainer.d.ts.map

@@ -0,0 +1 @@
+{"version":3,"file":"sft-trainer.d.ts","sourceRoot":"","sources":["../../src/trainers/sft-trainer.ts"],"names":[],"mappings":"AAAA;;;;;;;;;;;;;;;;;;;;;;;GAuBG;AAMH,OAAO,EAKL,cAAc,EAGd,KAAK,cAAc,EAEpB,MAAM,gBAAgB,CAAC;AACxB,OAAO,EAAa,KAAK,cAAc,EAAE,MAAM,cAAc,CAAC;AAE9D,OAAO,EAAE,UAAU,EAAkB,KAAK,QAAQ,EAAE,MAAM,wBAAwB,CAAC;AACnF,OAAO,KAAK,EAAE,gBAAgB,EAAE,MAAM,iBAAiB,CAAC;AAExD,OAAO,EAAwB,KAAK,cAAc,EAAE,MAAM,sBAAsB,CAAC;AAGjF,OAAO,EAAE,iBAAiB,EAAE,MAAM,gBAAgB,CAAC;AACnD,YAAY,EAAE,eAAe,EAAE,cAAc,EAAE,eAAe,EAAE,MAAM,gBAAgB,CAAC;AAEvF;;GAEG;AACH,MAAM,WAAW,gBAAgB;IAC/B,IAAI,EAAE,MAAM,CAAC;IACb,KAAK,EAAE,MAAM,CAAC;IACd,SAAS,EAAE,MAAM,CAAC;IAClB,WAAW,EAAE,KAAK,CAAC;CACpB;AAED;;GAEG;AACH,MAAM,WAAW,kBAAkB;IACjC,mBAAmB;IACnB,OAAO,EAAE,cAAc,CAAC;IACxB,oBAAoB;IACpB,KAAK,EAAE,MAAM,CAAC;CACf;AAED;;;;GAIG;AACH,qBAAa,UAAU;IACrB,OAAO,CAAC,MAAM,CAAoB;IAClC,OAAO,CAAC,KAAK,CAAiB;IAC9B,OAAO,CAAC,SAAS,CAAiB;IAClC,OAAO,CAAC,MAAM,CAAmB;IACjC,OAAO,CAAC,YAAY,CAAa;IACjC,OAAO,CAAC,WAAW,CAAa;IAChC,wEAAwE;IACxE,OAAO,CAAC,iBAAiB,CAAC,CAAS;IAGnC,OAAO,CAAC,MAAM,CAAkB;IAChC,OAAO,CAAC,aAAa,CAAkB;IACvC,OAAO,CAAC,cAAc,CAAC,CAAqB;IAC5C,OAAO,CAAC,MAAM,CAAiB;IAC/B,OAAO,CAAC,iBAAiB,CAA0C;IACnE,OAAO,CAAC,uBAAuB,CAAkB;IAEjD;;;;;;;OAOG;gBAED,KAAK,EAAE,cAAc,EACrB,SAAS,EAAE,cAAc,EACzB,MAAM,GAAE,OAAO,CAAC,gBAAgB,CAAM,EACtC,MAAM,CAAC,EAAE,cAAc;IAiDzB;;;OAGG;IACH,OAAO,CAAC,mBAAmB;IA8B3B;;OAEG;IACH,OAAO,CAAC,iBAAiB;IAezB;;OAEG;IACH,OAAO,CAAC,kBAAkB;IAmC1B;;OAEG;YACW,aAAa;IAM3B;;;;;OAKG;WACU,MAAM,CAAC,MAAM,EAAE,OAAO,CAAC,gBAAgB,CAAC,GAAG,OAAO,CAAC,UAAU,CAAC;IAgE3E;;OAEG;IACH,MAAM,CAAC,oBAAoB,CAAC,SAAS,CAAC,EAAE,MAAM,GAAG,MAAM,GAAG,IAAI;IAmB9D;;;;;OAKG;IACG,SAAS,CAAC,KAAK,EAAE,QAAQ,GAAG,OAAO,CAAC,kBAAkB,CAAC;IAqB7D;;;;OAIG;IACG,KAAK,CAAC,OAAO,EAAE,UAAU,GAAG,MAAM,GAAG,OAAO,CAAC,IAAI,CAAC;IAiLxD;;;;;OAKG;IACG,cAAc,CAAC,IAAI,CAAC,EAAE,MAAM,GAAG,OAAO,CAAC,MAAM,CAAC;IA+CpD;;OAEG;IACH,OAAO,CAAC,qBAAqB;IAiC7B;;OAEG;IACH,IAAI,IAAI,IAAI,MAAM,CAEjB;IAED;;OAEG;IACH,IAAI,KAAK,IAAI,MAAM,CAElB;IAED;;OAEG;IACH,QAAQ,IAAI,cAAc;IAU1B;;OAEG;IACH,YAAY,IAAI,cAAc;CAG/B"}