Training is GPU-only. The configuration arrives as structs mirroring the C ones field by field, so a field left at its default is zero, which the library derives:
let executor = unsafe { Executor::for_vram(&[0, 1], None, &[], None, &[], 0)? };
let mesh = Mesh::bootstrap(&executor, &MeshOptions {
host_domain_index: rank,
host_domain_count: nodes,
directory: Path::new("/shared/exchange"),
job: "run-17",
..MeshOptions::default()
})?;
let trainer = Trainer::init_vram(&executor, TrainerInitVramConfig {
mesh: Some(&mesh),
depth: 12,
hidden: 512,
objective: Objective::Joint,
batch: BatchShape { max_rows_per_device: 64, max_segments_per_device: 256, max_columns_per_row: 1024,
max_bags_per_device: 128, image: ImageShape { patch_side: 16, channels: 0 },
..BatchShape::default() },
schedule: Schedule { steps: 100_000, ..Schedule::default() },
seed: 42,
..TrainerInitVramConfig::default()
})?;
let mut stats = TrainerStepStats::default();
let relations = Relations { attract: Edges { offsets: &[0, 1, 2, 3, 4], ids: &[1, 0, 3, 2], weights: &[] },
..Relations::default() };
let config = TrainerStepConfig { bags: &[0, 0, 1, 1], relations, stats: Some(&mut stats), ..TrainerStepConfig::default() };
trainer.on(&executor).step(&modalities, &contents, config)?;
trainer.on(&executor).save("checkpoint.safetensors")?;BatchShape nests an ImageShape, a VideoShape and an AudioShape, while Masking, Contrastive, LateInteraction, Choices and Schedule configure the objectives, and PairLossForm picks each contrastive term’s pair loss. Relations holds three Edges sets as compressed sparse rows, which the binding checks against the step’s items before the call.
A step’s sides give each item a LateInteractionSide, and its choices and choice_targets give each choice item its legal bytes and their target probabilities as (data, offsets) tapes, the shapes text_choose takes and answers, checked the same way.
Paths cross as the OS’s bytes on Unix, not necessarily UTF-8, while the mesh’s job is a name and a &str. step answers the pooled embeddings the loss scored when its config names an arena, and None otherwise.
A trainer borrows its mesh, so the mesh outlives every trainer made over it, and mesh.describe() answers a MeshDescription whose multimem_tier is a MultimemTier and whose gpu_class is the GPUClass every GPU of the run dispatches as.
AudioShapeThe audio slice of a batch envelope; a zero chunk_samples leaves audio out of the run.pub struct AudioShapeMirrors uform_audio_shape_t.
chunk_samples | usize | Raw samples per frame, the fan-in of the learned frontend. |
|---|---|---|
sample_rate | usize | The mono rate every recording is resampled to; zero derives 16000. |
Clone, Copy, StructuralPartialEq, PartialEq, Eq, Debug, DefaultBatchShapeWhat one device's step carries at most, whatever the modality mix: capacity fixed at creation.pub struct BatchShapeText always takes part; a media modality takes part when its slice is live.
Mirrors uform_batch_shape_t.
max_rows_per_device | usize | Rows one device's step packs. |
|---|---|---|
max_segments_per_device | usize | Items one device's step carries across all its rows, which is also its pooling axis. |
max_columns_per_row | usize | Positions one row spans, the sequence length. |
max_bags_per_device | usize | Bags one device's step carries; zero keeps every item its own routing unit. |
max_choices_per_item | usize | 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. |
Clone, Copy, StructuralPartialEq, PartialEq, Eq, Debug, DefaultContrastiveHow the contrastive loss scores the pooled embeddings; every zero derives the default the C header names.pub struct Contrastive<'config>Mirrors uform_contrastive_t.
matryoshka_dimensions | &'config [usize] | Matryoshka prefix widths, strictly ascending, each a multiple of 8, at most 8 of them, every one scored; empty scores the full width alone. Copied, not borrowed. |
|---|---|---|
temperature_groups | usize | Learnable inverse temperature and bias pairs, one per group a step's items select; one. |
inverse_temperature | f32 | The starting inverse temperature of every group; 10 under the sigmoid form, SigLIP's, else 1. |
bias | f32 | The starting logit bias of every group; −10 under the sigmoid form, SigLIP's, else zero. |
negatives | Negatives | Which pairs repel. |
ambient_repulsion_weight | f32 | The weight of every pair no edge names, under Negatives::Ambient; 1. |
max_edges_per_device | usize | Edges one device's step names per edge set at most; max_segments_per_device. |
pair_loss | PairLossForm | The pair loss over each anchor's row; sigmoid. |
Clone, Copy, StructuralPartialEq, PartialEq, Debug, DefaultEdgesOne edge set over a step's items, as compressed sparse rows; empty offsets name no edges.pub struct Edges<'config>Mirrors uform_edges_view_t.
offsets | &'config [usize] | The step's count plus one non-decreasing offsets into ids. |
|---|---|---|
ids | &'config [u32] | Item i links to items ids[offsets[i]..offsets[i + 1]], indices into the step's items,
ascending within each run. |
weights | &'config [f32] | The weight of each edge, paired with ids; empty weighs every edge one. |
Clone, Copy, StructuralPartialEq, PartialEq, Debug, DefaultImageShapeThe image slice of a batch envelope; a zero patch_side leaves images out of the run.pub struct ImageShapeMirrors uform_image_shape_t.
patch_side | usize | The square patch side images tile into, in pixels. |
|---|---|---|
channels | usize | Channels every pixel carries; zero derives four, RGBA end to end. |
Clone, Copy, StructuralPartialEq, PartialEq, Eq, Debug, DefaultLateInteractionHow the late-interaction loss scores the per-token vectors of query and document items; every zero derives the default the C header names.pub struct LateInteractionMirrors uform_late_interaction_t.
weight | f32 | The loss weight; zero trains no token head. |
|---|---|---|
width | usize | The token vectors' width, a multiple of 8; 128. |
inverse_temperature | f32 | The starting inverse temperature of every contrastive group; 1. |
bias | f32 | The starting logit bias of every contrastive group; zero is zero. |
max_query_span | usize | Tokens one query item contributes at most, its first past the prefix; 64. |
max_query_tokens_per_device | usize | Query token rows one device's step holds at most; max_query_span of every item. |
max_documents_per_device | usize | Document items one device's step holds at most; max_segments_per_device. |
max_document_tokens_per_device | usize | Document token rows one device's step holds at most; the batch's tokens per device. |
pair_loss | PairLossForm | The pair loss over each query's row of hard MaxSim scores; ranked. |
Clone, Copy, StructuralPartialEq, PartialEq, Debug, DefaultLateInteractionSideWhich side of a late-interaction score an item takes: a query's tokens score a document's, and a symmetric pair plays both.pub enum LateInteractionSideMirrors uform_late_interaction_side_t.
None | 0 | The item takes no part in late interaction. |
|---|---|---|
Query | 1 | The item's tokens score the documents it attracts and repels. |
Document | 2 | The item's tokens are scored by the queries naming it. |
QueryAndDocument | 3 | The item plays both sides. |
Clone, Copy, StructuralPartialEq, PartialEq, Eq, Debug, Default, HashMaskingHow the masked objective corrupts a batch and weighs its heads; every zero derives the default the C header names.pub struct MaskingMirrors uform_masking_t.
token_mask_probability | f32 | Fraction of text positions covered by spans; 0.15. |
|---|---|---|
token_mean_span | f32 | Mean geometric span length in bytes; 16. |
replace_fraction | f32 | Of covered positions, the share shown as MASK; 0.8. |
random_fraction | f32 | Of covered positions, the share replaced by a random byte, the rest kept clean; 0.1. |
image_mask_ratio | f32 | Fraction of image patches masked; 0.75. |
video_mask_ratio | f32 | Fraction of video tubes masked, tubed across time; 0.75. |
audio_mask_ratio | f32 | Fraction of audio frames masked; 0.75. |
text_weight | f32 | The weight of the text head; 1. |
image_weight | f32 | The weight of the image head; 1. |
video_weight | f32 | The weight of the video head; 1. |
audio_weight | f32 | The weight of the audio head; 1. |
Clone, Copy, StructuralPartialEq, PartialEq, Debug, DefaultModalityWhat the bytes of one training item hold.pub enum ModalityMirrors uform_modality_t.
Text | 0 | UTF-8 text, a byte per token. |
|---|---|---|
Image | 1 | A JPEG, PNG or WebP file. |
Video | 2 | An MP4 or WebM file. |
Audio | 3 | A WAV, FLAC, MP3 or Opus file. |
Clone, Copy, StructuralPartialEq, PartialEq, Eq, Debug, Hash, Display, FromStrNegativesWhich pairs the contrastive loss pushes apart.pub enum NegativesMirrors uform_negatives_t.
Ambient | 0 | Repel edges, and every pair no edge names, at the ambient repulsion weight. |
|---|---|---|
Named | 1 | Repel edges alone. |
Clone, Copy, StructuralPartialEq, PartialEq, Eq, Debug, Default, Hash, Display, FromStrObjectiveWhich losses a step relaxes, one enumerator per combination the engine trains.pub enum ObjectiveMirrors uform_objective_t.
Masked | 0 | Masked reconstruction alone, as the Base pretrains. |
|---|---|---|
Contrastive | 1 | The pairwise contrastive loss alone, as the Embedder trains. |
Joint | 2 | Both, joined at the final-norm states. |
Clone, Copy, StructuralPartialEq, PartialEq, Eq, Debug, Default, Hash, Display, FromStrPairLossFormHow a contrastive term turns a row of scores into forces.pub enum PairLossFormMirrors uform_pair_loss_form_t.
Derived | 0 | The term's own form: sigmoid for pooled vertices, ranked for token vectors. |
|---|---|---|
Balanced | 1 | The sigmoid cross-entropy of every cell, its attract and repel halves normalized apart. |
Ranked | 2 | The sigmoid cross-entropy of every positive's score over every negative's, where the bias cancels. |
Sigmoid | 3 | SigLIP's sigmoid cross-entropy of every cell, each repel weighed like an attract. |
Clone, Copy, StructuralPartialEq, PartialEq, Eq, Debug, Default, HashRelationsThe 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.pub struct Relations<'config>Mirrors uform_relations_view_t.
attract | Edges<'config> | Pairs pulled together. |
|---|---|---|
repel | Edges<'config> | Pairs pushed apart. |
ignore | Edges<'config> | Pairs skipped, like known false negatives; their weights are not read. |
ScheduleThe learning-rate schedule: a linear warmup to the peak, then a cosine decay to the floor.pub struct ScheduleMirrors uform_schedule_t.
learning_rate | f32 | The peak learning rate; zero derives 5e-4. |
|---|---|---|
min_learning_rate | f32 | The floor the cosine settles at; zero derives 1e-5. |
warmup_steps | usize | Steps from zero to the peak; zero derives a tenth of steps, at most 400. |
steps | usize | Steps the run takes across resumes, which the cosine spans; zero is refused. |
Clone, Copy, StructuralPartialEq, PartialEq, Debug, DefaultTrainVerbsEvery trainer verb, written once for both receivers of Trainer::on and Trainer::deferred, which run on the executor the trainer was made on.4pub trait TrainVerbs<'trainer>: CopyRust-only.
impl<'trainer, Completion: Completion> TrainVerbs<'trainer> for On<'trainer, Trainer<'_>, Completion>CompletionBlocking or Deferring.type Completion: Completion;handlesThe trainer this receiver reaches and the executor its calls run on.fn handles(&self) -> HandlesstepApplies one optimizer step over this process's items and every peer's: forward, losses, backward, the world gradient reduction, and the update.fn step<'config, 'arena, Modalities: AsRef<[Modality]>, Contents: TapeView<u8>>(
self,
modalities: Modalities,
contents: Contents,
config: impl Into<TrainerStepConfig<'config, 'arena>>,
) -> <<Self as TrainVerbs<'trainer>>::Completion as Completion>::Answer<'trainer, TrainerStep<'config, 'arena, Modalities, Contents>>
where
Self::Completion: Accepts<Modalities> + Accepts<Contents> + Accepts<TrainerStepConfig<'config, 'arena>>,Every process of the run calls it
once per step, with one modality per item and each item's bytes in contents.
Answers the pooled embeddings when the config names an arena.
Calls uform_trainer_step.
saveWrites the weights and the step, which the model loads serve and Trainer::init_vram resumes; only the process of index zero writes, the rest return at once.fn save<'config, Location: AsRef<Path>>(
self,
config: impl Into<TrainerSaveConfig<'config, Location>>,
) -> <<Self as TrainVerbs<'trainer>>::Completion as Completion>::Answer<'trainer, TrainerSave<'config, Location>>
where
Self::Completion: Accepts<TrainerSaveConfig<'config, Location>>,Calls uform_trainer_save.
TrainerA training run: the tower replicas, gradients, optimizer moments, heads, temperatures and the collectives sized by them, on its executor's devices, whose verbs come from TrainVerbs.3pub struct Trainer<'mesh>use uform::*;
let executor = unsafe { Executor::for_vram(&[0], None, &[], None, &[], 0, 0)? };
let trainer = Trainer::init_vram(&executor, TrainerInitVramConfig {
depth: 12,
hidden: 512,
batch: BatchShape { max_rows_per_device: 64, max_segments_per_device: 256, max_columns_per_row: 1024,
..BatchShape::default() },
schedule: Schedule { steps: 100_000, ..Schedule::default() },
seed: 42,
..TrainerInitVramConfig::default()
})?;
let mut stats = TrainerStepStats::default();
let config = TrainerStepConfig { stats: Some(&mut stats), ..TrainerStepConfig::default() };
trainer.on(&executor).step([Modality::Text], ["a caption"], config)?;
trainer.on(&executor).save("checkpoint.safetensors")?;
Mirrors uform_trainer_t; dropping it calls uform_trainer_free. 'mesh is the borrow of the
Mesh it was made over.
Debug, Send, Dropinit_vramBuilds one replica per device of executor, the NVLink domain its consumers reserve over the mesh, and starts or resumes the run; every process of the run calls it with the same geometry, objective and seed.pub fn init_vram<'config>(
executor: &Executor,
config: impl Into<TrainerInitVramConfig<'config, 'mesh>>,
) -> Result<Self>Calls uform_trainer_init_vram.
onThe receiver whose verbs block on executor, the one the trainer was made on.pub fn on<'trainer>(&'trainer self, executor: &'trainer Executor) -> On<'trainer, Trainer<'mesh>>Rust-only.
TrainerInitVramConfigThe geometry, objective and starting point of a run; see Trainer::init_vram.pub struct TrainerInitVramConfig<'config, 'mesh>Mirrors uform_trainer_init_vram_t.
resume | Option<&'config Path> | The checkpoint to continue from, at its saved step; None starts from seed. |
|---|---|---|
depth | usize | Transformer blocks stacked. |
hidden | usize | 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<'config> | 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 beside every other term, read when its weight is positive, which
needs the batch's max_choices_per_item. |
schedule | Schedule | The learning-rate schedule. |
mesh | Option<&'mesh Mesh> | The mesh made over the same executor, borrowed until the trainer drops; None runs
this process alone. |
seed | u64 | The base of every keyed draw: initialization, corruption and the shuffle. |
progress | Option<&'config mut ProgressCallback<'config>> | Reports the tensors placed on a resume and stops wherever it says to. |
progress_steps | usize | Reports the call makes at most, the last at completion; zero derives to ten. |
DefaultTrainerSaveA checkpoint write; see TrainVerbs::save.TrainerSaveConfigThe path a checkpoint is written to and the handler counting its tensors; see TrainVerbs::save.pub struct TrainerSaveConfig<'config, Location>Mirrors uform_trainer_save_t.
path | Location | Destination of the safetensors checkpoint, written whole and renamed into place. |
|---|---|---|
progress | Option<&'config mut ProgressCallback<'config>> | Reports the tensors written and stops wherever it says to. |
progress_steps | usize | Reports the call makes at most, the last at completion; zero derives to ten. |
From<Location>TrainerStepOne optimizer step; see TrainVerbs::step.pub struct TrainerStep<'config, 'arena, Modalities, Contents>Calls uform_trainer_step.
TrainerStepConfigThe routing, targets and reports of TrainVerbs::step.pub struct TrainerStepConfig<'config, 'arena>Mirrors uform_trainer_step_t.
bags | &'config [u32] | The bag of each item, items sharing a bag landing on one device; empty keeps every item its own bag. |
|---|---|---|
temperature_groups | &'config [u32] | The temperature group of each item, each below the run's temperature_groups; empty puts
every item in group zero. |
sides | &'config [LateInteractionSide] | The late-interaction side of each item; empty makes every item
LateInteractionSide::None. |
choices | (&'config [u8], &'config [usize]) | Each item's legal choices, one byte each, as a (data, offsets) tape like the one
text_choose takes; empty offsets make no choices. Read
when the run trains them: a text item offering any is the state its answer follows, and the
call appends the mask byte, cutting a state longer than a row at its head. |
choice_targets | (&'config [f32], &'config [usize]) | The target probability of each choice, a (data, offsets) tape whose entries are as long as
those of choices and sum to one, as text_choose answers its probabilities. |
relations | Relations<'config> | The contrastive targets, read unless the objective is masked alone. |
arena | Option<&'arena mut Arena> | The arena the pooled embeddings land in, as the contrastive loss scored them before the
update; None skips the copy, and the step answers no rows. |
outputs_scalar | Scalar | The scalar type the embeddings land as, read only with an arena. |
stats | Option<&'config mut TrainerStepStats> | Where the losses land, overwritten by the call; None skips their readback too. |
Default, From<Relations<'config>>TrainerStepStatsWhat a step reports, every loss its world mean, overwritten by the call.pub struct TrainerStepStatsMirrors uform_trainer_step_stats_t.
step | usize | The optimizer step the call applied, counted from one across resumes. |
|---|---|---|
learning_rate | f32 | The learning rate it applied. |
loss | f32 | The total loss the step minimized. |
contrastive_loss | f32 | The contrastive term; zero when the objective is masked alone. |
masked_text_loss | f32 | The text reconstruction term, before its weight; zero when absent. |
masked_image_loss | f32 | The image reconstruction term, before its weight; zero when absent. |
masked_video_loss | f32 | The video reconstruction term, before its weight; zero when absent. |
masked_audio_loss | f32 | The audio reconstruction term, before its weight; zero when absent. |
late_interaction_loss | f32 | The late-interaction term, weighted; zero when absent. |
choices_loss | f32 | The choices term over the items that made a choice, before its weight; zero when absent. |
truncated | usize | Items cut at a row's width. |
dropped | usize | Segments dropped whole because the batch had no room for them, since the last step that read its stats. |
Clone, Copy, StructuralPartialEq, PartialEq, Debug, DefaultVideoShapeThe video slice of a batch envelope; a zero patch_side leaves video out of the run.pub struct VideoShapeMirrors uform_video_shape_t.
patch_side | usize | The square patch side one tube tiles into, in pixels. |
|---|---|---|
tube_frames | usize | Frames one tube folds ahead of its patch; zero derives two. |
frame_stride | usize | Source frames between sampled ones; zero derives one, every frame. |
channels | usize | Channels every pixel carries; zero derives three. |
Clone, Copy, StructuralPartialEq, PartialEq, Eq, Debug, Default