/**
 * trainer.ts – MambaTrainer class
 *
 * REAL backpropagation. Until 2026-08 `_trainStep` handed the SAME `dLogits`
 * buffer to AdamW for every parameter, so no parameter ever received its own
 * gradient: `learn()` was SAFE (WSLA freezes the backbone, the trust region
 * bounds every step, the NaN guard blocks a poisoned weight, and the session
 * rolls back a regression) but INERT. It now runs a genuine backward pass —
 * `cpuModelBackward` in `model_cpu.ts` — and each parameter is stepped with its
 * OWN gradient.
 *
 * Where the maths runs, and why:
 *   • the backward pass is the CPU reference (`cpu_ops.ts` + `mamba1_cpu.ts`),
 *     a line-for-line mirror of the WGSL forward kernels. It is exact and it is
 *     PROVEN — `tests/backprop.test.ts` checks every parameter gradient against
 *     finite differences of the real loss, with no GPU involved;
 *   • the optimiser is {@link adamwUpdateInPlace}, the CPU mirror of
 *     `WEIGHT_UPDATE_WGSL`'s `adamw_update`, keeping the trust region and the
 *     NaN guard the safety story depends on;
 *   • the model's GPU buffers are written back after every optimiser step, so
 *     inference (which IS GPU) always serves the trained weights.
 *
 * The residual is a WGSL backward pass so training itself runs on the GPU. The
 * backward kernels exist (`SELECTIVE_SCAN_BACKWARD_WGSL`, `SSD_BACKWARD_WGSL`,
 * `CONV1D_BACKWARD_WGSL`, `LINEAR_BACKWARD_WGSL`, `ACTIVATIONS_BACKWARD_WGSL`)
 * but wiring them cannot be validated numerically without a WebGPU device, and
 * an unvalidated GPU backward is worse than a slower exact one.
 */
