Train

train.TrainSpec

One fine-tuning request, backend-agnostic.

Usage

Source

train.TrainSpec(
    name,
    base,
    dataset,
    hyperparams,
    method="sft",
    lora=LoraSpec(),
    backend=None,
    max_usd=None,
    resume_from=None,
    resume_reference=None
)

Attributes

name: str

The registry family the trained artifact is registered under.

base: BaseModelSpec

The base model, addressed across backends.

dataset: DatasetSource

Where the training data comes from.

hyperparams: Hyperparams

Steps, batch, learning rate, sequence length, seed.

method: Method

SFT or DPO. local trains SFT only.

lora: LoraSpec

LoRA rank, alpha, and target modules.

backend: BackendName | None

An explicit backend override; None selects by availability.

max_usd: float | None

A per-run spend cap for the metered backends (tinker, modal); None falls back to that backend’s configured spend_cap_usd.

resume_from: StateHandle | None

The cross-run continuation seed — a prior checkpoint’s training state the run warm-starts its policy from. Set by from_checkpoint(); excluded from the run key so a continued run still recovers itself by identity.

resume_reference: StateHandle | None

The DPO reference anchor a continuation freezes its reference client at, also set by from_checkpoint() and excluded from the run key.

Methods

Name	Description
from_checkpoint()	A fresh continuation of prior: its training state warm-starts and anchors the new run.

from_checkpoint()

A fresh continuation of prior: its training state warm-starts and anchors the new run.

Usage

Source

from_checkpoint(prior, *, dataset, name, **overrides)

The seed is prior.state — a ~athome.train.state.StateHandle, so a sampler_path cannot be threaded as a training seed: the distinct-artifact rule is a type error here, not a convention. The DPO reference anchor is that same prior state, so a continued preference run freezes its reference at the model it is moving away from. base is inherited from prior; everything else (hyperparams, method, lora, max_usd) arrives through overrides.

Parameters

prior: Checkpoint

The checkpoint to continue from; its base and state seed the new spec.

dataset: DatasetSource

The new run’s corpus — fresh negatives for the next DITTO iteration.

name: str

The registry family the continuation registers under.

overrides: object = {}

Any other TrainSpec field (hyperparams is required).

Returns

TrainSpec

A spec whose policy and DPO reference both seed from prior.state.

train.TrainResult

The outcome of one run: the artifact, its eval scalar, and its registry entry.

Usage

Source

train.TrainResult(checkpoint, metric, leaderboard, version, promoted)

Attributes

checkpoint: Checkpoint

The servable artifact.

metric: float

The trained arm’s primary metric on the evaluation bake-off.

leaderboard: Leaderboard

The full bake-off result the metric came from.

version: VersionInfo

The registry entry the artifact was registered as.

promoted: bool

True when this artifact became the family’s current.

train.Checkpoint

The servable artifact one backend produced: a standalone 4-bit MLX model directory.

Usage

Source

train.Checkpoint(
    base,
    backend,
    method,
    step,
    mlx_path,
    adapter_dir,
    train_cost_usd,
    sampler_path,
    state
)

mlx_path is the only serve path — rapid-mlx serve takes a single model and has no adapter flag, so every backend fuses its LoRA into the base. adapter_dir is the intermediate adapter kept for provenance (tinker and modal train a PEFT adapter first); consumers with adapter support can also serve it directly.

Attributes

base: BaseModelSpec

The base model the LoRA was trained over.

backend: BackendName

The backend that produced it.

method: Method

The method it was trained with.

step: int

The step count the weights were saved at.

mlx_path: Path

The fused standalone 4-bit MLX model directory.

adapter_dir: Path | None

The mlx-lm adapter directory, when one was materialized.

train_cost_usd: float

What the run spent; 0.0 on the unmetered local backend.

sampler_path: str | None

The opaque tinker:// sampler checkpoint the fused adapter came from, or None for a backend that produces no hosted sampler checkpoint (local, modal).

state: StateHandle

The training-state handle this checkpoint can be continued from, at its backend’s declared fidelity (Tinker weights+optimizer; local/modal weights-only).

train.Adapter

A materialized non-fused MLX adapter that can be served or fused later.

Usage

Source

train.Adapter(step, adapter_dir, train_cost_usd, sampler_path, state)

Attributes

step: int

The training step whose weights this adapter carries.

