Training

Learner

XPM Taskxpm_torch.learner.Learner(*, random, trainer, model, max_epochs, steps_per_epoch, optimizers, listeners, checkpoint_interval, hooks)[source]

Model Learner

The learner task is generic, and takes two main arguments: (1) the model defines the model (e.g. DRMM), and (2) the trainer defines how the model should be trained (e.g. pointwise, pairwise, etc.)

When submitted, it returns a dictionary based on the listeners

random: xpm_torch.base.Random: The random generator

trainer: xpm_torch.trainers.Trainer: Specifies how to train the model

model: xpm_torch.module.Module: Defines the model to be learned. If multiple models are used, one can use MultipleModel.

max_epochs: int = 1000: Maximum number of epochs

steps_per_epoch: int = 128: Number of steps for one epoch (after each epoch results are logged)

optimizers: List[xpm_torch.optim.ParameterOptimizer]: The list of parameter optimizers

listeners: List[xpm_torch.learner.LearnerListener]: Listeners are in charge of handling the validation of the model, and saving the relevant checkpoints

checkpoint_interval: int = 1: Number of epochs between each checkpoint

logpath: pathgenerated: The path to the tensorboard logs

checkpointspath: pathgenerated: The path to the checkpoints

hooks: List[xpm_torch.context.Hook] = []: Global learning hooks Initialization hooks are called before and after the initialization of the trainer and listeners.

fabric_config: xpm_torch.configuration.FabricConfiguration = xpm_torch.configuration.FabricConfiguration(precision=32-true, torch_fp32_precision=None, num_nodes=1, devices=auto, strategy=auto, accelerator=auto)generated: Runtime configuration, managed by Fabric

XPM Configxpm_torch.learner.LearnerListener(*, id)[source]

Hook for learner

Performs some operations after a learning epoch

id: str: Unique ID to identify the listener (ignored for signature)

XPM Configxpm_torch.learner.CheckpointSettings(*, epoch)[source]

Settings for a checkpoint-specific ModuleLoader.

epoch: int: The epoch of the checkpoint

XPM Configxpm_torch.learner.Strategy[source]: A Lightning strategy

Trainers

XPM Configxpm_torch.trainers.Trainer(*, hooks, model)[source]

Generic trainer

hooks: List[xpm_torch.trainers.context.TrainingHook] = []: Hooks for this trainer: this includes the losses, but can be adapted for other uses The specific list of hooks depends on the specific trainer

model: xpm_torch.module.Module: If the model to optimize is different from the model passsed to Learn, this parameter can be used – initialization is still expected to be done at the learner level

XPM Configxpm_torch.trainers.LossTrainer(*, hooks, model, sampler, batch_size, num_workers)[source]

Trainer based on a loss function

Uses StatefulDataLoader + IterableDataset for data loading.

hooks: List[xpm_torch.trainers.context.TrainingHook] = []: Hooks for this trainer: this includes the losses, but can be adapted for other uses The specific list of hooks depends on the specific trainer

model: xpm_torch.module.Module: If the model to optimize is different from the model passsed to Learn, this parameter can be used – initialization is still expected to be done at the learner level

batcher: xpm_torch.batchers.Batchergenerated: How to batch samples together

sampler: xpm_torch.base.Sampler: The sampler to use

batch_size: int = 16: Number of samples per batch

num_workers: int = 2: Number of DataLoader workers

XPM Configxpm_torch.trainers.multiple.MultipleTrainer(*, hooks, model, trainers)[source]

This trainer can be used to combine various trainers

hooks: List[xpm_torch.trainers.context.TrainingHook] = []: Hooks for this trainer: this includes the losses, but can be adapted for other uses The specific list of hooks depends on the specific trainer

model: xpm_torch.module.Module: If the model to optimize is different from the model passsed to Learn, this parameter can be used – initialization is still expected to be done at the learner level

trainers: Dict[str, xpm_torch.trainers.Trainer]: The trainers

Training State

class xpm_torch.trainers.context.TrainState(model: Module, trainer: Trainer, optimizer: ScheduledOptimizer, epoch=0, steps=0)[source]

Represents a training state for serialization

epoch: int: The epoch

