Training | UForm for Mojo

A run joins its processes with Mesh.bootstrap over a device executor, then steps one Trainer per process:

var executor = Executor.for_vram(library, VramOptions.on(devices))
var mesh = Mesh.bootstrap(executor, options)
var config = TrainerInitVramConfig[origin_of(mesh)]()
config.depth = 12
config.hidden = 768
config.mesh = Pointer(to=mesh)
var trainer = Trainer.init_vram(executor, config)
trainer.on(executor).step(modalities, contents, TrainerStepConfig(relations=relations))
trainer.on(executor).save("checkpoint.safetensors")

API

StructModalityWhat the bytes of one training item hold, uform_modality_t.

struct Modality

Fields

`code`	`c_int`	The raw `uform_modality_t` value the ABI exchanges.

Values

`TEXT`	`Modality(Int32(0))`	UTF-8 text, a byte per token.
`IMAGE`	`Modality(Int32(1))`	A JPEG, PNG or WebP file.
`VIDEO`	`Modality(Int32(2))`	An MP4 or WebM file.
`AUDIO`	`Modality(Int32(3))`	A WAV, FLAC, MP3 or Opus file.

Conforms to

Copyable, Deinitable, Equatable, ImplicitlyCopyable, Movable, RegisterPassable, TrivialRegisterPassable

StructObjectiveWhich losses a step relaxes, uform_objective_t.

struct Objective

Fields

`code`	`c_int`	The raw `uform_objective_t` value the ABI exchanges.

Values

`MASKED`	`Objective(Int32(0))`	Masked reconstruction alone.
`CONTRASTIVE`	`Objective(Int32(1))`	The pairwise contrastive loss alone.
`JOINT`	`Objective(Int32(2))`	Both, joined at the final-norm states.

Conforms to

Copyable, Deinitable, Equatable, ImplicitlyCopyable, Movable, RegisterPassable, TrivialRegisterPassable

StructNegativesWhich pairs the contrastive loss pushes apart, uform_negatives_t.

struct Negatives

Fields

`code`	`c_int`	The raw `uform_negatives_t` value the ABI exchanges.

Values

`AMBIENT`	`Negatives(Int32(0))`	Repel edges, and every pair no edge names at the ambient weight.
`NAMED`	`Negatives(Int32(1))`	Repel edges alone.

Conforms to

Copyable, Deinitable, Equatable, ImplicitlyCopyable, Movable, RegisterPassable, TrivialRegisterPassable

StructPairLossFormHow a contrastive term turns a row of scores into forces, uform_pair_loss_form_t.

struct PairLossForm

Fields

`code`	`c_int`	The raw `uform_pair_loss_form_t` value the ABI exchanges.

Values

`DERIVED`	`PairLossForm(Int32(0))`	The term's own form: sigmoid for pooled vertices, ranked for token vectors.
`BALANCED`	`PairLossForm(Int32(1))`	The sigmoid cross-entropy of every cell, its attract and repel halves normalized apart.
`RANKED`	`PairLossForm(Int32(2))`	The sigmoid cross-entropy of every positive's score over every negative's.
`SIGMOID`	`PairLossForm(Int32(3))`	SigLIP's sigmoid cross-entropy of every cell, each repel weighed like an attract.

Conforms to

Copyable, Deinitable, Equatable, ImplicitlyCopyable, Movable, RegisterPassable, TrivialRegisterPassable

StructLateInteractionSideWhich side of a late-interaction score an item takes, uform_late_interaction_side_t.

struct LateInteractionSide

Fields

`code`	`c_int`	The raw `uform_late_interaction_side_t` value the ABI exchanges.

Values

`NONE`	`LateInteractionSide(Int32(0))`	The item takes no part in late interaction.
`QUERY`	`LateInteractionSide(Int32(1))`	The item's tokens score the documents it attracts and repels.
`DOCUMENT`	`LateInteractionSide(Int32(2))`	The item's tokens are scored by the queries naming it.
`QUERY_AND_DOCUMENT`	`LateInteractionSide(Int32(3))`	The item plays both sides.

Conforms to

Copyable, Deinitable, Equatable, ImplicitlyCopyable, Movable, RegisterPassable, TrivialRegisterPassable

