@mlx-node/trl 0.0.0 → 0.0.2

package/README.md

@@ -0,0 +1,389 @@
+# @mlx-node/trl
+
+Training library for language models on Apple Silicon. Supports GRPO (Group Relative Policy Optimization) and SFT (Supervised Fine-Tuning) with Metal GPU acceleration, built-in reward functions, dataset handling, and checkpoint management.
+
+## Requirements
+
+- macOS with Apple Silicon (M1 or later)
+- Node.js 18+
+
+## Installation
+
+```bash
+npm install @mlx-node/trl
+```
+
+## Quick Start
+
+### GRPO Training
+
+```typescript
+import { GRPOTrainer } from '@mlx-node/trl';
+
+const trainer = await GRPOTrainer.create({
+  modelPath: './models/Qwen3-0.6B',
+  outputDir: './output/grpo-run',
+  learningRate: 1e-6,
+  groupSize: 4,
+  maxCompletionLength: 256,
+  temperature: 0.8,
+  rewardFunction: async (outputs) => {
+    return outputs.map((o) => (o.text.includes('correct') ? 1.0 : 0.0));
+  },
+});
+
+const dataset = await loadDataset('train');
+await trainer.train(dataset);
+```
+
+### SFT Training
+
+```typescript
+import { SFTTrainer } from '@mlx-node/trl';
+
+const trainer = await SFTTrainer.create({
+  modelName: './models/Qwen3-0.6B',
+  outputDir: './output/sft-run',
+  learningRate: 2e-5,
+  batchSize: 4,
+  numEpochs: 3,
+  completionOnly: true,
+});
+
+await trainer.train('./data/training.jsonl');
+```
+
+## GRPO Training
+
+GRPO generates multiple completions per prompt, scores them with reward functions, and trains the model to prefer higher-reward outputs.
+
+### Loss Variants
+
+| Loss Type | Description                                 |
+| --------- | ------------------------------------------- |
+| `grpo`    | Standard Group Relative Policy Optimization |
+| `dapo`    | Dynamic sampling with adaptive clipping     |
+| `dr_grpo` | Dr.GRPO with improved gradient estimation   |
+| `bnpo`    | Batch-normalized policy optimization        |
+
+### Configuration
+
+```typescript
+import { GRPOTrainer, GRPOTrainerConfig } from '@mlx-node/trl';
+
+const config: GRPOTrainerConfig = {
+  // Model
+  modelPath: './models/Qwen3-0.6B',
+  outputDir: './output',
+
+  // Training
+  learningRate: 1e-6,
+  batchSize: 1,
+  numEpochs: 1,
+  gradientAccumulationSteps: 1,
+  gradientClipNorm: 1.0,
+  weightDecay: 0.01,
+
+  // GRPO
+  groupSize: 4, // completions per prompt
+  clipEpsilon: 0.2, // PPO clipping
+  klCoef: 0.0, // KL divergence coefficient
+  lossType: 'grpo', // grpo | dapo | dr_grpo | bnpo
+
+  // Generation
+  maxCompletionLength: 256,
+  temperature: 0.8,
+  topP: 0.95,
+  repetitionPenalty: 1.1,
+
+  // Tool calling
+  tools: [toolDef],
+  enableThinking: true,
+
+  // Rewards
+  rewardFunction: myRewardFn,
+
+  // Memory optimization
+  gradientCheckpointing: true,
+  lmHeadChunkSize: 2,
+  vocabChunkSize: 65536,
+
+  // Checkpointing
+  saveInterval: 100,
+  maxCheckpoints: 3,
+  resumeFromCheckpoint: './output/checkpoint-500',
+
+  // Optimizer
+  optimizerType: 'adamw', // adamw | sgd
+};
+```
+
+### TOML Configuration
+
+Load training config from a TOML file:
+
+```typescript
+import { loadTomlConfig, applyOverrides } from '@mlx-node/trl';
+
+const config = loadTomlConfig('./train.toml');
+applyOverrides(config, ['learningRate=2e-6', 'batchSize=2']);
+```
+
+### Built-in Rewards
+
+Register native Rust reward functions for high-performance scoring:
+
+```typescript
+trainer.registerBuiltinReward({
+  type: 'ToolUse',
+  weight: 1.0,
+  allowedTools: ['get_weather', 'search'],
+});
+
+trainer.registerBuiltinReward({
+  type: 'XmlFormat',
+  weight: 0.5,
+  requiredTags: ['reasoning', 'answer'],
+});
+
+trainer.registerBuiltinReward({
+  type: 'Length',
+  weight: 0.3,
+  min: 50,
+  max: 500,
+});
+
+trainer.registerBuiltinReward({
+  type: 'JsonSchema',
+  weight: 1.0,
+});
+```
+
+### Custom Reward Functions
+
+```typescript
+import { RewardFunction, RewardOutput } from '@mlx-node/trl';
+
+const reward: RewardFunction = async (outputs: RewardOutput[]) => {
+  return outputs.map((output) => {
+    let score = 0;
+    if (output.toolCalls?.length) score += 0.5;
+    if (output.text.length > 100) score += 0.3;
+    return score;
+  });
+};
+
+trainer.setRewardFunction(reward);
+```
+
+### Custom Training Loop
+
+For advanced use cases, use the low-level API:
+
+```typescript
+const trainer = await GRPOTrainer.create(config);
+
+for (const batch of dataset) {
+  const generations = await trainer.generateBatch(batch.prompts);
+  const rewards = await trainer.scoreGenerations(batch.prompts, generations.completions, context);
+  const metrics = trainer.trainStep(batch.prompts, context);
+  trainer.incrementStep();
+
+  if (metrics.step % 100 === 0) {
+    await trainer.saveCheckpoint();
+  }
+}
+```
+
+### Output Store (SQLite)
+
+Record all training generations and metrics to SQLite for analysis:
+
+```typescript
+const trainer = await GRPOTrainer.create({
+  ...config,
+  outputStore: {
+    enabled: true,
+    database: './output/training.db',
+  },
+});
+```
+
+## SFT Training
+
+Supervised fine-tuning with autograd, gradient accumulation, and completion-only masking.
+
+### Dataset Formats
+
+Two formats are auto-detected from JSONL files:
+
+**Prompt-Completion:**
+
+```json
+{ "prompt": [{ "role": "user", "content": "Hello" }], "completion": { "role": "assistant", "content": "Hi!" } }
+```
+
+**Conversation:**
+
+```json
+{
+  "messages": [
+    { "role": "user", "content": "Hello" },
+    { "role": "assistant", "content": "Hi!" }
+  ]
+}
+```
+
+### SFT Configuration
+
+```typescript
+import { SFTTrainer, SFTTrainerConfig } from '@mlx-node/trl';
+
+const config: SFTTrainerConfig = {
+  modelName: './models/Qwen3-0.6B',
+  outputDir: './output/sft',
+  learningRate: 2e-5,
+  batchSize: 4,
+  gradientAccumulationSteps: 8,
+  numEpochs: 3,
+  maxSeqLength: 2048,
+  completionOnly: true, // only compute loss on assistant tokens
+  labelSmoothing: 0.1,
+  maxGradNorm: 1.0,
+  weightDecay: 0.01,
+  loggingSteps: 10,
+  saveSteps: 100,
+  maxCheckpoints: 3,
+  gradientCheckpointing: true,
+};
+```
+
+### Programmatic Dataset
+
+```typescript
+import { SFTDataset, createSFTDataset } from '@mlx-node/trl';
+
+const dataset = createSFTDataset(examples, tokenizer, {
+  maxSeqLength: 2048,
+  completionOnly: true,
+});
+
+const trainer = await SFTTrainer.create(config);
+await trainer.train(dataset);
+```
+
+## Datasets
+
+### GSM8K Loader
+
+Built-in loader for the GSM8K math dataset:
+
+```typescript
+import { loadLocalGsm8kDataset, LocalGsm8kDatasetLoader } from '@mlx-node/trl';
+
+// Direct load
+const examples = await loadLocalGsm8kDataset('train', { limit: 1000 });
+
+// Via DatasetLoader interface
+const loader = new LocalGsm8kDatasetLoader('./data/gsm8k');
+const trainData = await loader.load('train');
+```
+
+### Custom Datasets
+
+Implement the `DatasetLoader` interface:
+
+```typescript
+import { DatasetLoader, DatasetExample } from '@mlx-node/trl';
+
+class MyDataset implements DatasetLoader {
+  async load(split: 'train' | 'test', limit?: number): Promise<DatasetExample[]> {
+    return examples.map((e) => ({
+      prompt: [
+        { role: 'system', content: 'You are helpful.' },
+        { role: 'user', content: e.question },
+      ],
+      metadata: { answer: e.answer },
+    }));
+  }
+}
+```
+
+## Utilities
+
+### XML Chain-of-Thought Parser
+
+Parse `<reasoning>...</reasoning><answer>...</answer>` format:
+
+```typescript
+import { parseXmlCot, extractXmlAnswer } from '@mlx-node/trl';
+
+const result = parseXmlCot(modelOutput);
+// { reasoning: "...", answer: "42", isStrictMatch: true, isSoftMatch: true, errors: [] }
+
+const answer = extractXmlAnswer(modelOutput);
+// "42"
+```
+
+### Model Conversion
+
+Re-exported from `@mlx-node/core`:
+
+```typescript
+import { convertModel, convertParquetToJsonl } from '@mlx-node/trl';
+```
+
+## Features
+
+- **Checkpoint resume** — automatic state restoration including optimizer, step count, and dataset position
+- **Emergency save** — catches NaN gradients and SIGTERM/SIGINT for safe recovery
+- **TUI mode** — interactive terminal UI with pause/resume/stop (via `mlx-tui` binary)
+- **JSONL logging** — structured training logs for external monitoring
+- **Multi-model** — supports Qwen3, Qwen3.5 Dense, and Qwen3.5 MoE architectures
+- **Reward timeout** — configurable timeout for async reward functions (default 60s)
+- **Path security** — traversal prevention for dataset file loading
+
+## API Reference
+
+### Trainers
+
+| Class         | Description                                                     |
+| ------------- | --------------------------------------------------------------- |
+| `GRPOTrainer` | GRPO training with generation, rewards, and policy optimization |
+| `SFTTrainer`  | Supervised fine-tuning with completion-only masking             |
+
+### Datasets
+
+| Export                    | Description                                      |
+| ------------------------- | ------------------------------------------------ |
+| `loadLocalGsm8kDataset()` | Load GSM8K JSONL dataset                         |
+| `LocalGsm8kDatasetLoader` | `DatasetLoader` implementation for GSM8K         |
+| `SFTDataset`              | Tokenized SFT dataset with padding and shuffling |
+| `loadSFTDataset()`        | Load SFT dataset from JSONL file                 |
+| `createSFTDataset()`      | Create SFT dataset from in-memory examples       |
+
+### Configuration
+
+| Export                  | Description                       |
+| ----------------------- | --------------------------------- |
+| `GRPOTrainerConfig`     | Full GRPO configuration interface |
+| `SFTTrainerConfig`      | Full SFT configuration interface  |
+| `loadTomlConfig()`      | Load GRPO config from TOML file   |
+| `loadSFTTomlConfig()`   | Load SFT config from TOML file    |
+| `getDefaultConfig()`    | Default GRPO config               |
+| `getDefaultSFTConfig()` | Default SFT config                |
+
+### Types
+
+| Type                  | Description                                        |
+| --------------------- | -------------------------------------------------- |
+| `DatasetExample`      | Training example with prompt messages and metadata |
+| `RewardFunction<T>`   | Custom reward function signature                   |
+| `RewardOutput`        | Structured completion data for reward scoring      |
+| `XmlParseResult`      | Result of XML chain-of-thought parsing             |
+| `TrainStepMetrics`    | Per-step training metrics                          |
+| `BuiltinRewardConfig` | Configuration for native reward functions          |
+
+## License
+
+[MIT](https://github.com/mlx-node/mlx-node/blob/main/LICENSE)

package/package.json

@@ -1,6 +1,16 @@
 {
   "name": "@mlx-node/trl",
-  "version": "0.0.0",
+  "version": "0.0.2",
+  "homepage": "https://github.com/mlx-node/mlx-node",
+  "bugs": {
+    "url": "https://github.com/mlx-node/mlx-node/issues"
+  },
+  "license": "MIT",
+  "repository": {
+    "type": "git",
+    "url": "https://github.com/mlx-node/mlx-node.git",
+    "directory": "packages/trl"
+  },
   "files": [
     "dist"
   ],
@@ -19,11 +29,12 @@
     "test:trainer": "TEST_TRAINER=1 vite test run"
   },
   "dependencies": {
-    "@mlx-node/core": "0.0.0",
-    "@mlx-node/lm": "0.0.0",
-    "@std/toml": "npm:@jsr/std__toml@^1.0.11"
+    "@mlx-node/core": "0.0.2",
+    "@mlx-node/lm": "0.0.2",
+    "@std/toml": "npm:@jsr/std__toml@^1.0.11",
+    "change-case": "^5.4.4"
   },
   "devDependencies": {
-    "@huggingface/hub": "^2.7.1"
+    "@huggingface/hub": "^2.10.7"
   }
 }

