/**
 * trainer.ts – MambaTrainer class
 */
import { HybridMambaModel, MambaModel } from '../model/mamba_model.js';
import { BPETokenizer } from '../tokenizer/bpe.js';
export interface TrainOptions {
    learningRate?: number;
    epochs?: number;
    batchSize?: number;
    seqLen?: number;
    maxGradNorm?: number;
    weightDecay?: number;
    beta1?: number;
    beta2?: number;
    eps?: number;
    wsla?: boolean;
    /** Trust region: max |Δθ| per optimizer step. 0 disables. Defaults to
     *  {@link WSLA_MAX_DELTA} in WSLA (write-through) mode, else 0 (full training). */
    maxDelta?: number;
    /** Read back the pre-clip gradient L2 norm each step (the `grad_norm_reduce`
     *  kernel already computes it — this just maps it to the CPU). OFF by default: the
     *  readback is a GPU sync point that would slow a long finetune. `adapt()` turns it
     *  ON so the norm reaches {@link AdaptResult} as an instability signal. */
    trackGradNorm?: boolean;
    onEpochEnd?: ((epoch: number, loss: number, gradNorm?: number) => void) | null;
}
/**
 * Default per-step trust region for WSLA / write-through adaptation. Small
 * enough that a single `adapt()` nudges the narrow params without lurching, so
 * repeated adapts stay stable (and any that regress are cheaply rolled back by
 * the session). Full-training callers pass `maxDelta: 0` (or omit it in
 * non-WSLA mode) for unbounded steps.
 */
export declare const WSLA_MAX_DELTA = 0.05;
export declare class MambaTrainer {
    model: HybridMambaModel;
    tokenizer: BPETokenizer | null;
    device: GPUDevice;
    /**
     * Adam moments keyed by parameter NAME (not array index). Name-keying is what
     * lets WSLA toggle safely: a narrow write-through step updates only the
     * `layer{i}.wXProj/bXProj` subset, a full fine-tune updates everything, and
     * both reuse the SAME `m`/`v` buffers per parameter. Index-keyed moments (the
     * old scheme) silently misaligned the moment with the wrong parameter the
     * moment the trainable set changed shape — corrupting the update.
     */
    private _moments;
    private _step;
    private _adamwPipeline;
    private _clipReducePipeline;
    private _clipScalePipeline;
    constructor(model: HybridMambaModel | MambaModel, tokenizer?: BPETokenizer | null);
    /** Get-or-create the Adam first/second moments for a parameter, by name. */
    private _momentFor;
    train(input: string | number[], opts?: TrainOptions): Promise<number[]>;
    private _trainStep;
    private _adamwStep;
    /**
     * Clip gradients to `maxNorm` in place. The reduce kernel writes Σg² into
     * `normSqBuf`; the scale kernel then rescales the grads. When `trackNorm` is set
     * we read Σg² back and return the true pre-clip L2 norm (√Σg²) — the kernel
     * already computed it, so this is only a small buffer readback (a GPU sync,
     * hence opt-in). Returns null when not tracking.
     */
    private _clipGradients;
    evaluate(input: string | number[]): Promise<number>;
}
//# sourceMappingURL=trainer.d.ts.map