# Training

> UForm for Python

Training is GPU-only.
The configuration arrives as keyword-built structures mirroring the C ones field by field, so a field left out is zero, which the library derives:

```python
executor = uform.Executor.for_vram(devices=[0, 1])
mesh = uform.Mesh.bootstrap(executor, uform.MeshOptions(host_domain_index=rank, host_domain_count=nodes,
                                                        directory="/shared/exchange", job="run-17"))
trainer = uform.Trainer.init_vram(executor, mesh=mesh, depth=12, hidden=512, objective=uform.Objective.JOINT,
                                  batch=uform.BatchShape(max_rows_per_device=64, max_segments_per_device=256,
                                                         max_columns_per_row=1024, max_bags_per_device=128,
                                                         image=uform.ImageShape(patch_side=16)),
                                  schedule=uform.Schedule(steps=100_000), seed=42)
stats = uform.TrainerStepStats()
trainer.step(modalities, contents, bags=bags, relations=uform.Relations(attract=(offsets, ids)), stats=stats)
trainer.save("checkpoint.safetensors")
```

`BatchShape` nests an `ImageShape`, a `VideoShape` and an `AudioShape`, each also taken as a tuple of its fields in order.
`Masking`, `Contrastive`, `LateInteraction`, `Choices` and `Schedule` configure the objectives, and `Relations` holds three `Edges` sets of buffers: `size_t` offsets, `uint32` IDs and optional `float32` weights.
`Contrastive` and `LateInteraction` take a `uform.PairLossForm`, and `trainer.step` takes each item's `uform.LateInteractionSide` as `sides`.
`choices`, taken as `text_choose` takes them, and `choice_targets`, a `float32` tape shaped as the `RowsTape` it answers, train the calibrated choices within the batch's `max_choices_per_item`, so `(probabilities, probabilities.offsets)` from one model can train another.
Paths, like the mesh's `directory`, cross as `os.fsencode` gives their bytes, while its `job` is a `str` name.
`trainer.step` answers the pooled embeddings the loss scored when given an `arena`, and `None` otherwise.
`mesh.describe()` answers a `uform.MeshDescription`, its `multimem_tier` a `uform.MultimemTier` and its `gpu_class` the `uform.GPUClass` every GPU of the run dispatches as.

`import uform`

`from uform import ImageShape, VideoShape, AudioShape, …`

## Modality

What the bytes of one training item hold, valued as `uform_modality_t`.

```python
class Modality(IntEnum)
```

- `TEXT` = `0`
- `IMAGE` = `1`
- `VIDEO` = `2`
- `AUDIO` = `3`

## Objective

Which losses a training step relaxes, valued as `uform_objective_t`.

```python
class Objective(IntEnum)
```

- `MASKED` = `0`
- `CONTRASTIVE` = `1`
- `JOINT` = `2`

## Negatives

Which pairs the contrastive loss pushes apart, valued as `uform_negatives_t`.

```python
class Negatives(IntEnum)
```

- `AMBIENT` = `0`
- `NAMED` = `1`

## PairLossForm

How a contrastive term turns a row of scores into forces, valued as `uform_pair_loss_form_t`.

```python
class PairLossForm(IntEnum)
```

- `DERIVED` = `0`
- `BALANCED` = `1`
- `RANKED` = `2`
- `SIGMOID` = `3`

## LateInteractionSide

Which side of a late-interaction score an item takes, valued as `uform_late_interaction_side_t`.

```python
class LateInteractionSide(IntEnum)
```

- `NONE` = `0`
- `QUERY` = `1`
- `DOCUMENT` = `2`
- `QUERY_AND_DOCUMENT` = `3`

## ImageShape