StructImageShapeThe image slice of a batch envelope, uform_image_shape_t; a zero patch_side leaves images out.1

struct ImageShape

Fields

`patch_side`	`c_size_t`	The square patch side images tile into, in pixels.
`channels`	`c_size_t`	Channels every pixel carries; zero derives four.

Conforms to

Copyable, Deinitable, ImplicitlyCopyable, Movable

__init__Names only the fields that differ from the derived ones.

def __init__(out self, *, patch_side: Int = Int(0), channels: Int = Int(0))

StructVideoShapeThe video slice of a batch envelope, uform_video_shape_t; a zero patch_side leaves video out.1

struct VideoShape

Fields

`patch_side`	`c_size_t`	The square patch side one tube tiles into, in pixels.
`tube_frames`	`c_size_t`	Frames one tube folds ahead of its patch; zero derives two.
`frame_stride`	`c_size_t`	Source frames between sampled ones; zero derives one.
`channels`	`c_size_t`	Channels every pixel carries; zero derives three.

Conforms to

Copyable, Deinitable, ImplicitlyCopyable, Movable

__init__Names only the fields that differ from the derived ones.

def __init__(
    out self,
    *,
    patch_side: Int = Int(0),
    tube_frames: Int = Int(0),
    frame_stride: Int = Int(0),
    channels: Int = Int(0)
)

StructAudioShapeThe audio slice of a batch envelope, uform_audio_shape_t; a zero chunk_samples leaves audio out.1

struct AudioShape

Fields

`chunk_samples`	`c_size_t`	Raw samples per frame, the fan-in of the learned frontend.
`sample_rate`	`c_size_t`	The mono rate every recording is resampled to; zero derives 16000.

Conforms to

Copyable, Deinitable, ImplicitlyCopyable, Movable

__init__Names only the fields that differ from the derived ones.

def __init__(out self, *, chunk_samples: Int = Int(0), sample_rate: Int = Int(0))

StructBatchShapeWhat one device's step carries at most, whatever the modality mix, uform_batch_shape_t.1

struct BatchShape

Fields

`max_rows_per_device`	`c_size_t`	Rows one device's step packs.
`max_segments_per_device`	`c_size_t`	Items one device's step carries across all its rows, which is also its pooling axis.
`max_columns_per_row`	`c_size_t`	Positions one row spans, the sequence length.
`max_bags_per_device`	`c_size_t`	Bags one device's step carries; zero keeps every item its own routing unit.
`max_choices_per_item`	`c_size_t`	Legal bytes one item offers to choose from at most; zero makes no choices.
`image`	`ImageShape`	The image geometry.
`video`	`VideoShape`	The video geometry.
`audio`	`AudioShape`	The audio geometry.

Conforms to

Copyable, Deinitable, ImplicitlyCopyable, Movable

__init__Names only the fields that differ from the derived ones.

def __init__(
    out self,
    *,
    max_rows_per_device: Int = Int(0),
    max_segments_per_device: Int = Int(0),
    max_columns_per_row: Int = Int(0),
    max_bags_per_device: Int = Int(0),
    max_choices_per_item: Int = Int(0),
    image: ImageShape = ImageShape(patch_side=Int(0), channels=Int(0)),
    video: VideoShape = VideoShape(patch_side=Int(0), tube_frames=Int(0), frame_stride=Int(0), channels=Int(0)),
    audio: AudioShape = AudioShape(chunk_samples=Int(0), sample_rate=Int(0))
)

StructMaskingHow the masked objective corrupts a batch and weighs its heads, uform_masking_t; every zero derives the default its field names.1

struct Masking

Fields

`token_mask_probability`	`Float32`	Fraction of text positions covered by spans; 0.15.
`token_mean_span`	`Float32`	Mean geometric span length in bytes; 16.
`replace_fraction`	`Float32`	Of covered positions, the share shown as MASK; 0.8.
`random_fraction`	`Float32`	Of covered positions, the share replaced by a random byte; 0.1.
`image_mask_ratio`	`Float32`	Fraction of image patches masked; 0.75.
`video_mask_ratio`	`Float32`	Fraction of video tubes masked; 0.75.
`audio_mask_ratio`	`Float32`	Fraction of audio frames masked; 0.75.
`text_weight`	`Float32`	The weight of the text head; 1.
`image_weight`	`Float32`	The weight of the image head; 1.
`video_weight`	`Float32`	The weight of the video head; 1.
`audio_weight`	`Float32`	The weight of the audio head; 1.