adapter_dir: Path

The mlx-lm adapter directory.

train_cost_usd: float

What the training run spent to produce the saved weights.

sampler_path: str

The opaque tinker:// sampler checkpoint the adapter was materialized from.

state: StateHandle

The training-state handle (weights+optimizer) the snapshot also saved, distinct from the sampler path — what a continuation re-seeds its optimizer from.

train.BaseModelSpec

One base model addressed across every backend, pinned to exact weights.

Usage

Source

train.BaseModelSpec(
    mlx,
    hf,
    hf_revision,
    mlx_revision,
    tinker,
    num_layers,
    serves_locally=True
)

Every repo id carries its commit: an unpinned id floats with the hub’s head, so two runs could train against different weights — or an adapter trained on one revision could be fused into an independently-updated MLX snapshot — while their fingerprints claimed to be the same run.

Attributes

mlx: MlxModelId

The 4-bit MLX id used for local serving and PEFT-to-MLX conversion.

hf: HfRepoId

The HuggingFace repo for modal (GPU) training and snapshotting.

hf_revision: str

The commit the hf weights are pinned to.

mlx_revision: str

The commit the mlx weights are pinned to.

tinker: TinkerModelId | None

The Tinker base id, or None when Tinker cannot train it.

num_layers: int

Layer count, needed to write an mlx-lm adapter_config.json.

serves_locally: bool

False when the LoRA cannot load into mlx-lm — a split linear_attn maps to a fused in_proj_qkv, so the adapter has no mlx-lm counterpart.

train.Hyperparams

The optimization knobs shared by every backend.

Usage

Source

train.Hyperparams(
    steps, batch_size=4, learning_rate=0.0001, max_seq_len=4096, seed=SEED
)

Parameter Attributes

steps: int

batch_size: int = 4

learning_rate: float = 0.0001

max_seq_len: int = 4096

seed: int = SEED


train.LoraSpec

The LoRA adapter’s shape: its rank, scale, and which modules it wraps.

Usage

Source

train.LoraSpec(
    rank=16,
    alpha=32,
    dropout=0.0,
    target_modules=STD_MODULES,
    train_mlp=True,
    train_attn=True,
    train_unembed=False
)

Parameter Attributes

rank: int = 16

alpha: int = 32

dropout: float = 0.0

target_modules: tuple[str, …] = STD_MODULES


train_mlp: bool = True

train_attn: bool = True

train_unembed: bool = False

train.HfDatasetRef

A cc-steer HuggingFace export: one config of one dataset repo.

Usage

Source

train.HfDatasetRef(repo, config, split="train")

Attributes

repo: HfRepoId

The owner/name dataset repo id.

config: str

The export config name (a cc-steer export ships sft and dpo).

split: str

The split to train on.

train.LocalJsonlRef

A local jsonl corpus: one .jsonl file, or a directory of them.

Usage

Source

train.LocalJsonlRef(path)

Attributes

path: Path

The jsonl file, or a directory whose *.jsonl files are read in sorted order. SFT rows are mlx-lm chat rows ({"messages": [...]}); DPO rows carry prompt/chosen/rejected.

train.Rows

An in-memory SFT corpus: examples a caller already holds, with no jsonl detour.

Usage

Source

train.Rows(examples)

Attributes

examples: tuple[SftExample, …]

The SFT examples to train on directly.

train.TrainSettings

Backend selection, the mlx-lm sidecar pin, and the train roots, bound to [train].

Usage

Source

train.TrainSettings()

train.TinkerSettings

Tinker credentials, spend cap, and per-model token prices, bound to [train.tinker].

Usage

Source

train.TinkerSettings()

train.TinkerPrice

One Tinker base model’s per-million-token price, split by the token class it bills.

Usage

Source

train.TinkerPrice(prefill, sample, train)

Tinker meters three classes at three rates, so a run costs the sum over the classes it actually billed — a flat rate would overcharge a forward-only pass at the training rate.

Attributes

prefill: float

Tokens through a forward-only pass, such as a frozen reference policy’s logprobs.

sample: float

Tokens generated by a sampling client.

train: float

Tokens through a forward_backward and its optim_step.

Example

>>> TinkerPrice(prefill=0.13, sample=0.40, train=0.40)

train.TrainBackend

What every training backend exposes: a cheap probe, its methods, and a train call.