package/dist/data/dataset.d.ts

@@ -1,22 +0,0 @@
-import type { DatasetExample, ChatMessage, DatasetSplit, PromptFormatterOptions, PromptTemplate, DatasetLoader } from '../types';
-import { type PathValidationOptions } from '../utils/path-security';
-export interface LocalDatasetOptions extends PromptFormatterOptions, PathValidationOptions {
-    basePath?: string;
-    promptTemplate?: PromptTemplate;
-    metadata?: Record<string, unknown>;
-}
-export declare const SYSTEM_PROMPT: string;
-export declare const XML_COT_FORMAT = "<reasoning>\n{reasoning}\n</reasoning>\n<answer>\n{answer}\n</answer>";
-export declare const defaultPromptTemplate: PromptTemplate;
-export declare function createDatasetExample(prompt: ChatMessage[], metadata?: Record<string, unknown>): DatasetExample;
-export declare function extractGsm8kAnswer(raw: string): string | null;
-export declare function validateDatasetExample(example: DatasetExample): void;
-export declare function loadLocalGsm8kDataset(split: DatasetSplit, options?: LocalDatasetOptions & {
-    limit?: number;
-}): Promise<DatasetExample[]>;
-export declare class LocalGsm8kDatasetLoader implements DatasetLoader {
-    private readonly options;
-    constructor(options?: LocalDatasetOptions);
-    load(split: DatasetSplit, limit?: number): Promise<DatasetExample[]>;
-}
-//# sourceMappingURL=dataset.d.ts.map

