Config Reference¶

All trainer configs share a common base (TrainConfig) extended by TorchTrainConfig which is further extended by CachedDistillConfig. EmbeddingDistillConfig is standalone with a similar set of fields.

TrainConfig (base)¶

Shared by TorchDistillTrainer and CachedDistillTrainer.

Field	Type	Default	Description
`learning_rate`	`float`	`1e-4`	AdamW learning rate
`epochs`	`int`	`3`	Number of training epochs
`batch_size`	`int`	`4`	Input batch size (informational; actual size comes from your dataset)
`alpha`	`float`	`0.3`	CE loss weight. KL weight = `1 - alpha`.
`fusion_strategy`	`str`	`"min_ce"`	How to fuse multiple teacher distributions: `"min_ce"` or `"mean_ce"` (also accepted as `"mean"`). An unrecognised name raises `ValueError`.
`top_k`	`int`	`64`	Top-k teacher logits to consider during fusion
`sparse_kl`	`bool`	`True`	Compute the distillation KL over the teacher's top-k support instead of a dense `(B, S, vocab)` target. See below.
`log_every`	`int`	`10`	Print / callback loss every N optimizer steps
`seed`	`int`	`42`	Random seed for torch, numpy, and random
`lr_scheduler`	`str`	`"constant"`	LR schedule: `"constant"` / `"cosine"` / `"linear"`
`warmup_steps`	`int`	`0`	Linear warmup steps before the main schedule
`eval_every`	`int`	`0`	Run eval loop every N steps. `0` = disabled.
`save_every`	`int`	`0`	Auto-checkpoint every N steps. `0` = disabled.
`save_dir`	`str`	`""`	Directory for auto-checkpoints

TorchTrainConfig¶

Extends TrainConfig with torch-specific fields.

Field	Type	Default	Description
`max_grad_norm`	`float`	`1.0`	Gradient clipping norm
`weight_decay`	`float`	`0.01`	AdamW weight decay
`device`	`str`	`"auto"`	`"auto"` (CUDA→MPS→CPU), `"cpu"`, `"cuda"`, `"cuda:0"`, `"mps"`
`grad_accumulation_steps`	`int`	`1`	Accumulate gradients over N batches before stepping
`torch_dtype`	`str`	`"float32"`	Mixed precision: `"float32"` / `"bfloat16"` / `"float16"`
`log_backend`	`str`	`"none"`	Logging backend: `"wandb"` / `"tensorboard"` / `"none"`
`run_name`	`str`	`""`	W&B or TensorBoard run name
`project`	`str`	`"olaverse-foundry"`	W&B project name

CachedDistillConfig¶

Extends TorchTrainConfig.

Field	Type	Default	Description
`cache_dir`	`str \\| None`	`None`	Directory to save/load `.npz` logit cache files. `None` = in-memory only.
`cache_top_k`	`int`	`64`	Top-k logits to cache per token position
`use_accelerate`	`bool`	`True`	Use `accelerate.Accelerator` for DDP/FSDP. Falls back to plain torch if not installed.

EmbeddingDistillConfig¶

Standalone config for EmbeddingDistillTrainer.

Field	Type	Default	Description
`loss`	`str`	`"cosine"`	Loss function: `"cosine"` or `"mse"`
`pool`	`str`	`"mean"`	Pooling: `"mean"` or `"cls"`
`normalize`	`bool`	`True`	L2-normalise embeddings before loss computation
`temperature`	`float`	`0.05`	Scaling factor applied before cosine similarity
`learning_rate`	`float`	`2e-5`	AdamW learning rate
`epochs`	`int`	`3`	Training epochs
`weight_decay`	`float`	`0.01`	AdamW weight decay
`max_grad_norm`	`float`	`1.0`	Gradient clipping norm
`device`	`str`	`"auto"`	Device selection
`log_every`	`int`	`10`	Log every N steps
`lr_scheduler`	`str`	`"constant"`	`"constant"` / `"cosine"` / `"linear"`
`warmup_steps`	`int`	`0`	Linear warmup steps
`grad_accumulation_steps`	`int`	`1`	Gradient accumulation steps
`torch_dtype`	`str`	`"float32"`	Mixed precision dtype
`eval_every`	`int`	`0`	Eval every N steps
`save_every`	`int`	`0`	Auto-checkpoint every N steps
`save_dir`	`str`	`""`	Checkpoint directory
`log_backend`	`str`	`"none"`	Logging backend
`run_name`	`str`	`""`	Run name
`project`	`str`	`"olaverse-foundry"`	W&B project name
`seed`	`int`	`42`	Random seed

