class Modality(IntEnum)Values
TEXT | 0 | |
|---|---|---|
IMAGE | 1 | |
VIDEO | 2 | |
AUDIO | 3 |
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:
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.
API
Types
ModalityWhat the bytes of one training item hold, valued as uform_modality_t.
ObjectiveWhich losses a training step relaxes, valued as uform_objective_t.
NegativesWhich pairs the contrastive loss pushes apart, valued as uform_negatives_t.
PairLossFormHow a contrastive term turns a row of scores into forces, valued as uform_pair_loss_form_t.
LateInteractionSideWhich side of a late-interaction score an item takes, valued as uform_late_interaction_side_t.
ClassImageShapeThe image slice of a BatchShape; a zero patch_side leaves images out of the run.3
class ImageShapeImageShape
ImageShape(*, patch_side: int = 0, channels: int = 0)patch_side
patch_side: Final[int]channels
channels: Final[int]ClassVideoShapeThe video slice of a BatchShape; a zero patch_side leaves video out of the run.5
class VideoShapeVideoShape
VideoShape(
*,
patch_side: int = 0,
tube_frames: int = 0,
frame_stride: int = 0,
channels: int = 0,
)patch_side
patch_side: Final[int]tube_frames
tube_frames: Final[int]frame_stride
frame_stride: Final[int]channels
channels: Final[int]ClassAudioShapeThe audio slice of a BatchShape; a zero chunk_samples leaves audio out of the run.3
class AudioShapeAudioShape
AudioShape(*, chunk_samples: int = 0, sample_rate: int = 0)chunk_samples
chunk_samples: Final[int]sample_rate
sample_rate: Final[int]ClassBatchShapeWhat one device's step carries at most, whatever the modality mix, fixed at creation.9
class BatchShapeBatchShape
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
max_rows_per_device: Final[int]max_segments_per_device
max_segments_per_device: Final[int]max_columns_per_row
max_columns_per_row: Final[int]max_bags_per_device
max_bags_per_device: Final[int]max_choices_per_item
max_choices_per_item: Final[int]image
image: Final[ImageShape | None]video
video: Final[VideoShape | None]audio
audio: Final[AudioShape | None]ClassMaskingHow the masked objective corrupts a batch and weighs its heads; every zero derives.12
class MaskingMasking
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
token_mask_probability: Final[float]token_mean_span
token_mean_span: Final[float]replace_fraction
replace_fraction: Final[float]random_fraction
random_fraction: Final[float]image_mask_ratio
image_mask_ratio: Final[float]video_mask_ratio
video_mask_ratio: Final[float]audio_mask_ratio
audio_mask_ratio: Final[float]text_weight
text_weight: Final[float]image_weight
image_weight: Final[float]video_weight
video_weight: Final[float]audio_weight
audio_weight: Final[float]ClassContrastiveHow the contrastive loss scores the pooled embeddings; every zero derives.9
class ContrastiveContrastive
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
matryoshka_dimensions: Final[Sequence[int] | None]temperature_groups
temperature_groups: Final[int]inverse_temperature
inverse_temperature: Final[float]bias
bias: Final[float]negatives
negatives: Final[Negatives]ambient_repulsion_weight
ambient_repulsion_weight: Final[float]max_edges_per_device
max_edges_per_device: Final[int]pair_loss
pair_loss: Final[PairLossForm]ClassLateInteractionHow the late-interaction loss scores the token vectors of queries and documents; every zero derives.10
class LateInteractionLateInteraction
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
weight: Final[float]width
width: Final[int]inverse_temperature
inverse_temperature: Final[float]bias
bias: Final[float]max_query_span
max_query_span: Final[int]max_query_tokens_per_device
max_query_tokens_per_device: Final[int]max_documents_per_device
max_documents_per_device: Final[int]max_document_tokens_per_device
max_document_tokens_per_device: Final[int]pair_loss
pair_loss: Final[PairLossForm]ClassScheduleThe learning-rate schedule: a linear warmup to the peak, then a cosine decay to the floor.5
class ScheduleSchedule
Schedule(
*,
learning_rate: float = 0.0,
min_learning_rate: float = 0.0,
warmup_steps: int = 0,
steps: int = 0,
)learning_rate
learning_rate: Final[float]min_learning_rate
min_learning_rate: Final[float]warmup_steps
warmup_steps: Final[int]steps
steps: Final[int]ClassEdgesOne edge set over a step's items, as compressed sparse rows of buffers borrowed until the step is final.4
class EdgesEdges
Edges(
*,
offsets: Buffer | None = None,
ids: Buffer | None = None,
weights: Buffer | None = None,
)offsets
offsets: Final[Buffer | None]ids
ids: Final[Buffer | None]weights
weights: Final[Buffer | None]ClassRelationsThe contrastive targets of one step: ignored, attracted, repelled, then the run's negatives.4
class RelationsRelations
Relations(
*,
attract: _EdgesLike = None,
repel: _EdgesLike = None,
ignore: _EdgesLike = None,
)attract
attract: Final[Edges | None]repel
repel: Final[Edges | None]ignore
ignore: Final[Edges | None]ClassTrainerStepStatsWhat step reports, overwritten by every call it is passed to, every loss its world mean.12
class TrainerStepStatsstepThe optimizer step the call applied, counted from one.
step: intlearning_rateThe learning rate it applied.
learning_rate: floatlossThe total loss the step minimized, its world mean.
loss: floatcontrastive_lossThe contrastive term; 0 when absent.
contrastive_loss: floatmasked_text_lossThe text reconstruction term, unweighted.
masked_text_loss: floatmasked_image_lossThe image reconstruction term, unweighted.
masked_image_loss: floatmasked_video_lossThe video reconstruction term, unweighted.
masked_video_loss: floatmasked_audio_lossThe audio reconstruction term, unweighted.
masked_audio_loss: floatlate_interaction_lossThe late-interaction term, weighted.
late_interaction_loss: floatchoices_lossThe choices term over the items that made a choice, unweighted.
choices_loss: floattruncatedItems cut at a row's width.
truncated: intdroppedSegments dropped whole because the batch had no room for them, since the last stats read.
dropped: intClassTrainerA training run on the devices of the executor it was created on, which its own verbs block on.5
class Trainerinit_vramBuilds one replica per device and starts or resumes the run.
@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,
) -> TrainerCalls uform_trainer_init_vram.
onA receiver whose verbs block on executor.
deferredA receiver whose verbs submit to executor and answer a Deferred.
stepApplies one optimizer step over this process's items and every peer's.
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 | NoneCalls uform_trainer_step.
saveWrites the weights and the step as a checkpoint.
def save(
path: _PathLike,
/,
*,
progress: _Progress | None = None,
progress_steps: int = 0,
) -> NoneCalls uform_trainer_save.