load(path, onlyinfo=False)[source]: Loads the state from disk

save(path)[source]: Save the state

property step: Returns the step for logging (number of steps)

steps: int: The number of steps (each epoch is composed of sptes)

class xpm_torch.trainers.context.TrainerContext(logpath: Path, path: Path, max_epoch: int, steps_per_epoch: int, trainer, model: Module, optimizer: ScheduledOptimizer, fabric: lightning.Fabric)[source]

Contains all the information about the training context for a spefic

This object is used when training to provide models and losses’ with extra information - as well as the possibility to add regularization losses

add_metric(metric: Metric)[source]

add a metric to be reported at the end of the step (e.g., for logging in tensorboard)

Parameters:: metric (Metric) – The metric to be added

copy(path: Path)[source]: Copy the state into another folder

load_bestcheckpoint(max_epoch: int)[source]: Try to find the best checkpoint to load (the highest lower than the epoch target)

metrics: Metrics | None: Metrics to be reported

property writer

Returns a tensorboard writer

by default, purges the entries beside the current epoch

Training Hooks

XPM Configxpm_torch.trainers.context.TrainingHook[source]: Base class for all training hooks

XPM Configxpm_torch.trainers.context.StepTrainingHook[source]: Base class for hooks called at each step (before/after)

XPM Configxpm_torch.trainers.context.ValidationHook[source]: Base class for all the validation hooks

XPM Configxpm_torch.trainers.context.InitializationTrainingHook[source]: Base class for hooks called at initialization

XPM Configxpm_torch.trainers.hooks.LayerFreezer(*, selector)[source]

This training hook class can be used to freeze a subset of model parameters

selector: xpm_torch.parameters.ParametersIterator: How to select the layers to freeze

XPM Configxpm_torch.trainers.hooks.LayerSharer(*, source, target)[source]

This training hook class can be used to share parameters

source: xpm_torch.parameters.ParametersIterator: The parameters to share

target: xpm_torch.parameters.ParametersIterator: The parameters to be shared

Validation

XPM Configxpm_torch.validation.ValidationSettings(*, listener, key)[source]

Settings for a validation-specific ModuleLoader.

Attached as settings on the loader to distinguish validation checkpoints from other loaders with the same model and path.

listener: xpm_torch.learner.LearnerListener: The listener (kept to change the loader identifier based on the learner listener configuration)

key: str: The metric key for this validation checkpoint

XPM Configxpm_torch.trainers.validation.TrainerValidationLoss(*, id, data, batch_size, trainer, warmup, validation_interval, early_stop)[source]

Generic trainer-based loss validation

id: str: Unique ID to identify the listener (ignored for signature)

data: xpm_torch.base.SampleIterator: The dataset to use

batcher: xpm_torch.batchers.Batchergenerated: How to batch samples together

batch_size: int: Batch size

trainer: xpm_torch.trainers.LossTrainer: The trainer

warmup: int = -1: How many epochs before actually computing the validation loss

bestpath: pathgenerated: Path to the best checkpoints

info: pathgenerated: Path to the JSON file that contains the metric values at each epoch

validation_interval: int = 1: Epochs between each validation

early_stop: int = 0: Number of epochs without improvement after which we stop learning. Should be a multiple of validation_interval or 0 (no early stopping)

Batching

A Batcher controls how an outer batch is split into microbatches before each forward/backward. Three strategies are available; the right choice depends on whether you train on a single GPU or with multi-GPU (DDP), and on whether your per-batch memory cost varies (e.g. variable sequence length in transformers).

