JavaScript · Training

Training

import * as uform from 'uform';import { Trainer } from 'uform'; 4.0.0-rc.1 · 3ee5e534

Training runs on GPUs through the Node addon. A run’s processes join a mesh once, and a trainer borrows it until freed.

import { Executor, Mesh, Trainer, Objective, Modality, PairLossForm, LateInteractionSide } from "@unum-cloud/uform";

const gpus = Executor.forVram();
const mesh = Mesh.bootstrap(gpus, { hostDomainIndex: 0, hostDomainCount: 2, directory: "/shared", job: "run-1" });
const trainer = Trainer.initVram(
  gpus,
  { depth: 12, hidden: 512, batch: { maxRowsPerDevice: 64 }, objective: Objective.Joint, schedule: { steps: 100000 } },
  mesh,
);
const { stats } = trainer.step({
  modalities: [Modality.Text, Modality.Image],
  contents: ["a small red panda", jpeg],
  relations: { attract: { offsets: [0, 1, 2], ids: [1, 0] } },
});
await trainer.deferred().save("checkpoint.safetensors"); // or: call.wait()

A lateInteraction term trains the token head over items whose sides name them queries or documents, and a choices term trains the byte decoder on each item’s legal bytes and their target probabilities:

const reranker = Trainer.initVram(gpus, {
  depth: 12,
  hidden: 512,
  objective: Objective.Contrastive,
  schedule: { steps: 100000 },
  batch: { maxRowsPerDevice: 64, maxChoicesPerItem: 16 },
  lateInteraction: { weight: 1, pairLoss: PairLossForm.Ranked },
  choices: { weight: 1 },
});
reranker.step({
  modalities: [Modality.Text, Modality.Text],
  contents: ["red panda diet", "Red pandas eat bamboo."],
  sides: [LateInteractionSide.Query, LateInteractionSide.Document],
  choices: ["yn", ""],
  choiceTargets: [Float32Array.of(0.9, 0.1), new Float32Array()],
  relations: { attract: { offsets: [0, 1, 2], ids: [1, 0] } },
});

The choices take the shape textChoose takes, and their targets the shape of the Probabilities it answers, so one model’s probabilities can train another as they land. Every batch, contents and choices alike, also takes a packed { data, offsets } tape, which the Node addon lends to the library without a copy when its offsets are a BigUint64Array; an array of entries is first gathered into one tape, the one copy that path makes.

API

Types

BatchShapeWhat one device's step carries at most.
uform.ts:98 ↗
interface BatchShape

Mirrors uform_batch_shape_t; every absent field reads zero.

Fields

maxRowsPerDevice?number
maxSegmentsPerDevice?number
maxColumnsPerRow?number
maxBagsPerDevice?number
maxChoicesPerItem?numberLegal bytes one item offers to choose from at most; absent makes no choices.
image?{ patchSide?: number; channels?: number }
video?{ patchSide?: number; tubeFrames?: number; frameStride?: number; channels?: number }
audio?{ chunkSamples?: number; sampleRate?: number }
MaskingHow the masked objective corrupts a batch and weighs its heads.
uform.ts:112 ↗
interface Masking

Mirrors uform_masking_t; absent fields derive.

Fields

tokenMaskProbability?number
tokenMeanSpan?number
replaceFraction?number
randomFraction?number
imageMaskRatio?number
videoMaskRatio?number
audioMaskRatio?number
textWeight?number
imageWeight?number
videoWeight?number
audioWeight?number
ContrastiveHow the contrastive loss scores the pooled embeddings.
uform.ts:128 ↗
interface Contrastive

Mirrors uform_contrastive_t; absent fields derive.

Fields

matryoshkaDimensions?number[]
temperatureGroups?number
inverseTemperature?number
bias?number
negatives?Negatives
ambientRepulsionWeight?number
maxEdgesPerDevice?number
pairLoss?PairLossFormThe pair loss over each anchor's row; absent keeps sigmoid.
LateInteractionHow the late-interaction loss scores token vectors.
uform.ts:142 ↗
interface LateInteraction

Mirrors uform_late_interaction_t; absent fields derive.

Fields

weight?numberThe loss weight; absent trains no token head.
width?numberThe token vectors' width, a multiple of 8; absent derives 128.
inverseTemperature?number
bias?number
maxQuerySpan?number
maxQueryTokensPerDevice?number
maxDocumentsPerDevice?number
maxDocumentTokensPerDevice?number
pairLoss?PairLossFormThe pair loss over each query's row of hard MaxSim scores; absent keeps ranked.
ScheduleThe learning-rate schedule.
uform.ts:167 ↗
interface Schedule

Mirrors uform_schedule_t.

Fields

learningRate?number
minLearningRate?number
warmupSteps?number
stepsnumberSteps the run takes, which the cosine spans; required.
TrainerConfigA training run's geometry and objectives.
uform.ts:176 ↗
interface TrainerConfig

Mirrors uform_trainer_init_vram_t.

Fields

resume?stringThe checkpoint to continue from; absent starts from seed.
depthnumberTransformer blocks stacked.
hiddennumberEmbedding width, a multiple of 128.
batchBatchShape
objectiveObjective
masking?Masking
contrastive?Contrastive
lateInteraction?LateInteraction
choices?ChoicesRead when its weight is positive, which needs the batch's maxChoicesPerItem.
scheduleSchedule
seed?number | bigintThe base of every keyed draw.
EdgesOne edge set over a step's items, as compressed sparse rows.
uform.ts:196 ↗
interface Edges