Usage

Source

train.TrainBackend()

available and supports are static so select() can probe a backend without constructing its settings — a required secret would raise before the backend was even chosen.

Methods

Name	Description
available()	Whether this backend can run here: its credentials or hardware are present.
from_settings()	Construct the backend from its configuration section.
supports()	Whether this backend trains method.
train()	Train spec, journaling progress to sink, and converge on a servable artifact.

available()

Whether this backend can run here: its credentials or hardware are present.

Usage

Source

available()

from_settings()

Construct the backend from its configuration section.

Usage

Source

from_settings()

supports()

Whether this backend trains method.

Usage

Source

supports(method)

train()

Train spec, journaling progress to sink, and converge on a servable artifact.

Usage

Source

train(spec, *, sink, work_dir, resume=None, store=None)

Every file the run writes — data splits, adapters, the fused model — goes under work_dir, which athome.train.run() mints fresh per run. A backend never derives its own output directory: two runs of one family would then overwrite each other’s weights, including weights the registry has already registered.

resume seeds the run from a prior training state — a same-run crash recovery or a cross-run continuation — and store is the run-state ledger progress persists to; both are None for a fresh, unledgered run. The produced Checkpoint always carries a state handle at this backend’s declared state_fidelity.

train.tinker.TinkerBackend

LoRA fine-tuning on Tinker’s managed trainers, converging on a fused local MLX model.

Usage

Source

train.tinker.TinkerBackend(settings)

SFT trains against Tinker’s cross_entropy loss over prompt-masked Datum s. DPO has no named Tinker loss, so it runs through forward_backward_custom: the trained policy’s logprobs come back as torch tensors, dpo_loss() scores them against a frozen reference policy’s cached logprobs, and the SDK backpropagates the client-side gradient.

Every op runs through the pipelined ~athome.train.engine.execute() stream, which keeps the server’s turnstile full so a step costs one clock cycle instead of three, and lets checkpoint saves and their eval scoring ride between steps without stalling. Every run is projected against spend_cap_usd before its first billable call and charged step by step, so a run that cannot fit the cap aborts having spent nothing.

Parameter Attributes

settings: TinkerSettings

Example

>>> checkpoint = await TinkerBackend.from_settings().train(spec, sink=sink)

Methods

Name	Description
available()	Whether the Tinker SDK is installed and keyed: a key in the env, or ~/.athome/tinker.env.
cost()	What model bills for these token counts, each class at its own per-Mtok rate.
fit()	Train spec on Tinker, journaling each step and saving checkpoints at the cadence.
from_settings()	Load [train.tinker], taking the API key from ~/.athome/tinker.env when unset.
fuse()	Fuse a materialized adapter into a standalone servable MLX checkpoint.
lora_client()	A training client for spec: fresh from model, or seeded from a prior Tinker state.
materialize()	Download and convert one saved Tinker checkpoint into a non-fused MLX adapter.
require_matches()	Refuse a resume whose saved LoRA shape disagrees with what spec.lora restores to.
sample()	Sample free-form completions from a Tinker checkpoint, or the base model when path is None.
score()	Score weighted token sequences against a saved Tinker sampling checkpoint.
supports()	Whether Tinker trains method here: SFT always, DPO only where torch can back it.
train()	Train spec on Tinker and fuse the resulting final adapter into a servable MLX model.

available()

Whether the Tinker SDK is installed and keyed: a key in the env, or ~/.athome/tinker.env.

Usage

Source

available()

The SDK is half of it. Selection treats an available backend as one that can actually run, so claiming availability without tinker on the path picks this backend and then dies importing it, instead of falling through to one that would have worked.

cost()

What model bills for these token counts, each class at its own per-Mtok rate.

Usage

Source

cost(*, model, prefill=0, sample=0, train=0)

Parameters

model: TinkerModelId

The Tinker base model whose price sheet applies.

prefill: int = 0

Tokens through a forward-only pass.

sample: int = 0

Tokens generated by a sampling client.

train: int = 0

Tokens through a forward_backward and its optim_step.

Returns

float

The USD total across the three classes.

fit()

Train spec on Tinker, journaling each step and saving checkpoints at the cadence.

Usage

Source

fit(
    spec,
    *,
    sink,
    budget,
    checkpoints=CheckpointPolicy(),
    eval_rows=None,
    resume=None,
    store=None,
    run_tag=None
)

