transformer_lens.tools.training module

Train loop for TransformerLens models.

Utilities for training models on autoregressive language modeling tasks. Typed against the model_protocol surface (__call__ with return_type="loss" plus standard nn.Module parameter access), so any conforming model — TransformerBridge foremost — works through this loop.

class transformer_lens.tools.training.TrainConfig(num_epochs: int, batch_size: int, lr: float = 0.001, seed: int = 0, momentum: float = 0.0, max_grad_norm: float | None = None, weight_decay: float | None = None, optimizer_name: str = 'Adam', device: str | device | None = None, warmup_steps: int = 0, save_every: int | None = None, save_dir: str | None = None, wandb: bool = False, wandb_project_name: str | None = None, print_every: int | None = 50, max_steps: int | None = None)

Bases: object

Configuration class to store training hyperparameters for a training run.

Parameters:

  • num_epochs (int) – Number of epochs to train for

  • batch_size (int) – Size of batches to use for training

  • lr (float) – Learning rate to use for training

  • seed (int) – Random seed to use for training

  • momentum (float) – Momentum to use for training

  • max_grad_norm (float, optional) – Maximum gradient norm to use for

  • weight_decay (float, optional) – Weight decay to use for training

  • optimizer_name (str) – The name of the optimizer to use

  • device (str or torch.device, optional) – Device to use for training

  • warmup_steps (int, optional) – Number of warmup steps to use for training

  • save_every (int, optional) – After how many batches should a checkpoint be saved

  • save_dir (str, optional) – Where to save checkpoints

:param : Where to save checkpoints :type : str, optional :param wandb: Whether to use Weights and Biases for logging :type wandb: bool :param wandb_project: Name of the Weights and Biases project to use :type wandb_project: str, optional :param print_every: Print the loss every n steps :type print_every: int, optional :param max_steps: Terminate the epoch after this many steps. Used for debugging. :type max_steps: int, optional

batch_size: int
device: str | device | None = None
lr: float = 0.001
max_grad_norm: float | None = None
max_steps: int | None = None
momentum: float = 0.0
num_epochs: int
optimizer_name: str = 'Adam'
print_every: int | None = 50
save_dir: str | None = None
save_every: int | None = None
seed: int = 0
wandb: bool = False
wandb_project_name: str | None = None
warmup_steps: int = 0
weight_decay: float | None = None
transformer_lens.tools.training.train(model: TrainableTransformerLensModel, config: TrainConfig, dataset: Dataset) → TrainableTransformerLensModel

Train a model on an autoregressive language modeling task.

Parameters:

  • model – The model to train (TrainableTransformerLensModel: callable with return_type="loss" and exposing torch parameters)

  • config – The training configuration

  • dataset – The dataset to train on - assumed set up for autoregressive language modeling.

Returns:

The trained model