mixle.models.train_search module

Design-of-experiments helpers for language-model training recipes.

Training a language model is an expensive-objective, low-fidelity-proxy setting: a short run over fewer steps or a data subset is a noisy estimate of the full run’s loss. tune_training() wraps mixle.doe.multi_fidelity_minimize so the search uses low-budget runs to locate promising recipes and reserves full-budget runs to refine them.

The objective is a caller-supplied training callback train(recipe, budget) -> held-out loss where budget in (0, 1] is the fraction of full training. lm_train_fn() provides a callback for LM, and extrapolate_learning_curve() predicts full-budget loss from a partial run’s curve for early stopping.

class TrainingSpace(d_model_choices=(64, 128, 256, 512), n_layer_range=(2, 12), log10_lr_range=(-4.0, -2.0), batch_choices=(16, 32, 64, 128))[source]

Bases: object

The tunable axes of an LM training recipe and how a unit-cube point decodes into concrete knobs.

Parameters:
dims()[source]

Return the dimensionality of the unit-cube recipe search space.

Return type:

int

bounds()[source]

Return unit-cube bounds for the DOE optimizer.

Return type:

list[tuple[float, float]]

decode(point)[source]

Decode a unit-cube point into concrete LM training hyperparameters.

Parameters:

point (ndarray)

Return type:

dict[str, Any]

class TrainingSearchResult(recipe, loss, history=None, seed=None)[source]

Bases: object

The outcome of a multi-fidelity training search: the best recipe, its full-budget loss, and the history.

The result OWNS what it reports (MXR-080-1890). It used to store the caller’s recipe mapping and history container by reference and validate nothing, so:

recipe = {"d_model": 64}
res = TrainingSearchResult(recipe=recipe, loss=1.0, history=trace)
recipe["d_model"] = 999          # the recorded winning recipe silently became a different one

rewrote a decision that had already been recorded, and loss accepted nan or the string "not a number" as a search outcome. __post_init__ now copies both containers with detach_receipt_container() and checks that the fields are the kind of thing the search actually produces.

seed records the search’s common-random-numbers seed, so the reported loss is attributable to a reproducible run rather than to one unlabelled draw.

Deliberately NOT checked, and why:

  • The numeric contents of ``history``. detach_receipt_container copies containers but passes elements through by identity, so the X/Y blocks a search accumulates stay the same arrays rather than being duplicated – they are the optimizer’s own freshly built output, nobody else holds them, and copying every evaluation block would add a real cost to a result that is often only glanced at. A holder can still edit those arrays in place; this defends against the caller who kept a reference, which is the aliasing defect that was reproduced.

  • That ``recipe`` decodes to ``history[“x”]``. Inside tune_training() both come from the same optimizer result, so the check would be tautological; for a directly constructed result there is no space to decode against.

Parameters:
tune_training(train, space=None, *, fidelities=(0.25, 1.0), costs=None, max_cost=20.0, n_init=None, seed=0)[source]

Run multi-fidelity BO over a training recipe.

train(recipe, budget) returns held-out loss, where lower is better. fidelities are the training-budget fractions the search may run at. Returns the recipe with the best full-budget loss and the full BO history.

seed is recorded on the returned result (MXR-080-1890) so the reported loss is attributable to a reproducible search rather than to one unlabelled draw. Whether the objective is itself reproducible is the callback’s contract, not this function’s: lm_train_fn() supplies one that pins its own randomness, but an arbitrary caller-supplied train is opaque here and is deliberately not second-guessed.

Parameters:
Return type:

TrainingSearchResult

class LMFidelity(budget, epochs, tokens_per_epoch, total_tokens, seed)[source]

Bases: object

The work one declared budget actually buys, counted in TRAINING TOKENS (MXR-080-1890).

Fidelity used to be max(1, round(max_epochs * budget)) whole epochs, which makes the declared budget and the executed work different quantities related by a coarse many-to-one map. Measured over the library’s own default max_epochs=3, budgets 0.05, 0.1, 0.2, 0.25, 0.3, 0.4 all executed exactly ONE epoch, and at max_epochs=1 every budget from 0.05 to 1.0 executed one epoch – so a “cheap screening run” and a “full run” were byte-identical work and the multi-fidelity premise did not hold at all.

Counting in tokens restores resolution: the corpus is truncated to a prefix as well as repeated for whole epochs, so total_tokens separates budgets that whole epochs could not. epochs and tokens_per_epoch are what is handed to fit(); total_tokens = epochs * tokens_per_epoch is the comparable work figure. seed is the common-random-numbers seed the run used.

Parameters:
class LMTrainFn(train_ids, validation_ids, *, vocab, block, max_epochs, device, seed)[source]

Bases: object

A (recipe, budget) -> held-out nats/token callback for LM that records what it ran.

Built by lm_train_fn(). Two properties the bare closure did not have (MXR-080-1890):

  • Fidelity is inspectable before it is paid for. plan() resolves a budget into the concrete LMFidelity it will execute, so a caller can see whether two declared budgets are actually different runs without training anything.

  • Every executed run is recorded. trials returns the fidelity (work AND seed) of each completed call, so a finished search can be audited for budgets that collapsed onto the same work rather than the collapse being invisible.

Parameters:
property seed: int

The common-random-numbers seed every run uses (see __call__()).

property trials: tuple[LMFidelity, ...]

The fidelity of every completed run, oldest first, as an immutable snapshot.

plan(budget)[source]

Resolve budget in (0, 1] into the concrete work it will execute, without running it.

Parameters:

budget (float)

Return type:

LMFidelity

lm_train_fn(token_ids, val_ids, *, vocab, block=64, max_epochs=3, device='cpu', seed=0, fidelities=None)[source]

Return a training callback (recipe, budget) -> held-out nats/token for LM.

budget in (0, 1] is a fraction of the FULL run’s training tokens (max_epochs * len(token_ids)), not a fraction of the epoch count. budget=1.0 is unchanged – max_epochs epochs over the whole corpus – while a smaller budget shortens the corpus prefix as well as the epoch count, so distinct budgets buy distinct work. See LMFidelity for the measurement that motivated this (MXR-080-1890).

seed is the common-random-numbers seed shared by every run, which is what makes two recipes comparable; see LMTrainFn.__call__().

fidelities is optional and off by default. Pass the exact budget ladder the search will use and it is checked, before any training happens, for budgets that still resolve to identical work – the failure this finding is about. Token resolution runs out on a short corpus, so this can legitimately happen; it is reported rather than assumed away.

Raises:

ValueError – If fidelities is supplied and two of its budgets execute the same work.

Parameters:
Return type:

LMTrainFn

extrapolate_learning_curve(steps, losses, *, at)[source]

Predict the loss at budget/step at from a partial run’s (steps, losses) via a power-law fit.

Fits loss(t) = a + b * t^(-c) and evaluates it at at so a partial run can estimate the full-budget loss for early stopping. Invalid data, an underdetermined curve, or a failed/non-finite fit raises explicitly.

Parameters:
Return type:

float