Conforms to

Copyable, Defaultable, Deinitable, ImplicitlyCopyable, Movable

__init__Derives every default.

def __init__(out self)

StructLateInteractionHow the late-interaction loss scores per-token vectors, uform_late_interaction_t; every zero derives the default its field names.1

struct LateInteraction

Fields

`weight`	`Float32`	The loss weight; zero trains no token head.
`width`	`c_size_t`	The token vectors' width, a multiple of 8; 128.
`inverse_temperature`	`Float32`	The starting inverse temperature of every contrastive group; 1.
`bias`	`Float32`	The starting logit bias of every contrastive group; zero.
`max_query_span`	`c_size_t`	Tokens one query segment contributes at most; 64.
`max_query_tokens_per_device`	`c_size_t`	Query token rows one device's step holds at most.
`max_documents_per_device`	`c_size_t`	Document segments one device's step holds at most.
`max_document_tokens_per_device`	`c_size_t`	Document token rows one device's step holds at most.
`pair_loss`	`PairLossForm`	The pair loss over each query's row of hard MaxSim scores; ranked.

Conforms to

Copyable, Defaultable, Deinitable, ImplicitlyCopyable, Movable

__init__Trains no token head.

def __init__(out self)

StructScheduleThe learning-rate schedule, a linear warmup then a cosine decay, uform_schedule_t.1

struct Schedule

Fields

`learning_rate`	`Float32`	The peak learning rate; zero derives 5e-4.
`min_learning_rate`	`Float32`	The floor the cosine settles at; zero derives 1e-5.
`warmup_steps`	`c_size_t`	Steps from zero to the peak; zero derives a tenth of `steps`, at most 400.
`steps`	`c_size_t`	Steps the run takes across resumes, which the cosine spans; zero is refused.

Conforms to

Copyable, Defaultable, Deinitable, ImplicitlyCopyable, Movable

__init__Derives every default but steps, which a run must name.

def __init__(out self)

StructTrainerStepStatsWhat a step reports, every loss its world mean, uform_trainer_step_stats_t.1

struct TrainerStepStats

Fields

`step`	`c_size_t`	The optimizer step the call applied, counted from one across resumes.
`learning_rate`	`Float32`	The learning rate it applied.
`loss`	`Float32`	The total loss the step minimized.
`contrastive_loss`	`Float32`	The contrastive term; zero when the objective is masked alone.
`masked_text_loss`	`Float32`	The text reconstruction term, before its weight.
`masked_image_loss`	`Float32`	The image reconstruction term, before its weight.
`masked_video_loss`	`Float32`	The video reconstruction term, before its weight.
`masked_audio_loss`	`Float32`	The audio reconstruction term, before its weight.
`late_interaction_loss`	`Float32`	The late-interaction term, weighted.
`choices_loss`	`Float32`	The choices term over the items that made a choice, before its weight.
`truncated`	`c_size_t`	Items cut at a row's width.
`dropped`	`c_size_t`	Segments dropped whole for lack of room, since the last step that read its stats.

Conforms to

Copyable, Defaultable, Deinitable, ImplicitlyCopyable, Movable

__init__Zeroes every field, which is what a failed step leaves.

def __init__(out self)

StructContrastiveHow the contrastive loss scores the pooled embeddings, uform_contrastive_t; every zero derives the default its field names.1

struct Contrastive

Fields