import { type ShardSpec } from '../optim/adamw.js';
import { LoRAAdapter, type LoRAConfig } from './lora.js';
import { type CpuModelDims } from './model_cpu.js';
import { HybridMambaModel, MambaModel } from '../model/mamba_model.js';
import { BPETokenizer } from '../tokenizer/bpe.js';
export interface TrainOptions {
    learningRate?: number;
    epochs?: number;
    /** Sequences trained together per optimiser step. Chunks of equal length are
     *  grouped into batches of this size; a ragged tail trains on its own. */
    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;
    /**
     * @deprecated The pre-clip gradient L2 norm is now computed exactly on the way
     * through the optimiser and is ALWAYS reported — it no longer costs a GPU
     * readback, so there is nothing to opt into. Accepted for API compatibility.
     */
    trackGradNorm?: boolean;
    /**
     * Train LOW-RANK ADAPTERS over the block projections instead of the weights
     * themselves (see {@link LoRAAdapter}). The base weights are frozen; only A and
     * B train, so an adapt costs a fraction of the optimiser state and the result is
     * a small, mergeable delta rather than a whole checkpoint.
     *
     * `targets` selects which parameters get an adapter, by their unqualified name
     * (`wInProj`, `wXProj`, `wDtProj`, `wOutProj`, `wConv`, `A_log`, `embedding`);
     * omit for {@link DEFAULT_LORA_TARGETS}. A rank larger than a target's smaller
     * dimension is clamped to it rather than throwing — the small projections
     * (`wDtProj` is `dInner x dtRank`) would otherwise reject a sensible global rank.
     */
    lora?: LoRAConfig & {
        targets?: readonly string[];
    };
    /**
     * Batches to accumulate gradients over before one optimiser step. Default 1.
     * The effective batch size is `batchSize * gradientAccumulation` at the memory
     * cost of `batchSize` — the standard way to train at a batch size the host
     * cannot hold at once.
     */
    gradientAccumulation?: number;
    /**
     * Shard the optimiser state (the Adam moments — 2x the parameter bytes, the
     * dominant cost of a full fine-tune). This instance then OWNS only parameters
     * where `index % count === index`, and allocates moments for those alone. Run
     * `count` trainers over the same model and the union of their steps equals one
     * unsharded step. Same contract as {@link ShardSpec} on {@link AdamW}.
     */
    shard?: ShardSpec;
    /**
     * Recompute each layer's activations in the backward pass instead of retaining
     * all of them. Gradients are identical; peak memory drops from O(layers)
     * activation caches to one, at the cost of a second forward per layer.
     */
    activationCheckpointing?: 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;
/** One training example: a token window and its next-token targets. */
interface TrainChunk {
    inputs: number[];
    targets: number[];
}
/** A batch of equal-length chunks trained in one optimiser step. */
interface TrainBatch {
    inputs: number[];
    targets: number[];
    batch: number;
    seqLen: number;
}
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` 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 readonly _dims;
    /** CPU mirror of every model tensor — the gradient engine's working set. */
    private _mirror;
    private _mirrorByName;
    /** Low-rank adapters by parameter name, when training with `lora`. */
    private _adapters;
    /** Frozen base weights the adapters sit on top of (LoRA only). */
    private _loraBase;
    private _shard;
    constructor(model: HybridMambaModel | MambaModel, tokenizer?: BPETokenizer | null);
    /**
     * The low-rank adapters this trainer is training, by parameter name — empty
     * unless `train({ lora })` was used. Exposed so a caller can serialise just the
     * adapter (a few hundred KB) instead of the whole checkpoint.
     */
    get adapters(): ReadonlyMap<string, LoRAAdapter>;
    /** Trainable scalar count for the current mode — the LoRA saving, measured. */
    trainableParamCount(): number;
    /** Whether this trainer owns parameter `index` under the active shard. */
    private _owns;
    /**
     * Build (once) an adapter for every LoRA-target parameter, snapshot the frozen
     * base, and return the set. Rank is clamped per parameter to its smaller
     * dimension so one global rank works across projections of very different shape.
     */
    private _installAdapters;
    /** Write `base + adapter.delta()` into the mirror, so the forward sees the adapted weights. */
    private _materialiseAdapters;
    /** Get-or-create the Adam first/second moments for a parameter, by name. */
    private _momentFor;
    /**
     * Pull every model tensor off the GPU into the CPU mirror. Done once per
     * `train()` call (not per step) so an externally loaded checkpoint or a
     * previous session's weights are picked up, after which the mirror is
     * authoritative and each step writes the updated tensors back.
     */
    private _syncMirrorFromGpu;
    train(input: string | number[], opts?: TrainOptions): Promise<number[]>;
    /**
     * One optimiser step over one batch: real forward + backward, a global
     * gradient-norm clip over the trainable set, then AdamW per parameter with
     * its OWN gradient, written straight back to the GPU buffer.
     */
    private _trainStep;
    /**
     * AdamW over the trainable set. Each parameter is paired with its OWN
     * gradient by NAME — the bug this replaces passed one shared `dLogits`
     * buffer to every parameter — and the updated tensor is uploaded to the GPU
     * so inference serves it.
     */
    private _adamwStep;
    /** Get-or-create moments for a synthetic parameter (a LoRA factor), by name. */
    private _momentNamed;
    /**
     * Bytes of optimiser state (Adam m + v) this trainer holds. Under a shard that
     * is 1/count of the full state; under LoRA it is the adapters' state, not the
     * weights'. The measurable point of both features.
     */
    optimizerStateBytes(): number;
    evaluate(input: string | number[]): Promise<number>;
}
/**
 * Derive the gradient engine's dimensions from a model, rejecting layer types
 * the CPU reference cannot differentiate.
 *
 * A hybrid schedule (`jamba`/`zamba`, or any explicit schedule containing
 * `mamba2`/`mamba3`/`attention`) fails LOUDLY here rather than silently training
 * on gradients that do not exist — which is exactly what the old trainer did.
 */
export declare function cpuDimsFor(model: HybridMambaModel): CpuModelDims;
export declare function buildChunks(ids: number[], seqLen: number): TrainChunk[];
/**
 * Group equal-length chunks into batches of at most `batchSize`.
 *
 * `batchSize` used to be accepted and ignored — every chunk trained alone. The
 * batch dimension the model and the gradient engine both take is real, so it is
 * honoured here; a ragged tail chunk simply trains as its own batch of one.
 */
export declare function buildBatches(chunks: TrainChunk[], batchSize: number): TrainBatch[];
export {};
//# sourceMappingURL=trainer.d.ts.map