package/dist/data/dataset.d.ts.map

@@ -1 +0,0 @@
-{"version":3,"file":"dataset.d.ts","sourceRoot":"","sources":["../../src/data/dataset.ts"],"names":[],"mappings":"AAEA,OAAO,KAAK,EACV,cAAc,EACd,WAAW,EAEX,YAAY,EACZ,sBAAsB,EACtB,cAAc,EACd,aAAa,EACd,MAAM,UAAU,CAAC;AAElB,OAAO,EAA2C,KAAK,qBAAqB,EAAE,MAAM,wBAAwB,CAAC;AAE7G,MAAM,WAAW,mBAAoB,SAAQ,sBAAsB,EAAE,qBAAqB;IACxF,QAAQ,CAAC,EAAE,MAAM,CAAC;IAClB,cAAc,CAAC,EAAE,cAAc,CAAC;IAChC,QAAQ,CAAC,EAAE,MAAM,CAAC,MAAM,EAAE,OAAO,CAAC,CAAC;CACpC;AAUD,eAAO,MAAM,aAAa,QASlB,CAAC;AAET,eAAO,MAAM,cAAc,0EAKjB,CAAC;AAWX,eAAO,MAAM,qBAAqB,EAAE,cAYnC,CAAC;AAEF,wBAAgB,oBAAoB,CAAC,MAAM,EAAE,WAAW,EAAE,EAAE,QAAQ,CAAC,EAAE,MAAM,CAAC,MAAM,EAAE,OAAO,CAAC,GAAG,cAAc,CAK9G;AAED,wBAAgB,kBAAkB,CAAC,GAAG,EAAE,MAAM,GAAG,MAAM,GAAG,IAAI,CAE7D;AAED,wBAAgB,sBAAsB,CAAC,OAAO,EAAE,cAAc,GAAG,IAAI,CAYpE;AAwDD,wBAAsB,qBAAqB,CACzC,KAAK,EAAE,YAAY,EACnB,OAAO,GAAE,mBAAmB,GAAG;IAAE,KAAK,CAAC,EAAE,MAAM,CAAA;CAAO,GACrD,OAAO,CAAC,cAAc,EAAE,CAAC,CA4B3B;AAED,qBAAa,uBAAwB,YAAW,aAAa;IAC3D,OAAO,CAAC,QAAQ,CAAC,OAAO,CAAsB;gBAElC,OAAO,GAAE,mBAAwB;IAIvC,IAAI,CAAC,KAAK,EAAE,YAAY,EAAE,KAAK,CAAC,EAAE,MAAM,GAAG,OAAO,CAAC,cAAc,EAAE,CAAC;CAG3E"}