`matryoshka_dimensions`	`List[UInt]`	Matryoshka prefix widths, strictly ascending, each a multiple of 8, at most 8; empty scores the full width alone.
`temperature_groups`	`Int`	Learnable inverse temperature and bias pairs, one per group a step's items select; one.
`inverse_temperature`	`Float32`	The starting inverse temperature of every group; 10 under the sigmoid form, else 1.
`bias`	`Float32`	The starting logit bias of every group; -10 under the sigmoid form, else zero.
`negatives`	`Negatives`	Which pairs repel.
`ambient_repulsion_weight`	`Float32`	The weight of every pair no edge names, under `Negatives.AMBIENT`; 1.
`max_edges_per_device`	`Int`	Edges one device's step names per edge set at most; zero derives it.
`pair_loss`	`PairLossForm`	The pair loss over each anchor's row; sigmoid.

Conforms to

Copyable, Defaultable, Deinitable, Movable

__init__Derives every default.

def __init__(out self)

StructTrainerInitVramConfigThe geometry, objective and starting point of a run, uform_trainer_init_vram_t.1

struct TrainerInitVramConfig[mesh_origin: ImmOrigin = ImmUntrackedOrigin]

Parameters

`mesh_originImmOrigin`	comptime	The borrowed mesh's origin, which the trainer then keeps alive.

Fields

`resume`	`String`	The checkpoint path to continue from, at its saved step; empty starts from `seed`.
`depth`	`Int`	Transformer blocks stacked.
`hidden`	`Int`	Embedding width, a multiple of 128.
`batch`	`BatchShape`	What one device's step carries at most.
`objective`	`Objective`	Which losses a step relaxes.
`masking`	`Masking`	The masked objective's corruption and weights, read unless the objective is contrastive alone.
`contrastive`	`Contrastive`	The contrastive loss's scoring, read unless the objective is masked alone.
`late_interaction`	`LateInteraction`	The late-interaction loss beside the pooled one, read when its weight is positive.
`choices`	`Choices`	The calibrated choices, read when its weight is positive.
`schedule`	`Schedule`	The learning-rate schedule.
`mesh`	`Optional[Pointer[Mesh, mesh_origin]]`	The mesh made over the same executor; absent runs this process alone.
`seed`	`UInt64`	The base of every keyed draw: initialization, corruption and the shuffle.
`progress`	`Progress`	Reports the tensors placed on a resume, and stops the call where it says to.

Conforms to

Copyable, Defaultable, Deinitable, Movable

__init__Every field zeroed or derived, which a run then overrides by name.

def __init__(out self)

StructEdgesOne edge set over a step's items as compressed sparse rows; empty offsets name no edges.2

struct Edges[origin: ImmOrigin = ImmUntrackedOrigin]

Parameters

`originImmOrigin`	comptime	The borrowed offsets, IDs and weights, kept alive for as long as the edges are used.

Fields

`offsets`	`Span[UInt, origin]`	The step's count plus one non-decreasing offsets into `ids`.
`ids`	`Span[UInt32, origin]`	Item `i` links to items `ids[offsets[i]]` up to `ids[offsets[i + 1]]`, ascending within each run.
`weights`	`Span[Float32, origin]`	The weight of each edge, paired with `ids`; empty weighs every edge one.

Conforms to

Copyable, Defaultable, Deinitable, ImplicitlyCopyable, Movable

__init__Names no edges.

def __init__(out self)

ofBorrows the three arrays of one edge set; an empty weights weighs every edge one.

def of[offsets_origin: ImmOrigin, ids_origin: ImmOrigin, weights_origin: ImmOrigin, //](
    offsets: Span[UInt, offsets_origin],
    ids: Span[UInt32, ids_origin],
    weights: Span[Float32, weights_origin]
) -> Edges[origin_of(ids_origin, offsets_origin, weights_origin)]

StructRelationsThe contrastive targets of one step, resolved per pair in this order: a pair ignore names is skipped, one attract names pulls together, one repel names pushes apart, and every other pair follows the run's negatives.2

struct Relations[origin: ImmOrigin = ImmUntrackedOrigin]

Parameters

`originImmOrigin`	comptime	The borrowed edge sets, kept alive for as long as the relations are used.

Fields

`attract`	`Edges[origin]`	Pairs pulled together.
`repel`	`Edges[origin]`	Pairs pushed apart.
`ignore`	`Edges[origin]`	Pairs skipped, like known false negatives; their weights are not read.

Conforms to

Copyable, Defaultable, Deinitable, ImplicitlyCopyable, Movable