The pool is rendered and under-filled runs aborted before anything billable; the whole schedule — training, the DPO reference pass, and every checkpoint’s eval scoring — is projected against budget in one reservation. The stream then runs, spend and the journal reconciling at each drain point.

Parameters

spec: TrainSpec

The fine-tuning request: base, dataset, hyperparams, method, LoRA.

sink: RunSink

The run journal; one record lands per step, plus a checkpoint event per save.

budget: SpendGuard

The spend envelope; the whole schedule’s projected cost is reserved against it before the first billable call and reconciled to actuals as the run drains.

checkpoints: CheckpointPolicy = CheckpointPolicy()


Which fractions of the run to snapshot as intermediates; the final step is always snapshotted and kept forever.

eval_rows: Sequence[EvalRow] | None = None

Pre-tokenized rows scored against every checkpoint’s weights, or None.

resume: Resume | None = None

The restore input — seeds the training client, slices the plan to the steps a same-run crash left unrun, and carries the DPO reference anchor — or None for a run from base.

store: RunStateStore | None = None

The run-state store the fold persists this run’s most progressed RunState to at each snapshot, or None to train without a resume ledger (an eval-only fit through observe/retrain).

run_tag: str | None = None

This run’s identity namespace, folded into every snapshot name so no two runs of one spec ever share a state name. The managed lifecycle (~athome.train.run() and ~athome.train.retrain()) passes its per-run work_dir name; None mints a fresh nonce for a standalone non-resumable fit (observe); a direct resume keeps its deterministic -r{base_step} names.

Returns

TrainReport

The report: per-step records, the saved checkpoints with their eval scores, the drop

count, the wall-clock, and the total metered spend.

Raises

UnservableBase

The base has no Tinker id.

UnsupportedLoraShape

The LoRA shape asks for something Tinker cannot train.

TorchRequired

The method is dpo and torch is not installed locally.

InsufficientData

The surviving pool is smaller than one batch.

OverlongEvalRows

An eval row exceeds max_seq_len.

SpendExceeded

The projected run cost crosses the spend cap.

from_settings()

Load [train.tinker], taking the API key from ~/.athome/tinker.env when unset.

Usage

Source

from_settings()

fuse()

Fuse a materialized adapter into a standalone servable MLX checkpoint.

Usage

Source

fuse(adapter, spec, *, work_dir)

Parameters

adapter: Adapter

The non-fused MLX adapter to merge into its base model.

spec: TrainSpec

The fine-tuning request whose base and method describe the checkpoint.

work_dir: Path

This fusion’s directory; the standalone model is written under it.

Returns

Checkpoint

The fused checkpoint, preserving the adapter’s step and training cost.

lora_client()

A training client for spec: fresh from model, or seeded from a prior Tinker state.

Usage

Source

lora_client(service, spec, *, model, seed=None)

With no seed the client is a fresh LoRA over the base. A TinkerState seed restores the weights and optimizer via create_training_client_from_state_with_optimizer_async — after require_matches() refuses a saved shape that disagrees with spec.lora. A foreign handle (LocalState/ModalState) cannot seed a Tinker client and is refused loudly.

materialize()

Download and convert one saved Tinker checkpoint into a non-fused MLX adapter.

Usage

Source

materialize(saved, spec, *, work_dir, cost)

Parameters

saved: SavedCheckpoint

Any saved checkpoint from fit(), including an intermediate.

spec: TrainSpec

The fine-tuning request whose base model the adapter targets.

work_dir: Path

This materialization’s directory; PEFT and MLX adapter files land under it.

cost: float

The training spend attributed to the adapter.

Returns

Adapter

The non-fused MLX adapter, ready for direct adapter serving or a later fuse().

require_matches()

Refuse a resume whose saved LoRA shape disagrees with what spec.lora restores to.

Usage

Source

require_matches(service, state_path, spec)

The SDK reads the rank and module toggles off the checkpoint’s weights_info and ignores spec.lora on a from-state creation, so a silent disagreement would train a different adapter than the spec asked for. This reads that saved shape and fails loudly on a mismatch — or, when the saved state has expired or vanished, refuses just as loudly rather than falling back to a fresh run.

sample()

Sample free-form completions from a Tinker checkpoint, or the base model when path is None.

Usage

Source