Mirrors uform_edges_view_t.

Fields

offsetsUint32Array | number[]The step's count + 1 offsets into ids.
idsUint32Array | number[]Item indices each item links to, ascending within each run.
weights?Float32Array | number[]The weight of each edge; absent weighs every edge one.
TrainerBatchThis process's share of one step.
uform.ts:206 ↗
interface TrainerBatch

Mirrors the inputs of uform_trainer_step_t.

Fields

modalitiesreadonly Modality[]The modality of each item.
contentsreadonly (string | Uint8Array)[] | Tape<Uint8Array>Each item's bytes, UTF-8 text or an encoded file, or a tape of them.
bags?Uint32Array | number[]The bag of each item; items sharing a bag land on one device.
temperatureGroups?Uint32Array | number[]The temperature group of each item.
sides?readonly LateInteractionSide[]The late-interaction side of each item; absent makes every item LateInteractionSide.None.
choices?readonly (string | Uint8Array)[] | Tape<Uint8Array>Each item's legal choices, one byte each, as textChoose takes them; an item offering none makes no choice.
choiceTargets?readonly Float32Array[] | Tape<Float32Array>The target probability of each choice, an entry per item as long as its choices and summing to one, as textChoose answers its Probabilities; required with choices.
relations?{ attract?: Edges; repel?: Edges; ignore?: Edges }The contrastive targets.
TrainerStepConfigWhere a step lands the embeddings it scored, when a caller wants them.
uform.ts:231 ↗
interface TrainerStepConfig

Fields

arena?ArenaThe memory the embeddings land in; absent skips the copy.
outputsScalar?ScalarThe scalar they land as; absent lands f32.
SteppedWhat one step reports, and the embeddings it scored when an arena was given.
uform.ts:239 ↗
interface Stepped

Fields

statsCompiledStepped["stats"]
embeddings?Rows
ModalityWhat the bytes of one training item hold.
uform.ts:86 ↗
enum Modality

Values

Text&quot;text&quot;
Image&quot;image&quot;
Video&quot;video&quot;
Audio&quot;audio&quot;
ObjectiveWhich losses a training step relaxes.
uform.ts:98 ↗
enum Objective

Values

Masked&quot;masked&quot;Masked reconstruction alone.
Contrastive&quot;contrastive&quot;The pairwise contrastive loss alone.
Joint&quot;joint&quot;Both, joined at the final-norm states.
NegativesWhich pairs the contrastive loss pushes apart.
uform.ts:112 ↗
enum Negatives

Values

Ambient&quot;ambient&quot;Repel edges, and every pair no edge names.
Named&quot;named&quot;Repel edges alone.
PairLossFormHow a contrastive term turns a row of scores into forces, numbered as uform_pair_loss_form_t numbers them.
uform.ts:135 ↗
enum PairLossForm

Values

Derived0The term's own form: sigmoid for pooled embeddings, ranked for token vectors.
Balanced1The sigmoid cross-entropy of every cell, its attract and repel halves normalized apart.
Ranked2The sigmoid cross-entropy of every positive's score over every negative's, where the bias cancels.
Sigmoid3SigLIP's sigmoid cross-entropy of every cell, each repel weighed like an attract.
LateInteractionSideWhich side of a late-interaction score an item takes, numbered as uform_late_interaction_side_t numbers them.
uform.ts:153 ↗
enum LateInteractionSide

Values

None0The item takes no part in late interaction.
Query1The item's tokens score the documents it attracts and repels.
Document2The item's tokens are scored by the queries naming it.
QueryAndDocument3The item plays both sides.
ClassTrainerA training run on a VRAM executor's devices, every verb blocking on that executor; deferred() submits them there instead.5
train.ts:340 ↗
class Trainer extends TrainerVerbs<"blocking"> implements Disposable
initVramBuilds one replica per device and starts or resumes the run.
train.ts:353 ↗
static initVram(
    executor: Executor,
    config: TrainerConfig,
    mesh?: Mesh,
): Trainer

Calls uform_trainer_init_vram.

Parameters

executorExecutorA VRAM executor, every device of which holds one replica.
configTrainerConfigThe geometry, objectives and schedule.
meshMeshThe mesh to borrow until the trainer is freed; absent runs this process alone.
deferredThe same verbs, each submitted to the trainer's executor and answered as a Deferred.
train.ts:367 ↗
deferred(): TrainerDeferredOn

Binding-only.

freeFrees the trainer and lets go of its mesh.
train.ts:372 ↗
free(): void

Calls uform_trainer_free.

stepApplies one optimizer step over this process's items and every peer's.
train.ts:266 ↗

Inherited from TrainerVerbs

step(
    batch: TrainerBatch,
    config: TrainerStepConfig = {},
): Answer<Receiving, Stepped>

Calls uform_trainer_step.

Parameters

batchTrainerBatchThis process's items, their modalities, bags and relations.
configTrainerStepConfigWhere the scored embeddings land.
saveWrites the weights and the step as a safetensors checkpoint; only process zero writes.
train.ts:315 ↗

Inherited from TrainerVerbs

save(path: string): Answer<Receiving, void>

Calls uform_trainer_save.

Parameters

pathstringThe destination, written whole and renamed into place.