__init__Names no pairs.

def __init__(out self)

ofGathers the three edge sets of one step.

def of[attract_origin: ImmOrigin, repel_origin: ImmOrigin, ignore_origin: ImmOrigin, //](
    attract: Edges[attract_origin],
    repel: Edges[repel_origin] = Edges(),
    ignore: Edges[ignore_origin] = Edges()
) -> Relations[origin_of(attract_origin, ignore_origin, repel_origin)]

StructTrainerStepConfigThe routing, targets and report slot of a training step, uform_trainer_step_t; every array is optional and, when given, holds one entry per item.1

struct TrainerStepConfig[bags_origin: ImmOrigin = ImmUntrackedOrigin, groups_origin: ImmOrigin = ImmUntrackedOrigin, sides_origin: ImmOrigin = ImmUntrackedOrigin, choices_origin: ImmOrigin = ImmUntrackedOrigin, targets_origin: ImmOrigin = ImmUntrackedOrigin, relations_origin: ImmOrigin = ImmUntrackedOrigin, stats_origin: MutOrigin = MutUntrackedOrigin]

Fields

`bags`	`Span[UInt32, bags_origin]`	The bag of each item, items sharing a bag landing on one device; empty keeps every item its own.
`temperature_groups`	`Span[UInt32, groups_origin]`	The temperature group of each item; empty puts every item in group zero.
`sides`	`Span[LateInteractionSide, sides_origin]`	The late-interaction side of each item; empty makes every item `LateInteractionSide.NONE`.
`choices`	`Optional[TapeView[choices_origin]]`	Each item's legal choices, one byte each, as `text_choose` reads them; absent makes no choices.
`choice_targets`	`Optional[TapeView[targets_origin]]`	The target probability of each choice, `F32`, as `text_choose` answers them; read with `choices`.
`relations`	`Relations[relations_origin]`	The contrastive targets, read unless the objective is masked alone.
`outputs_scalar`	`ScalarKind`	The scalar type the pooled embeddings land as, read only with an arena.
`stats`	`Optional[Pointer[TrainerStepStats, stats_origin]]`	Receives the losses, overwritten by the call; absent skips their readback too.

Conforms to

Copyable, Deinitable, ImplicitlyCopyable, Movable

__init__Names only the arrays and knobs a step uses.

def __init__(
    out self,
    *,
    bags: Span[UInt32, bags_origin] = Span(),
    temperature_groups: Span[UInt32, groups_origin] = Span(),
    sides: Span[LateInteractionSide, sides_origin] = Span(),
    choices: Optional[TapeView[choices_origin]] = Optional(),
    choice_targets: Optional[TapeView[targets_origin]] = Optional(),
    relations: Relations[relations_origin] = Relations(),
    outputs_scalar: ScalarKind = ScalarKind.F32,
    stats: Optional[Pointer[TrainerStepStats, stats_origin]] = Optional()
)

StructTrainerA training run on a device executor's devices, released when it goes out of scope.5

struct Trainer[mesh_origin: ImmOrigin = ImmUntrackedOrigin]

Its verbs run only on an executor: trainer.on(executor) blocks there, and trainer.deferred(executor) submits there and answers a Deferred.

Parameters

`mesh_originImmOrigin`	comptime	The mesh it borrows, kept alive for as long as the trainer is.

Fields

`library`	`Library`	The loaded engine, held so its symbol table outlives every call made through the handle.
`address`	`Address`	The opaque `uform_trainer_t`.
`mesh`	`Optional[Pointer[Mesh, mesh_origin]]`	The mesh the run spans, absent for a process running alone.

Conforms to

Deinitable, Movable

__init__Adopts a handle the engine just created.

def __init__(
    out self,
    *,
    library: Library,
    address: Int,
    mesh: Optional[Pointer[Mesh, mesh_origin]]
)

__deinit__Frees the trainer and leaves the mesh; every call using it must be final. Calls uform_trainer_free.

def __deinit__(deinit self)

init_vramBuilds one replica per device and starts or resumes the run; every process of the run calls it with the same geometry, objective and seed. Calls uform_trainer_init_vram.

