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))

Bases: object

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

Parameters:

• d_model_choices (Sequence[int])

• n_layer_range (tuple[int, int])

• log10_lr_range (tuple[float, float])

• batch_choices (Sequence[int])

dims()

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

Return type:

int

bounds()

Return unit-cube bounds for the DOE optimizer.

Return type:

list[tuple[float, float]]

decode(point)

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

Parameters:

point (ndarray)

Return type:

dict[str, Any]

class TrainingSearchResult(recipe, loss, history=None)

Bases: object

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

Parameters:

• recipe (dict[str, Any])

• loss (float)

• history (Any)

tune_training(train, space=None, *, fidelities=(0.25, 1.0), costs=None, max_cost=20.0, n_init=None, seed=0)

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.

Parameters:

• train (Callable[[dict[str, Any], float], float])

• space (TrainingSpace | None)

• fidelities (tuple[float, ...])

• costs (tuple[float, ...] | None)

• max_cost (float)

• n_init (int | None)

• seed (int)

Return type:

TrainingSearchResult

lm_train_fn(token_ids, val_ids, *, vocab, block=64, max_epochs=3, device='cpu')

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

budget in (0, 1] scales the number of epochs. A larger pretraining loop can use the same convention to scale steps or token subsets.

Parameters:

• token_ids (Sequence[int])

• val_ids (Sequence[int])

• vocab (int)

• block (int)

• max_epochs (int)

• device (str)

Return type:

Callable[[dict[str, Any], float], float]

extrapolate_learning_curve(steps, losses, *, at)

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. Falls back to the last observed loss if the fit fails.

Parameters:

• steps (Sequence[float])

• losses (Sequence[float])

• at (float)

Return type:

float