sample(path, prompts, *, base, budget, max_tokens, temperature, seed=None)

Every prompt is rendered with the base’s training chat template and a generation prompt, so a sample matches what the policy was trained to continue. The whole batch is projected against budget in one conservative reservation — full prefill plus max_tokens per prompt — before any sampling client exists, then reconciled to the real generated token counts, so a batch that cannot fit the envelope aborts having spent nothing and a short generation bills less.

Parameters

path: str | None

The opaque tinker:// sampler checkpoint to sample from, or None to sample the base model directly (iteration-zero negatives, before any checkpoint exists).

prompts: Sequence[Sequence[Message]]

The chat prompts to complete, one completion returned per prompt, in input order.

base: BaseModelSpec

The base model identity: its chat template tokenizes the prompts and its price sheet bills the run, and, when path is None, its Tinker id is the model sampled.

budget: SpendGuard

The spend envelope; the conservative full-prefill-plus-max_tokens projection is reserved against it before any client exists, then reconciled to the real generated cost.

max_tokens: int

The generation cap per prompt; also the per-prompt sample count the projection reserves.

temperature: float

The sampling temperature applied to every prompt.

seed: int | None = None

The base seed; prompt i samples with seed + i for reproducibility, or None to leave every prompt unseeded.

Returns

SampledSequence

One sampled sequence per prompt, preserving input order, each with its decoded text, token

…

counts, and billed cost.

Raises

UnservableBase

The base has no Tinker identity.

SpendExceeded

The projected prefill and max_tokens per prompt cross the envelope.

score()

Score weighted token sequences against a saved Tinker sampling checkpoint.

Usage

Source

score(path, rows, *, base, budget)

Parameters

path: str

The opaque tinker:// sampler checkpoint path.

rows: Sequence[EvalRow]

Pre-tokenized weighted rows to score, in return order.

base: BaseModelSpec

The base model identity used to price the sampling requests.

budget: SpendGuard

The spend envelope; the projected prefill and one sampled token per row are reserved against it before any client exists, then reconciled once the rows score.

Returns

tuple[ScoredSequence, …]

One weighted sequence score per row, preserving input order.

Raises

UnservableBase

The base has no Tinker identity.

SpendExceeded

The projected prefill and one sampled token per row cross the envelope.

supports()

Whether Tinker trains method here: SFT always, DPO only where torch can back it.

Usage

Source

supports(method)

Tinker has no named preference loss, so DPO runs through forward_backward_custom and dpo_loss() backprops it in-process — without torch there is nothing to run it with.

train()

Train spec on Tinker and fuse the resulting final adapter into a servable MLX model.

Usage

Source

train(spec, *, sink, work_dir, resume=None, store=None)

The spend cap binds per run, not per attempt: a resume seeds the SpendGuard with the cost already billed under this run’s key, so a crash-looping run never bills N times the cap.

Parameters

spec: TrainSpec

The fine-tuning request: base, dataset, hyperparams, method, LoRA, spend cap.

sink: RunSink

The run journal; one record lands per step with its loss, tokens, and spend.

work_dir: Path

This run’s own directory; the adapter and the fused model are written under it.

resume: Resume | None = None

The restore input for a same-run recovery or a cross-run continuation, or None to train from base.

store: RunStateStore | None = None

The run-state store the fit persists its progress to, or None for no resume ledger.

Returns

Checkpoint

The checkpoint, whose mlx_path is the fused standalone 4-bit MLX model.

Raises

UnservableBase

The base has no Tinker id, or no mlx-lm counterpart to fuse into.

UnsupportedLoraShape

The LoRA shape asks for something Tinker cannot train.

LoraShapeDrift

A resume seed’s saved LoRA shape disagrees with spec.lora.

TorchRequired

The method is dpo and torch is not installed locally.

InsufficientData

The surviving pool is smaller than one batch.

SpendExceeded

The projected run cost crosses the spend cap.

train.CheckpointPolicy

Which fractions of a run to snapshot as intermediate checkpoints, and for how long.

Usage

Source

train.CheckpointPolicy(at=(), ttl_seconds=INTERMEDIATE_TTL_S)

The final step is always snapshotted separately and kept forever; these fractions place the intermediate saves before it. at holds fractions in (0, 1], strictly increasing.

Parameter Attributes

at: tuple[float, …] = ()


