Training¶

olaverse-foundry ships a family of production-ready trainers that share the same config base and feature set. They fall into two groups: distillation trainers (learn from teachers) and encoder / head trainers (pretrain a base and add task heads).

Not sure which one you need?

Start from your goal, not the class list: Which trainer do I need? →

Trainer	Builds	Page
`TorchDistillTrainer`	causal LM (CE+KL, teachers live)	reference
`CachedDistillTrainer`	causal LM (cached logits, multi-GPU)	reference
`EmbeddingDistillTrainer`	embedding model (pooled MSE/cosine)	reference
`MLMTrainer`	encoder base from scratch (masked LM)	reference
`EncoderDistillTrainer`	encoder base (token-level distillation)	reference
`DistilMLMTrainer`	encoder base (distillation + MLM combined, DistilBERT-style)	reference
`ContrastiveTrainer`	retrieval embedding model (InfoNCE)	reference
`SequenceClassificationTrainer`	classification / langID / moderation head	reference
`TokenClassificationTrainer`	NER / token-classification head	reference

Distillation trainers¶

TorchDistillTrainer

Single GPU

CE + KL loss against one or many teachers
Teachers run live every step
Best for: short runs, small datasets, prototyping
Full reference →

CachedDistillTrainer

Multi-GPU via accelerate

Teachers run once; cached to disk
Subsequent epochs are free
Best for: multi-epoch training on large datasets
Full reference →

EmbeddingDistillTrainer

Single GPU

MSE / cosine loss on pooled sentence vectors
For bi-encoder and reranker distillation
Best for: embedding model compression
Full reference →

Shared feature set

All trainers

Mixed precision (bfloat16 / float16)
Gradient accumulation
LR scheduler with linear warmup
Checkpoint save / resume
Eval loop & W&B / TensorBoard logging

Shared production features¶

Mixed precision¶

Set torch_dtype in any config:

TorchTrainConfig(torch_dtype="bfloat16")   # recommended for A100/H100
TorchTrainConfig(torch_dtype="float16")    # for older GPUs
TorchTrainConfig(torch_dtype="float32")    # default — CPU safe

Gradient accumulation¶

TorchTrainConfig(
    batch_size              = 4,
    grad_accumulation_steps = 8,   # effective batch = 32
)

LR scheduler¶

TorchTrainConfig(
    lr_scheduler = "cosine",   # "constant" | "linear" | "cosine"
    warmup_steps = 500,
)

The scheduler wraps a LambdaLR over the optimizer. Linear warmup runs for warmup_steps, then the chosen schedule decays over the remaining steps. "constant" with no warmup returns None (no scheduler overhead).

Reproducible seed¶

# Set BEFORE creating your model for full reproducibility
import torch, numpy as np
torch.manual_seed(42)
np.random.seed(42)

student = MyModel()
trainer = TorchDistillTrainer(student, teachers, TorchTrainConfig(seed=42))

The trainer calls _seed_everything() at the start of train(), which re-seeds torch, numpy, and random. For the model weights to also be reproducible, seed before construction.

Checkpoint save / resume¶

# Manual save
trainer.save_checkpoint("/checkpoints/step_500")

# Auto-checkpoint every N optimizer steps
TorchTrainConfig(save_every=500, save_dir="/checkpoints/run1")

# Resume
trainer.resume_from_checkpoint("/checkpoints/run1")
result = trainer.train(dataset)

checkpoint.pt contains model weights, optimizer state, and the config dict. Checkpoints are loaded with torch.load(..., weights_only=True), so resuming never executes arbitrary pickled code.

Eval loop¶

TorchTrainConfig(eval_every=100)   # evaluate every 100 optimizer steps

result = trainer.train(train_pipe, eval_dataset=eval_pipe)
print(result["eval_losses"])   # {100: 1.23, 200: 1.18, ...}

W&B / TensorBoard logging¶

pip install "olaverse-foundry[logging]"   # W&B
pip install tensorboard                   # TensorBoard

TorchTrainConfig(
    log_backend = "wandb",        # "wandb" | "tensorboard" | "none"
    project     = "my-project",
    run_name    = "exp-001",
)

The logger silently degrades to no-op if the backend is not installed.

Config reference¶

See Config Reference → for a full table of every field across all three trainers.