Inheritance diagram¶

TrainConfig
└── TorchTrainConfig
    └── CachedDistillConfig

EmbeddingDistillConfig  (standalone, same fields)

LR scheduler details¶

`lr_scheduler`	`warmup_steps`	Behaviour
`"constant"`	`0`	Fixed LR throughout. No scheduler object created.
`"constant"`	`> 0`	Linear warmup, then constant LR.
`"cosine"`	any	Linear warmup, then cosine decay to 0.
`"linear"`	any	Linear warmup, then linear decay to 0.

total_steps for the decay phase is computed as:

total_steps = ceil(len(dataset) × epochs / grad_accumulation_steps)

For streaming datasets (len() raises TypeError), total_steps defaults to 0 which keeps LR constant after warmup. Override with trainer.train(..., total_steps=N).

Fusion strategies¶

`fusion_strategy`	Description
`"min_ce"`	Per-position, pick the teacher distribution with the lowest CE against ground truth. Encourages the student to learn from the most confident teacher.
`"mean_ce"`	Weighted average of all teacher distributions (by `teacher.weight`). Smooths over teacher disagreements.

Other configs¶

The encoder, head, and quantization trainers have their own configs (same shared production fields — mixed precision, grad accumulation, LR schedule, eval, checkpointing). Each is documented on its page:

Config	Trainer	Page
`MLMConfig`	`MLMTrainer`	MLM pretraining
`EncoderDistillConfig`	`EncoderDistillTrainer`	Encoder distillation
`DistilMLMConfig`	`DistilMLMTrainer`	Distillation + MLM
`ContrastiveConfig`	`ContrastiveTrainer`	Contrastive retrieval
`HeadTrainConfig`	`SequenceClassificationTrainer` / `TokenClassificationTrainer`	Task heads
`QATConfig`	`prepare_qat`	Quantization

`sparse_kl` — why it defaults to `True`¶

A teacher returns top-k distributions: (B, S, K) indices and probabilities, ~8MB at B=8, S=2048, K=64. The dense path scattered those into a full (B, S, vocab) array before computing the KL. At a 152k-token vocabulary that is ~10GB per teacher per step — allocated in numpy, then copied to the device as a second tensor of the same size. The teacher signal did not get richer on the way; only the zeros did.

KL(T‖S) = Σ_v T(v)·(log T(v) − log S(v)) has no contribution from any v where T(v) = 0, so the dense sum is doing arithmetic on zeros. The sparse path gathers the student's log-probs at exactly the teacher's top-k indices and sums there.

Measured on CPU, single teacher, top_k=64:

shape	dense	sparse
`B=4 S=256 V=32k`	131 MB, 3.5 s	0.9 MB, 0.08 s	154× less memory
`B=4 S=256 V=152k`	622 MB, 12.5 s	0.9 MB, 0.22 s	730× less memory

(Memory is the numpy allocation; the dense path also copies the same array to the device.)

The two paths differ very slightly

The dense path adds 1e-9 to every vocab entry before renormalising — across 152k tokens that smears ~1.5e-4 of probability mass over the whole vocabulary and pulls the target fractionally off what the teacher said. The sparse path renormalises the top-k mass over its own support, which is the standard top-k distillation formulation. Against a dense reference without that smoothing, the two agree to float32 precision.

Set sparse_kl=False only to reproduce a run made before this existed.

Collision semantics¶

If the same student token appears twice in one fused support — two teachers nominating it, or a MinEDAlignment mapping two teacher tokens onto one student token — those probabilities are mass on the same token and are summed.

IdentityAlignment previously used plain assignment here while EMAlignment used np.add.at, so the two disagreed and the identity path silently dropped the earlier probability. Both now sum.