ttl_seconds: int = INTERMEDIATE_TTL_S


Example

>>> CheckpointPolicy(at=(0.25, 0.5, 0.75)).steps_for(12)

(3, 6, 9)

Methods

Name	Description
steps_for()	The intermediate step numbers for a run of total steps, ascending and unique.

steps_for()

The intermediate step numbers for a run of total steps, ascending and unique.

Usage

Source

steps_for(total)

Each fraction rounds up to at least step 1; the final step is excluded because it is always snapshotted on its own. Fractions that collapse onto the same step or onto the final step drop out.

train.EvalRow

A pre-tokenized weighted scoring row, mapping one-to-one onto a Tinker Datum.

Usage

Source

train.EvalRow(tokens, weights)

The row’s tokens are the full sequence; the weight-carrying positions are the ones a checkpoint’s forward pass scores. Unlike a message-level example, this can weight a single sentinel-token position, which is what a next-token discriminator scores.

Attributes

tokens: tuple[int, …]

The full token sequence, prompt and completion.

weights: tuple[float, …]

One weight per token; a nonzero weight marks a scored position.

train.ScoredSequence

One eval row’s score against a checkpoint’s weights.

Usage

Source

train.ScoredSequence(logprob, weight)

Attributes

logprob: float

The weighted sum of per-token logprobs over the scored positions.

weight: float

The exact weight mass, so a caller can normalize logprob per scored weight.

train.SavedCheckpoint

A Tinker sampler checkpoint saved mid-run, with its eval scores against those weights.

Usage

Source

train.SavedCheckpoint(step, sampler_path, state, final, scores)

Attributes

step: int

The training step whose weights this snapshot captured.

sampler_path: str

The opaque tinker:// address the sampler weights were saved to.

state: StateHandle

The training state (weights+optimizer) saved atomically alongside the sampler weights, a distinct artifact — a snapshot that saved sampler weights without a training state has no representation.

final: bool

True for the run’s last checkpoint, which is kept forever.

scores: tuple[ScoredSequence, …] | None

The eval rows scored against these weights, order-preserving, or None when the run carried no eval rows.

train.StepRecord

One training step’s telemetry.

Usage

Source

train.StepRecord(step, tokens, usd, metrics)

Attributes

step: int

The one-based step number.

tokens: int

The step’s trained token count.

usd: float

What this step alone cost, not the running total.

metrics: dict[str, float]

The loss and any method-specific metrics (SFT loss; DPO loss/margin/accuracy).

train.TrainReport

Everything one fit produced: its per-step records and its saved checkpoints.

Usage

Source

train.TrainReport(method, steps, checkpoints, dropped, wall_s, train_cost_usd)

Attributes

method: Method

The method trained.

steps: tuple[StepRecord, …]

The per-step records, ascending.

checkpoints: tuple[SavedCheckpoint, …]

The saved checkpoints, ascending by step; the last is the final.

dropped: int

How many examples were dropped for exceeding max_seq_len.

wall_s: float

Wall-clock seconds spent executing the schedule.

train_cost_usd: float

The run’s total metered spend.

Attributes

Name	Description
final	The run’s final checkpoint.

final

The run’s final checkpoint.

final: SavedCheckpoint

train.retrain

train.retrain

Usage

Classes

Name	Description
RetrainOutcome	Everything one retrain() produced, with no side effect taken on the caller’s behalf.

RetrainOutcome

Everything one retrain() produced, with no side effect taken on the caller’s behalf.

Usage

Source

RetrainOutcome(report, best, adapter, served, verdict)

The consumer decides what to do with the verdict — register, promote, journal, prune — after retrain returns; the orchestrator itself writes to no registry.

Attributes

report: TrainReport

The full training report fit produced, including every saved checkpoint.

best: SavedCheckpoint

The checkpoint select argmaxed over report.checkpoints, the one materialized.

adapter: Adapter

The servable adapter best was materialized into.

served: dict[str, float]

The scores artifact_scorer read off the materialized adapter.

verdict: GateVerdict

The gate’s decision over served.

Functions

Name	Description
retrain()	Train, pick the strongest checkpoint, materialize it, score the artifact, and gate it.

retrain()

Train, pick the strongest checkpoint, materialize it, score the artifact, and gate it.

Usage

Source

