@mlx-node/trl 0.0.0

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

@@ -0,0 +1,32 @@
+export declare class SFTConfigError extends Error {
+    constructor(message: string);
+}
+export interface SFTTrainerConfig {
+    model_name: string;
+    output_dir: string;
+    run_name: string;
+    learning_rate: number;
+    batch_size: number;
+    gradient_accumulation_steps: number;
+    num_epochs: number;
+    max_train_samples: number;
+    max_grad_norm: number;
+    weight_decay: number;
+    max_seq_length: number;
+    completion_only: boolean;
+    label_smoothing: number;
+    logging_steps: number;
+    save_steps: number;
+    max_checkpoints: number;
+    log_jsonl: boolean;
+    tui_mode: boolean;
+    seed: number;
+    resume_from_checkpoint: 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":"AAIA,qBAAa,cAAe,SAAQ,KAAK;gBAC3B,OAAO,EAAE,MAAM;CAI5B;AAED,MAAM,WAAW,gBAAgB;IAE/B,UAAU,EAAE,MAAM,CAAC;IACnB,UAAU,EAAE,MAAM,CAAC;IACnB,QAAQ,EAAE,MAAM,CAAC;IAGjB,aAAa,EAAE,MAAM,CAAC;IACtB,UAAU,EAAE,MAAM,CAAC;IACnB,2BAA2B,EAAE,MAAM,CAAC;IACpC,UAAU,EAAE,MAAM,CAAC;IACnB,iBAAiB,EAAE,MAAM,CAAC;IAC1B,aAAa,EAAE,MAAM,CAAC;IACtB,YAAY,EAAE,MAAM,CAAC;IAGrB,cAAc,EAAE,MAAM,CAAC;IACvB,eAAe,EAAE,OAAO,CAAC;IACzB,eAAe,EAAE,MAAM,CAAC;IAGxB,aAAa,EAAE,MAAM,CAAC;IACtB,UAAU,EAAE,MAAM,CAAC;IACnB,eAAe,EAAE,MAAM,CAAC;IACxB,SAAS,EAAE,OAAO,CAAC;IACnB,QAAQ,EAAE,OAAO,CAAC;IAGlB,IAAI,EAAE,MAAM,CAAC;IACb,sBAAsB,EAAE,MAAM,CAAC;CAChC;AAED,QAAA,MAAM,kBAAkB,EAAE,gBAyBxB,CAAC;AA0GH,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,CAkBjG;AAED,OAAO,EAAE,kBAAkB,EAAE,CAAC"}

package/dist/trainers/sft-config.js

@@ -0,0 +1,186 @@
+import { parse as parseToml } from '@std/toml';
+import { readFileSync } from 'node:fs';
+import { resolve as resolvePath } from 'node:path';
+export class SFTConfigError extends Error {
+    constructor(message) {
+        super(message);
+        this.name = 'SFTConfigError';
+    }
+}
+const DEFAULT_SFT_CONFIG = Object.freeze({
+    model_name: 'Qwen/Qwen3-0.6B',
+    output_dir: 'outputs/sft',
+    run_name: 'sft-run',
+    learning_rate: 2e-5,
+    batch_size: 4,
+    gradient_accumulation_steps: 1,
+    num_epochs: 3,
+    max_train_samples: 0,
+    max_grad_norm: 1.0,
+    weight_decay: 0.01,
+    max_seq_length: 2048,
+    completion_only: false, // Changed to false for TRL parity
+    label_smoothing: 0.0,
+    logging_steps: 10,
+    save_steps: 100,
+    max_checkpoints: 3,
+    log_jsonl: true,
+    tui_mode: false,
+    seed: 42,
+    resume_from_checkpoint: '',
+});
+const SFT_CONFIG_VALUE_TYPES = {
+    model_name: 'string',
+    output_dir: 'string',
+    run_name: 'string',
+    learning_rate: 'number',
+    batch_size: 'number',
+    gradient_accumulation_steps: 'number',
+    num_epochs: 'number',
+    max_train_samples: 'number',
+    max_grad_norm: 'number',
+    weight_decay: 'number',
+    max_seq_length: 'number',
+    completion_only: 'boolean',
+    label_smoothing: 'number',
+    logging_steps: 'number',
+    save_steps: 'number',
+    max_checkpoints: 'number',
+    log_jsonl: 'boolean',
+    tui_mode: 'boolean',
+    seed: 'number',
+    resume_from_checkpoint: 'string',
+};
+const SFT_INTEGER_KEYS = new Set([
+    'batch_size',
+    'gradient_accumulation_steps',
+    'num_epochs',
+    'max_train_samples',
+    'max_seq_length',
+    'logging_steps',
+    'save_steps',
+    'max_checkpoints',
+    '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)) {
+        if (!isConfigKey(rawKey)) {
+            continue;
+        }
+        setConfigValue(normalized, rawKey, coerceValue(rawKey, 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 key = entry.slice(0, idx).trim();
+        const rawValue = entry.slice(idx + 1).trim();
+        if (!isConfigKey(key)) {
+            throw new SFTConfigError(`Unknown configuration key in override: ${key}`);
+        }
+        setConfigValue(accumulated, key, coerceValue(key, rawValue));
+    }
+    return mergeSFTConfig(config, accumulated);
+}
+export { DEFAULT_SFT_CONFIG };

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

@@ -0,0 +1,141 @@
+/**
+ * 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 { Qwen3Model, Qwen3Tokenizer, type SftStepMetrics } from '@mlx-node/core';
+import type { SFTTrainerConfig } from './sft-config';
+import { SFTDataset, type SFTBatch } from '../data/sft-dataset';
+import { type TrainingLogger } from './training-logger';
+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;
+    /**
+     * 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: Qwen3Model, tokenizer: Qwen3Tokenizer, config?: Partial<SFTTrainerConfig>, logger?: TrainingLogger);
+    /**
+     * 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(): Qwen3Model;
+    /**
+     * 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,EAEL,UAAU,EACV,cAAc,EAGd,KAAK,cAAc,EAEpB,MAAM,gBAAgB,CAAC;AAExB,OAAO,KAAK,EAAE,gBAAgB,EAAE,MAAM,cAAc,CAAC;AAErD,OAAO,EAAE,UAAU,EAAkB,KAAK,QAAQ,EAAE,MAAM,qBAAqB,CAAC;AAChF,OAAO,EAAwB,KAAK,cAAc,EAAE,MAAM,mBAAmB,CAAC;AAG9E,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,CAAa;IAC1B,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;IAEnE;;;;;;;OAOG;gBAED,KAAK,EAAE,UAAU,EACjB,SAAS,EAAE,cAAc,EACzB,MAAM,GAAE,OAAO,CAAC,gBAAgB,CAAM,EACtC,MAAM,CAAC,EAAE,cAAc;IAwCzB;;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;IA+D3E;;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;IA8KxD;;;;;OAKG;IACG,cAAc,CAAC,IAAI,CAAC,EAAE,MAAM,GAAG,OAAO,CAAC,MAAM,CAAC;IAgDpD;;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,UAAU;IAItB;;OAEG;IACH,YAAY,IAAI,cAAc;CAG/B"}