Batcher — base, no-op. Uses the configured batch_size directly, no microbatching, no fallback. Good when you have already manually picked a microbatch size that always fits.
PowerAdaptativeBatcher — single-GPU only. Catches CUDA OOM errors and shrinks the microbatch (batch_size // n for increasing n) until the call succeeds. Do not use under DDP: when one rank shrinks its microbatch and the others don’t, the ranks desynchronize at the next NCCL collective and the job hangs (see https://github.com/pytorch/pytorch/issues/50820).
PredictiveBatcher — fits a user-supplied memory cost formula (e.g. transformer attention: a*bs*l**2 + b*bs*l + c) from a small set of probe batches run before DDP wrapping, then chooses the largest safe microbatch per outer batch using the fitted model. Avoids OOM-driven recovery entirely, so DDP cannot desynchronize from per-rank shrinkage. Recommended whenever you train on more than one GPU.

XPM Configxpm_torch.batchers.Batcher[source]

Responsible for micro-batching when the batch does not fit in memory

The base class just does nothing (no adaptation)

XPM Configxpm_torch.batchers.PowerAdaptativeBatcher[source]: Starts with the provided batch size, and then divides in 2, 3, etc. until there is no more OOM

XPM Configxpm_torch.batchers.PredictiveBatcher(*, formula, coefficient_names, variable_names, batch_size_variable, memory_fraction, extra_overhead_bytes, min_batch_size, dims_provider)[source]

Predictive batcher.

Fits a user-supplied memory cost formula linear in coefficients (e.g. “a*bs*l**2 + b*bs*l + c”) from a small set of probes run before DDP wrapping, then chooses the largest microbatch that fits the remaining GPU memory budget at train time. Avoids the multi-GPU NCCL hang that an OOM-reactive batcher causes when one rank shrinks its microbatch and desynchronizes from the others.

Important: the trainer must invoke pre_train_setup on the worker before the model is wrapped with fabric.setup(model). LossTrainer does this in its initialize().

formula: str: Cost formula, linear in coefficients. Variables and coefficients must be declared via variable_names and coefficient_names. Canonicalized in __validate__ so equivalent inputs share an identifier.

coefficient_names: set[str]: Names of the unknown coefficients in formula, fit via least squares.

variable_names: set[str]: Names of the input variables in formula (e.g. {“bs”, “l”}).

batch_size_variable: str = bs: Which variable in variable_names is the microbatch size to maximize.

memory_fraction: float = 0.75: Safety margin on free GPU memory. Default is conservative to absorb DDP gradient buckets / NCCL buffers that the pre-DDP probe doesn’t measure.

extra_overhead_bytes: int = 0: Additional bytes to subtract from the budget at train time, in case memory_fraction alone isn’t enough on a given setup.

min_batch_size: int = 1: Floor on the chosen microbatch size.

dims_provider: xpm_torch.batchers.BatchDimsProvider: How to read dims from a runtime batch and how to build a probe batch.

packer: xpm_torch.batchers.MicroBatchPacker = xpm_torch.batchers.UniformPacker()generated: How to split an outer batch into microbatches. Default UniformPacker matches conservative slicing on outer-batch dims; use AdaptivePacker for heterogeneous-cost batches (e.g. variable sequence length).

XPM Configxpm_torch.batchers.BatchDimsProvider(*, probe_dims)[source]

Contract between a PredictiveBatcher and its surrounding trainer.

Implementations are responsible for:

reading the formula’s input variables from a runtime batch
constructing a representative probe batch matching given dims

probe_dims lives here (and not on the batcher) because probes are intrinsically tied to the data shape the provider knows how to construct.

probe_dims: set[xpm_torch.batchers.ProbePoint]

XPM Configxpm_torch.batchers.ProbePoint(*, values)[source]

One probe configuration: a mapping {variable_name -> value} for a single point at which to measure peak GPU memory during the probe phase.

values: dict[str, int]

Microbatch packing

PredictiveBatcher delegates the “how to split an outer batch into microbatches” decision to a swappable MicroBatchPacker. Two packers ship with xpm-torch:

UniformPacker (default) — compute dims once from the outer batch, choose a single microbatch size, slice equally. Optimal when records are homogeneous; conservative when they aren’t.
AdaptivePacker — accumulate records into a microbatch incrementally, emit when adding one more would exceed the budget. Optionally sort by a per-record cost proxy descending (sort_by) so heavy records go in one early microbatch and light records pack densely into later ones.

Use AdaptivePacker when records have heterogeneous cost (e.g. variable sequence length) — each microbatch is then sized for its actual contents rather than the worst-case record in the outer batch.

from xpm_torch.batchers import AdaptivePacker, PredictiveBatcher

PredictiveBatcher.C(
    formula="a*bs*l**2 + b*bs*l + c",
    coefficient_names={"a", "b", "c"},
    variable_names={"bs", "l"},
    dims_provider=TextBatchDimsProvider.C(probe_dims=...),
    packer=AdaptivePacker.C(sort_by="l"),
)

XPM Configxpm_torch.batchers.MicroBatchPacker[source]

Strategy for splitting an outer batch into microbatches under a PredictiveBatcherWorker.

Different implementations trade off simplicity against packing density. Plug into a PredictiveBatcher via its packer Param.

XPM Configxpm_torch.batchers.UniformPacker[source]

Default packer: compute dims once from the outer batch, choose a single microbatch size, slice equally.

Conservative: each microbatch is sized as if it contained the worst-case-cost record from the outer batch. For homogeneous batches (all records same shape) this is optimal; for heterogeneous batches AdaptivePacker packs denser.

XPM Configxpm_torch.batchers.AdaptivePacker(*, sort_by)[source]

Pack records greedily: incrementally accumulate into a microbatch and emit when adding one more record would exceed the memory budget.

Two improvements over UniformPacker:

Each microbatch is sized for its own dims (recomputed per-microbatch), so chunks containing only short-cost records can be much larger.
Optionally sort by a per-record cost proxy (sort_by) descending so heavy records go into one early microbatch and light records pack densely into later ones.

sort_by: str: Variable name (must be produced by compute_batch_dims([record])) to sort records by, descending. Reduces microbatch count for batches with heterogeneous per-record cost. Default: no sort (preserve sampler order).

Predictive batching: worked example

For a transformer-style model where memory scales with batch size and the square of sequence length, you would:

Write a BatchDimsProvider for your record type, exposing the batch size and max sequence length and the set of probe points.
Plug it into a PredictiveBatcher and pass that to your trainer’s batcher parameter.

from experimaestro import Param, sealed_set
from xpm_torch.batchers import (
    BatchDimsProvider,
    PredictiveBatcher,
    ProbePoint,
)


class TextBatchDimsProvider(BatchDimsProvider):
    """Reads (bs, l) from a list-of-records batch where each record has a
    ``tokens`` attribute."""

    def compute_batch_dims(self, batch) -> dict[str, int]:
        return {
            "bs": len(batch),
            "l": max(len(r.tokens) for r in batch),
        }

    def build_probe_batch(self, dims: dict[str, int]):
        # Return a synthetic batch matching the requested dims. Replace
        # this with your project's record constructor.
        return [DummyRecord(tokens=[0] * dims["l"]) for _ in range(dims["bs"])]


trainer = LossTrainer.C(
    batcher=PredictiveBatcher.C(
        formula="a*bs*l**2 + b*bs*l + c",
        coefficient_names={"a", "b", "c"},
        variable_names={"bs", "l"},
        batch_size_variable="bs",
        dims_provider=TextBatchDimsProvider.C(
            probe_dims=sealed_set(
                ProbePoint.C(values={"bs": 1, "l": 64}),
                ProbePoint.C(values={"bs": 1, "l": 128}),
                ProbePoint.C(values={"bs": 2, "l": 128}),
                ProbePoint.C(values={"bs": 4, "l": 64}),
                ProbePoint.C(values={"bs": 2, "l": 256}),
            ),
        ),
    ),
    batch_size=32,
    sampler=...,
)

Notes:

The formula is canonicalized in __validate__ (terms sorted, whitespace normalized) so equivalent inputs share an experimaestro task identifier.
Probes run on rank 0 only; the fitted coefficients are broadcast to other ranks. This is safe because probing happens before fabric.setup(model) — a rank-0-only forward+backward on a DDP-wrapped model would deadlock at the first gradient bucket all-reduce.
The default memory_fraction=0.75 is conservative to absorb DDP gradient buckets and NCCL communication buffers that the pre-DDP probe doesn’t see. If you still OOM at train time, prefer lowering memory_fraction or setting extra_overhead_bytes over silently retrying — silent retry re-introduces the very NCCL desync risk this batcher exists to avoid.

Batchwise Losses

XPM Configxpm_torch.losses.batchwise.BatchwiseLoss(*, weight)[source]

weight: float = 1.0: The weight of this loss

XPM Configxpm_torch.losses.batchwise.CrossEntropyLoss(*, weight)[source]

weight: float = 1.0: The weight of this loss

XPM Configxpm_torch.losses.batchwise.SoftmaxCrossEntropy(*, weight)[source]

weight: float = 1.0: The weight of this loss

Pairwise Losses

XPM Configxpm_torch.losses.pairwise.PairwiseLoss(*, weight)[source]

Base class for any pairwise loss

weight: float = 1.0: The weight \(w\) with which the loss is multiplied (useful when combining with other ones)

XPM Configxpm_torch.losses.pairwise.CrossEntropyLoss(*, weight)[source]

Cross-Entropy Loss

Computes the cross-entropy loss

Classification loss (relevant vs non-relevant) where the logit is equal to the difference between the relevant and the non relevant document (or equivalently, softmax then mean log probability of relevant documents) Reference: C. Burges et al., “Learning to rank using gradient descent,” 2005.

warning: this loss assumes the score returned by the scorer is a logit

\[\frac{w}{N} \sum_{(s^+,s-)} \log \frac{\exp(s^+)}{\exp(s^+)+\exp(s^-)}\]

weight: float = 1.0: The weight \(w\) with which the loss is multiplied (useful when combining with other ones)

XPM Configxpm_torch.losses.pairwise.HingeLoss(*, weight, margin)[source]

Hinge (or max-margin) loss

\[\frac{w}{N} \sum_{(s^+,s-)} \max(0, m - (s^+ - s^-))\]

weight: float = 1.0: The weight \(w\) with which the loss is multiplied (useful when combining with other ones)

margin: float = 1.0: The margin for the Hinge loss

XPM Configxpm_torch.losses.pairwise.PointwiseCrossEntropyLoss(*, weight)[source]

Point-wise cross-entropy loss

This is a point-wise loss adapted as a pairwise one.

This loss adapts to the ranker output type:

If real, uses a BCELossWithLogits (sigmoid transformation)
If probability, uses the BCELoss
If log probability, uses a BCEWithLogLoss

\[\frac{w}{2N} \sum_{(s^+,s-)} \log \frac{\exp(s^+)}{\exp(s^+)+\exp(s^-)} + \log \frac{\exp(s^-)}{\exp(s^+)+\exp(s^-)}\]

weight: float = 1.0: The weight \(w\) with which the loss is multiplied (useful when combining with other ones)

Fabric Configuration

XPM Configxpm_torch.configuration.Strategy[source]

XPM Configxpm_torch.configuration.FabricConfigurationBase[source]

XPM Configxpm_torch.configuration.FabricConfiguration(*, precision, torch_fp32_precision, num_nodes, devices, strategy, accelerator)[source]

Describe the computation device

The backend is fabric, so the complete documentation can be found on https://lightning.ai/docs/fabric/stable/api/fabric_args.html

precision: str = 32-true: Precision to use, e.g., ‘16-mixed’, ‘bf16-mixed’, ‘32-true’: see Lightning documentation at https://lightning.ai/docs/fabric/stable/api/fabric_args.html#precision

torch_fp32_precision: str: Torch precision for torch.float32 operations, see https://pytorch.org/docs/stable/generated/torch.set_float32_matmul_precision.html#torch.set_float32_matmul_precision Automatically set depending on fabric_config.precision if not set, but can be overridden if needed (e.g., to force TF32 on Ampere GPUs while using bf16 precision for other operations)

num_nodes: int = 1: Number of nodes

devices: str = auto: Configure the devices to run on. See https://lightning.ai/docs/fabric/stable/api/fabric_args.html#devices for more details and options. Note that for multi-node training, you should specify the devices per node, e.g., devices=”4” for 4 GPUs per node, not devices=”16” for a total of 16 GPUs across 4 nodes.

strategy: str = auto: The strategy to use See https://lightning.ai/docs/fabric/stable/api/fabric_args.html#strategy for more details and options.

accelerator: str = auto: The accelerator to use See https://lightning.ai/docs/fabric/stable/api/fabric_args.html#accelerator for more details and options.