package/dist/data/dataset.js

@@ -1,142 +0,0 @@
-import { readFileSync } from 'node:fs';
-import { resolve as resolvePath } from 'node:path';
-import { extractHashAnswer } from '../utils/xml-parser';
-import { validatePathContainment, getAllowedRoot } from '../utils/path-security';
-const DEFAULT_BASE_PATH = resolvePath(process.cwd(), 'data/gsm8k');
-const VALID_SPLITS = new Set(['train', 'test']);
-export const SYSTEM_PROMPT = `
-Respond in the following format:
-
-<reasoning>
-...
-</reasoning>
-<answer>
-...
-</answer>
-`.trim();
-export const XML_COT_FORMAT = `<reasoning>
-{reasoning}
-</reasoning>
-<answer>
-{answer}
-</answer>`;
-const SYSTEM_MESSAGE = {
-    role: 'system',
-    content: SYSTEM_PROMPT,
-};
-function createMessage(role, content) {
-    return { role, content };
-}
-export const defaultPromptTemplate = (question, options) => {
-    const messages = [SYSTEM_MESSAGE];
-    if (options?.includeOneShot && options.oneShotExample) {
-        const { question: exampleQuestion, reasoning, answer } = options.oneShotExample;
-        messages.push(createMessage('user', exampleQuestion), createMessage('assistant', XML_COT_FORMAT.replace('{reasoning}', reasoning).replace('{answer}', answer)));
-    }
-    messages.push(createMessage('user', question));
-    return messages;
-};
-export function createDatasetExample(prompt, metadata) {
-    return {
-        prompt: prompt.map((message) => ({ ...message })), // defensive copy
-        metadata: metadata ? { ...metadata } : undefined,
-    };
-}
-export function extractGsm8kAnswer(raw) {
-    return extractHashAnswer(raw);
-}
-export function validateDatasetExample(example) {
-    if (!Array.isArray(example.prompt) || example.prompt.length === 0) {
-        throw new Error('Dataset example must contain at least one prompt message.');
-    }
-    for (const message of example.prompt) {
-        if (!message || typeof message.content !== 'string' || message.content.trim() === '') {
-            throw new Error('Prompt messages must include non-empty textual content.');
-        }
-        if (message.role !== 'system' && message.role !== 'user' && message.role !== 'assistant') {
-            throw new Error(`Unsupported chat role: ${String(message.role)}`);
-        }
-    }
-}
-function resolveBasePath(optionPath, options) {
-    const allowedRoot = getAllowedRoot(options);
-    if (!optionPath) {
-        // Default path - validate it's within allowed root
-        validatePathContainment(DEFAULT_BASE_PATH, allowedRoot);
-        return DEFAULT_BASE_PATH;
-    }
-    // Resolve and validate user-provided path
-    const resolved = resolvePath(allowedRoot, optionPath);
-    validatePathContainment(resolved, allowedRoot);
-    return resolved;
-}
-function datasetFileForSplit(split) {
-    if (!VALID_SPLITS.has(split)) {
-        throw new Error(`Unsupported GSM8K split "${split}". Expected one of: ${Array.from(VALID_SPLITS).join(', ')}`);
-    }
-    return `${split}.jsonl`;
-}
-function readDatasetFile(filePath) {
-    try {
-        return readFileSync(filePath, 'utf8');
-    }
-    catch (error) {
-        const message = error instanceof Error ? error.message : String(error);
-        throw new Error(`Failed to read dataset file at ${filePath}: ${message}`);
-    }
-}
-function readJsonl(path, limit) {
-    const fileContents = readDatasetFile(path);
-    const lines = fileContents.split(/\r?\n/).filter((line) => line.trim().length > 0);
-    const records = [];
-    const max = typeof limit === 'number' && limit >= 0 ? limit : Number.POSITIVE_INFINITY;
-    for (let i = 0; i < lines.length && records.length < max; i += 1) {
-        const line = lines[i];
-        try {
-            const parsed = JSON.parse(line);
-            if (typeof parsed.question !== 'string' || typeof parsed.answer !== 'string') {
-                throw new Error('Record must include string "question" and "answer" fields.');
-            }
-            records.push({ question: parsed.question, answer: parsed.answer });
-        }
-        catch (error) {
-            const message = error instanceof Error ? error.message : String(error);
-            throw new Error(`Failed to parse JSONL record at ${path}:${i + 1} - ${message}`);
-        }
-    }
-    return records;
-}
-export async function loadLocalGsm8kDataset(split, options = {}) {
-    const basePath = resolveBasePath(options.basePath, options);
-    const fileName = datasetFileForSplit(split);
-    const filePath = resolvePath(basePath, fileName);
-    // Additional validation: ensure the final file path stays within the base path
-    // This protects against any edge cases where the filename could escape
-    validatePathContainment(filePath, basePath);
-    const promptTemplate = options.promptTemplate ?? defaultPromptTemplate;
-    const records = readJsonl(filePath, options.limit);
-    const examples = records.map((record, index) => {
-        const prompt = promptTemplate(record.question, {
-            includeOneShot: options.includeOneShot,
-            oneShotExample: options.oneShotExample,
-        });
-        const example = createDatasetExample(prompt, {
-            split,
-            index,
-            raw_answer: record.answer,
-            ...options.metadata,
-        });
-        validateDatasetExample(example);
-        return example;
-    });
-    return examples;
-}
-export class LocalGsm8kDatasetLoader {
-    options;
-    constructor(options = {}) {
-        this.options = { ...options };
-    }
-    async load(split, limit) {
-        return loadLocalGsm8kDataset(split, { ...this.options, limit });
-    }
-}