The image slice of a [`BatchShape`](#BatchShape); a zero [`patch_side`](#patch_side) leaves images out of the run.

```python
class ImageShape
```

## ImageShape

```python
ImageShape(*, patch_side: int = 0, channels: int = 0)
```

## patch_side

```python
patch_side: Final[int]
```

## channels

```python
channels: Final[int]
```

## VideoShape

The video slice of a [`BatchShape`](#BatchShape); a zero [`patch_side`](#patch_side) leaves video out of the run.

```python
class VideoShape
```

## VideoShape

```python
VideoShape(
    *,
    patch_side: int = 0,
    tube_frames: int = 0,
    frame_stride: int = 0,
    channels: int = 0,
)
```

## patch_side

```python
patch_side: Final[int]
```

## tube_frames

```python
tube_frames: Final[int]
```

## frame_stride

```python
frame_stride: Final[int]
```

## channels

```python
channels: Final[int]
```

## AudioShape

The audio slice of a [`BatchShape`](#BatchShape); a zero [`chunk_samples`](#chunk_samples) leaves audio out of the run.

```python
class AudioShape
```

## AudioShape

```python
AudioShape(*, chunk_samples: int = 0, sample_rate: int = 0)
```

## chunk_samples

```python
chunk_samples: Final[int]
```

## sample_rate

```python
sample_rate: Final[int]
```

## BatchShape

What one device's step carries at most, whatever the modality mix, fixed at creation.

```python
class BatchShape
```

## BatchShape

```python
BatchShape(
    *,
    max_rows_per_device: int = 0,
    max_segments_per_device: int = 0,
    max_columns_per_row: int = 0,
    max_bags_per_device: int = 0,
    max_choices_per_item: int = 0,
    image: ImageShape | tuple[int, ...] | None = None,
    video: VideoShape | tuple[int, ...] | None = None,
    audio: AudioShape | tuple[int, ...] | None = None,
)
```

## max_rows_per_device

```python
max_rows_per_device: Final[int]
```

## max_segments_per_device

```python
max_segments_per_device: Final[int]
```

## max_columns_per_row

```python
max_columns_per_row: Final[int]
```

## max_bags_per_device

```python
max_bags_per_device: Final[int]
```

## max_choices_per_item

```python
max_choices_per_item: Final[int]
```

## image

```python
image: Final[ImageShape | None]
```

## video

```python
video: Final[VideoShape | None]
```

## audio

```python
audio: Final[AudioShape | None]
```

## Masking

How the masked objective corrupts a batch and weighs its heads; every zero derives.

```python
class Masking
```

## Masking

```python
Masking(
    *,
    token_mask_probability: float = 0.0,
    token_mean_span: float = 0.0,
    replace_fraction: float = 0.0,
    random_fraction: float = 0.0,
    image_mask_ratio: float = 0.0,
    video_mask_ratio: float = 0.0,
    audio_mask_ratio: float = 0.0,
    text_weight: float = 0.0,
    image_weight: float = 0.0,
    video_weight: float = 0.0,
    audio_weight: float = 0.0,
)
```

## token_mask_probability

```python
token_mask_probability: Final[float]
```

## token_mean_span

```python
token_mean_span: Final[float]
```

## replace_fraction

```python
replace_fraction: Final[float]
```

## random_fraction

```python
random_fraction: Final[float]
```

## image_mask_ratio

```python
image_mask_ratio: Final[float]
```

## video_mask_ratio

```python
video_mask_ratio: Final[float]
```

## audio_mask_ratio

```python
audio_mask_ratio: Final[float]
```

## text_weight

```python
text_weight: Final[float]
```

## image_weight

```python
image_weight: Final[float]
```

## video_weight

```python
video_weight: Final[float]
```

## audio_weight

```python
audio_weight: Final[float]
```

## Contrastive

How the contrastive loss scores the pooled embeddings; every zero derives.

```python
class Contrastive
```

## Contrastive

```python
Contrastive(
    *,
    matryoshka_dimensions: Sequence[int] | None = None,
    temperature_groups: int = 0,
    inverse_temperature: float = 0.0,
    bias: float = 0.0,
    negatives: Negatives = ...,
    ambient_repulsion_weight: float = 0.0,
    max_edges_per_device: int = 0,
    pair_loss: PairLossForm = ...,
)
```

## matryoshka_dimensions

```python
matryoshka_dimensions: Final[Sequence[int] | None]
```

## temperature_groups

```python
temperature_groups: Final[int]
```

## inverse_temperature

```python
inverse_temperature: Final[float]
```

## bias

```python
bias: Final[float]
```

## negatives

```python
negatives: Final[Negatives]
```

## ambient_repulsion_weight

```python
ambient_repulsion_weight: Final[float]
```

## max_edges_per_device

```python
max_edges_per_device: Final[int]
```

## pair_loss

```python
pair_loss: Final[PairLossForm]
```

## LateInteraction

How the late-interaction loss scores the token vectors of queries and documents; every zero derives.

```python
class LateInteraction
```

## LateInteraction

```python
LateInteraction(
    *,
    weight: float = 0.0,
    width: int = 0,
    inverse_temperature: float = 0.0,
    bias: float = 0.0,
    max_query_span: int = 0,
    max_query_tokens_per_device: int = 0,
    max_documents_per_device: int = 0,
    max_document_tokens_per_device: int = 0,
    pair_loss: PairLossForm = ...,
)
```

## weight

```python
weight: Final[float]
```

## width

```python
width: Final[int]
```

## inverse_temperature

```python
inverse_temperature: Final[float]
```

## bias

```python
bias: Final[float]
```

## max_query_span

```python
max_query_span: Final[int]
```

## max_query_tokens_per_device

```python
max_query_tokens_per_device: Final[int]
```

## max_documents_per_device

```python
max_documents_per_device: Final[int]
```

## max_document_tokens_per_device

```python
max_document_tokens_per_device: Final[int]
```

## pair_loss

```python
pair_loss: Final[PairLossForm]
```

## Schedule

The learning-rate schedule: a linear warmup to the peak, then a cosine decay to the floor.

```python
class Schedule
```

## Schedule

```python
Schedule(
    *,
    learning_rate: float = 0.0,
    min_learning_rate: float = 0.0,
    warmup_steps: int = 0,
    steps: int = 0,
)
```

## learning_rate

```python
learning_rate: Final[float]
```

## min_learning_rate

```python
min_learning_rate: Final[float]
```

## warmup_steps

```python
warmup_steps: Final[int]
```

## steps

```python
steps: Final[int]
```

## Edges

One edge set over a step's items, as compressed sparse rows of buffers borrowed until the step is final.

```python
class Edges
```

## Edges

```python
Edges(
    *,
    offsets: Buffer | None = None,
    ids: Buffer | None = None,
    weights: Buffer | None = None,
)
```

## offsets

```python
offsets: Final[Buffer | None]
```

## ids

```python
ids: Final[Buffer | None]
```

## weights

```python
weights: Final[Buffer | None]
```

## Relations

The contrastive targets of one step: ignored, attracted, repelled, then the run's [`negatives`](#negatives).

```python
class Relations
```

## Relations

```python
Relations(
    *,
    attract: _EdgesLike = None,
    repel: _EdgesLike = None,
    ignore: _EdgesLike = None,
)
```

## attract

```python
attract: Final[Edges | None]
```

## repel

```python
repel: Final[Edges | None]
```

## ignore

```python
ignore: Final[Edges | None]
```

## TrainerStepStats

What [`step`](#step) reports, overwritten by every call it is passed to, every loss its world mean.

```python
class TrainerStepStats
```

## step

The optimizer step the call applied, counted from one.

```python
step: int
```

## learning_rate

The learning rate it applied.

```python
learning_rate: float
```

## loss

The total loss the step minimized, its world mean.

```python
loss: float
```

## contrastive_loss

The contrastive term; 0 when absent.

```python
contrastive_loss: float
```

## masked_text_loss

The text reconstruction term, unweighted.

```python
masked_text_loss: float
```

## masked_image_loss

The image reconstruction term, unweighted.

```python
masked_image_loss: float
```

## masked_video_loss

The video reconstruction term, unweighted.

```python
masked_video_loss: float
```

## masked_audio_loss

The audio reconstruction term, unweighted.

```python
masked_audio_loss: float
```

## late_interaction_loss

The late-interaction term, weighted.

```python
late_interaction_loss: float
```

## choices_loss

The choices term over the items that made a choice, unweighted.

```python
choices_loss: float
```

## truncated

Items cut at a row's width.

```python
truncated: int
```

## dropped

Segments dropped whole because the batch had no room for them, since the last stats read.

```python
dropped: int
```

## Trainer

A training run on the devices of the executor it was created on, which its own verbs block on.

```python
class Trainer
```

## init_vram

Builds one replica per device and starts or resumes the run.

```python
@classmethod
def init_vram(
    executor: Executor,
    /,
    *,
    mesh: Mesh | None = None,
    resume: _PathLike | None = None,
    depth: int = 0,
    hidden: int = 0,
    batch: BatchShape | None = None,
    objective: Objective = ...,
    masking: Masking | None = None,
    contrastive: Contrastive | None = None,
    late_interaction: LateInteraction | None = None,
    choices: Choices | None = None,
    schedule: Schedule | None = None,
    seed: int = 0,
    progress: _Progress | None = None,
    progress_steps: int = 0,
) -> Trainer
```

Calls `uform_trainer_init_vram`.

## on

A receiver whose verbs block on `executor`.

```python
def on(executor: Executor, /) -> On[Trainer]
```

Python-only.

## deferred

A receiver whose verbs submit to `executor` and answer a [`Deferred`](https://www.unum.cloud/docs/uform/python/receivers-and-deferred-calls.md#Deferred).

```python
def deferred(executor: Executor, /) -> DeferredOn[Trainer]
```

Python-only.

## step

Applies one optimizer step over this process's items and every peer's.

```python
def step(
    modalities: Sequence[Modality],
    contents: _Texts,
    /,
    *,
    bags: _Words | None = None,
    temperature_groups: _Words | None = None,
    sides: Sequence[LateInteractionSide] | None = None,
    choices: _Texts | None = None,
    choice_targets: Sequence[Buffer] | _Tape | None = None,
    relations: Relations | None = None,
    arena: Arena | None = None,
    outputs_scalar: Scalar = ...,
    stats: TrainerStepStats | None = None,
) -> ArenaBuffer | None
```

Calls `uform_trainer_step`.

## save

Writes the weights and the step as a checkpoint.

```python
def save(
    path: _PathLike,
    /,
    *,
    progress: _Progress | None = None,
    progress_steps: int = 0,
) -> None
```

Calls `uform_trainer_save`.