retrain(
    backend,
    spec,
    *,
    checkpoints,
    eval_rows,
    budget,
    select,
    artifact_scorer,
    gate,
    work_dir,
    sink
)

The pipeline shape lives here; the domain arrives as callables. fit runs the whole training schedule and scores every checkpoint’s eval rows on the turnstile; select scores each of the resulting checkpoints and the argmax is materialized; artifact_scorer reads the servable adapter locally once fit’s stream has fully drained; gate turns those scores into a verdict. Nothing here registers, promotes, or journals a decision — those side effects belong to the caller, after this returns, against the caller’s own registry.

Parameters

backend: TinkerBackend

The Tinker backend that trains and materializes; checkpoint selection is Tinker-shaped, so no other backend qualifies.

spec: TrainSpec

The fine-tuning request forwarded to fit and materialize.

checkpoints: CheckpointPolicy

Which fractions of the run to snapshot, passed straight to fit.

eval_rows: Sequence[EvalRow] | None

Pre-tokenized rows scored against every checkpoint’s weights, or None.

budget: SpendGuard

The spend envelope threaded verbatim into fit; the training schedule’s projected cost reserves against it before the first billable call and reconciles to actuals as the run drains. retrain mints no envelope of its own, so this caller-owned cap is the one ledger every billable call in the run debits.

select: Callable[[SavedCheckpoint], float]

Scores one checkpoint; retrain materializes the argmax over report.checkpoints. A raising select propagates and fails the run.

artifact_scorer: Callable[[Adapter], dict[str, float]]

Reads scores off the materialized adapter, synchronously and locally.

gate: Callable[[dict[str, float]], GateVerdict]

Turns the artifact’s scores into a promote/reject verdict.

work_dir: Path

This run’s directory, where materialize writes the adapter.

sink: RunSink

The run journal fit appends each step and checkpoint to.

Returns

RetrainOutcome

The report, the selected checkpoint, its materialized adapter, the artifact scores, and the

gate’s verdict.

Raises

UnservableBase

The base has no mlx-lm LoRA counterpart to materialize into, so the run would bill a full hosted fit only to refuse at materialize; it aborts before fit.

train.RetrainOutcome

Everything one retrain() produced, with no side effect taken on the caller’s behalf.

Usage

Source

train.RetrainOutcome(report, best, adapter, served, verdict)

The consumer decides what to do with the verdict — register, promote, journal, prune — after retrain returns; the orchestrator itself writes to no registry.

Attributes

report: TrainReport

The full training report fit produced, including every saved checkpoint.

best: SavedCheckpoint

The checkpoint select argmaxed over report.checkpoints, the one materialized.

adapter: Adapter

The servable adapter best was materialized into.

served: dict[str, float]

The scores artifact_scorer read off the materialized adapter.

verdict: GateVerdict

The gate’s decision over served.

train.run()

Train spec, score the artifact against evaluation, and register what came out.

Usage

Source

train.run(spec, *, evaluation)

The run gets a fresh new_work_dir() and a free_port() of its own, so concurrent runs — of one family or several — share neither artifacts nor a server. The selected backend trains the LoRA there and converges on a fused standalone MLX directory. That directory is served locally and bakes off against evaluation’s arms as the trained arm; the run is registered under spec.name — weights and all, copied into the registry — and promoted to the family’s current only when the trained arm wins the bake-off and clears its statistical gate.

Progress journals to <work_dir>/progress.jsonl, and the metric lands in .athome-metric.json in the working directory — the structured channel athome.research reads — so athome train run drops straight into a research loop as a metric command.

Parameters

spec: TrainSpec

The fine-tuning request; it also picks the backend.

evaluation: BakeoffSpec

The bake-off the trained artifact is scored against. Its arms[0] is the baseline the gate measures lift over, and its primary_metric is the scalar reported as ~athome.train.TrainResult.metric.

Returns

TrainResult

The checkpoint, its primary metric, the leaderboard it came from, the registered

version, and whether that version was promoted.

Raises

PreflightFailure

A mandatory backend, dataset, cost, or evaluation probe fails.

train.evaluate()

Serve checkpoint’s fused artifact and bake it off against evaluation’s arms.

Usage

Source

train.evaluate(checkpoint, evaluation, *, port)

The fused MLX directory is served with rapid-mlx on port — an override that leaves the configured server alone — and appended to the bake-off as the trained arm, whose client is bound to that endpoint. The server is torn down before returning.