def init_vram[config_mesh_origin: ImmOrigin, //](
    executor: Executor,
    config: TrainerInitVramConfig[config_mesh_origin]
) -> Trainer[config_mesh_origin]

Raises

Error with CONFIG_MISMATCH for a checkpoint of another geometry, or UNSUPPORTED in a build without the trainer.

onThe trainer's verbs, blocking on executor, the one it was made over. Binding-only.

def on(ref self, ref executor: Executor) -> TrainerOn[origin_of(executor, self)]

deferredThe trainer's verbs, each submitted to executor and answering a Deferred at once. Binding-only.

def deferred(ref self, ref executor: Executor) -> TrainerDeferredOn[origin_of(executor, self)]

StructTrainerOnA trainer bound to the executor its verbs block on, as trainer.on(executor) answers it. Binding-only.1

struct TrainerOn[origin: ImmOrigin]

Parameters

`originImmOrigin`	comptime	The trainer and the executor, which this receiver borrows.

Fields

`library`	`Library`	The loaded engine every verb crosses into.
`trainer`	`Address`	The `uform_trainer_t` every verb addresses.
`executor`	`Address`	The `uform_executor_t` every verb runs on.

Conforms to

Copyable, Deinitable, ImplicitlyCopyable, Movable, TrainVerbs

__init__Binds a trainer to an executor.

def __init__(out self, library: Library, trainer: Int, executor: Int)

TraitTrainVerbsThe verbs every blocking trainer receiver offers, each written once over the receiver's _route.2

trait TrainVerbs

stepApplies one optimizer step over this process's items and every peer's; every process of the run calls it once per step. Calls uform_trainer_step.

def step[modalities_origin: ImmOrigin, contents_origin: ImmOrigin, bags_origin: ImmOrigin, groups_origin: ImmOrigin, sides_origin: ImmOrigin, choices_origin: ImmOrigin, targets_origin: ImmOrigin, relations_origin: ImmOrigin, stats_origin: MutOrigin, //](
    self,
    modalities: Span[Modality, modalities_origin],
    contents: TapeView[contents_origin],
    config: TrainerStepConfig[bags_origin, groups_origin, sides_origin, choices_origin, targets_origin, relations_origin, stats_origin] = TrainerStepConfig(bags=Span(), temperature_groups=Span(), sides=Span(), choices=Optional(), choice_targets=Optional(), relations=Relations(), outputs_scalar=ScalarKind.F32, stats=Optional())
)

def step[modalities_origin: ImmOrigin, contents_origin: ImmOrigin, bags_origin: ImmOrigin, groups_origin: ImmOrigin, sides_origin: ImmOrigin, choices_origin: ImmOrigin, targets_origin: ImmOrigin, relations_origin: ImmOrigin, stats_origin: MutOrigin, //](
    self,
    modalities: Span[Modality, modalities_origin],
    contents: TapeView[contents_origin],
    mut arena: Arena,
    config: TrainerStepConfig[bags_origin, groups_origin, sides_origin, choices_origin, targets_origin, relations_origin, stats_origin] = TrainerStepConfig(bags=Span(), temperature_groups=Span(), sides=Span(), choices=Optional(), choice_targets=Optional(), relations=Relations(), outputs_scalar=ScalarKind.F32, stats=Optional())
) -> Rows[origin_of(arena)]

The same step, also landing the pooled embedding of every item in arena, as the contrastive loss scored it before the update. Calls uform_trainer_step.

Parameters

`modalitiesSpan[Modality, modalities_origin]`		The modality of each item.
`contentsTapeView[contents_origin]`		Each item's bytes, UTF-8 text or an encoded file as `modalities` tags it.
`configTrainerStepConfig[bags_origin, groups_origin, sides_origin, choices_origin, targets_origin, relations_origin, stats_origin]`		The routing, the targets and the report slot.

Raises

Error with SHAPE_MISMATCH when the per-item arrays disagree with contents.

saveWrites the weights and the step, which the model loads serve and init_vram resumes; only the process of index zero writes. Calls uform_trainer_save.

def save[path_origin: ImmOrigin, //](
    self,
    path: StringSpan[path_origin],
    progress: Progress = Progress()
)