mixle.models.train_search module

DoE for training LLMs – multi-fidelity Bayesian optimization of the training recipe, budget as the fidelity.

Training a language model is the expensive-objective, cheap-proxy setting multi-fidelity BO is built for: a short run (few steps / a data subset) is a noisy, cheap estimate of the full run’s loss. tune_training() wraps mixle.doe.multi_fidelity_minimize so the search spends cheap low-budget runs to locate good recipes and reserves full-budget runs to refine – reaching a good recipe for a fraction of the compute of grid search.

The objective is your training callback train(recipe, budget) -> held-out loss (budget in (0, 1] is the fraction of full training), so it scales from a tiny local model to a real pretraining loop unchanged – lm_train_fn() gives a ready callback that trains a mixle LM. extrapolate_learning_curve() predicts a 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:
d_model_choices: Sequence[int] = (64, 128, 256, 512)
n_layer_range: tuple[int, int] = (2, 12)
log10_lr_range: tuple[float, float] = (-4.0, -2.0)
batch_choices: Sequence[int] = (16, 32, 64, 128)
dims()[source]
Return type:

int

bounds()[source]
Return type:

list[tuple[float, float]]

decode(point)[source]
Parameters:

point (ndarray)

Return type:

dict[str, Any]

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

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 = None
tune_training(train, space=None, *, fidelities=(0.25, 1.0), costs=None, max_cost=20.0, n_init=None, seed=0)[source]

Multi-fidelity BO of the training recipe; train(recipe, budget) returns held-out loss (lower is better).

fidelities are the training-budget fractions the search may run at (cheap first). Returns the recipe with the best full-budget loss and the full BO history. Scale the objective by swapping train for a real launcher – the search machinery is identical.

Parameters:
Return type:

TrainingSearchResult

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

A ready training callback (recipe, budget) -> held-out nats/token that trains a real mixle LM.

budget in (0, 1] scales the number of epochs (the cheap-fidelity axis is training length); a real pretraining loop would scale steps or the token subset the same way.

Parameters:
Return type:

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

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) (the standard learning-curve form) and evaluates it at at – so a cheap partial run can estimate the full-budget loss for early stopping. Falls back to the last observed loss if the fit fails.

Parameters:
Return type:

float