Parameters

checkpoint: Checkpoint

The artifact under test; its mlx_path is what gets served.

evaluation: BakeoffSpec

The bake-off supplying the task, corpus, and baseline arms.

port: int

The port to serve on. It, with the model, names the managed server, so pass a free_port() per run: a port shared by two concurrent runs is one server identity two models contend for.

Returns

Leaderboard

The bake-off leaderboard, with the trained arm ranked among the baselines.

train.register()

Register the fused MLX directory mlx_path under the name family in root.

Usage

Source

train.register(name, mlx_path, metadata, *, root)

The registry takes ownership of the weights: mlx_path is copied into the version directory as MODEL_DIR and frozen read-only, so the registered artifact is the registry’s own and a later run of the family cannot mutate it. model_path(version) resolves it; metadata["source_mlx_path"] records the run directory it was fused in. Registration never promotes.

Parameters

name: str

The artifact family.

mlx_path: Path

The fused standalone MLX model directory to copy in.

metadata: Mapping[str, object]

The version’s provenance — backend, method, metric, cost.

root: Path

The registry root; athome.train uses [train].registry_root.

Returns

VersionInfo

The registered version.

train.model_path()

The registry’s own copy of version’s fused MLX model — the artifact to serve.

Usage

Source

train.model_path(version)

train.select()

Pick the backend that runs spec, preferring tinker, then local, then modal.

Usage

Source

train.select(spec, settings)

An override — spec.backend first, else settings.backend — pins the choice; otherwise the first backend that is ~TrainBackend.available() and ~TrainBackend.supports() the spec’s method wins. Only the winner is constructed.

Parameters

spec: TrainSpec

The fine-tuning request; its method and backend drive the choice.

settings: TrainSettings

The [train] settings, whose backend is the fallback override.

Returns

TrainBackend

The selected backend, constructed from its settings section.

Raises

NoBackendAvailable

The override cannot run the spec, or nothing qualifies.

train.lora_keys()

The modules an adapter actually wraps: target_modules filtered by the LoRA toggles.

Usage

Source

train.lora_keys(lora)

Every backend converges on a fused mlx-lm artifact, so every backend is bound by what mlx-lm can express — which makes this the one shape rule, not a per-backend one.

Parameters

lora: LoraSpec

The requested LoRA shape.

Returns

tuple[str, …]

The subset of target_modules the toggles select, in the spec’s order.

Raises

UnsupportedLoraShape

train_unembed is set, or the toggles leave nothing to train. mlx-lm reaches the unembedding through a base-dependent module path — lm_head, which a base that ties its embeddings does not have at all — and BaseModelSpec does not carry it, so no backend can fuse an unembedding LoRA: training one only burns money on weights the PEFT-to-MLX conversion would drop.

train.std_lora_keys()

The standard modules lora’s toggles select, ignoring its target_modules.

Usage

Source

train.std_lora_keys(lora)

What a backend that takes only toggles — Tinker — trains for this spec.

train.spend_cap()

The dollar cap binding this run: the spec’s max_usd when it set one, else default.

Usage

Source

train.spend_cap(spec, default)

Only None means “unset”. max_usd=0.0 is an explicit “spend nothing” and binds as such — a falsy-zero fallback here would silently authorize the whole configured cap.

train.NoBackendAvailable

Raised when no backend can run the request: none is available, or the override cannot.

Usage

Source

train.NoBackendAvailable()

train.UnservableBase

Raised when the spec’s base model cannot produce a servable fused MLX artifact.

Usage

Source

train.UnservableBase()

train.UnsupportedLoraShape

Raised when a LoraSpec asks for an adapter the MLX fuse path cannot express.

Usage

Source

train.UnsupportedLoraShape()

train.InsufficientData

Raised when the surviving training pool holds fewer examples than one batch.

Usage

Source

train.InsufficientData(count, batch_size)

Carries the surviving count so the caller can report how far short the pool fell. It fires before the spend guard reserves anything and before any training client is constructed, so an under-filled run aborts having spent nothing.

train.OverlongEvalRows

Raised when any eval row exceeds max_seq_len; the whole frozen set fails up-front.

Usage

Source

train.OverlongEvalRows(count, max_seq_len)

Carries how many rows overran, checked before any billable call so an ill-sized eval set never rides a snapshot’s prefill.