validation

Model validation and loss calculation utilities for single-GPU and multi-GPU setups.

Also provides helper functions for loss metrics, KL divergence, JS divergence, and similarity losses for knowledge distillation.

Classes

LowMemorySparseTensor

Functions

calculate_losses

Do model forward on each batch and calculate LM loss.

calculate_batch_outputs

calculate_batch_outputs_flash_kd

Memory-efficient teacher-similarity scoring for large-vocab models.

cosine_embedding_loss

normalized_mse_loss

mse_loss

kl_div

Kullback-Leibler Divergence for a single sample.
class LowMemorySparseTensor

Bases: object

__init__(x)
Parameters:

x (Tensor)

to(*args)
Return type:

Self

to_dense()
Return type:

Tensor

calculate_batch_outputs(hidden_states, target_hidden_states, logits, target_logits, targets, return_hidden_states, calculate_full_score_ablations, calc_on_cpu)
Parameters:

  • hidden_states (Tensor | None)

  • target_hidden_states (Tensor | None)

  • logits (Tensor)

  • target_logits (Tensor | None)

  • targets (Tensor)

  • return_hidden_states (bool)

  • calculate_full_score_ablations (bool)

  • calc_on_cpu (bool)

Return type:

dict

calculate_batch_outputs_flash_kd(hidden_states, target_hidden_states, logits, teacher_lm_head_weight, targets, temperature=1.0, ignore_index=-1, chunk_size=16384, calc_on_cpu=False)

Memory-efficient teacher-similarity scoring for large-vocab models.

Equivalent in spirit to calculate_batch_outputs() with teacher targets, but it never materializes the full [b, t, vocab] teacher logits or fp32 softmax tensors. Instead CE (lm_loss) and KD (kl_div) are computed by streaming over the vocab dimension via flash_ce_kd_loss(), while the cheap hidden-state similarity metrics (including the MIP objective cosine_embedding_loss_hidden_states) are computed as before.

Shapes: hidden_states/target_hidden_states are [b, t, d], logits is [b, t, vocab], teacher_lm_head_weight is [vocab, d], targets is [b, t]. Per-sample reductions match the originals: lm_loss and kl_div are the per-token losses averaged over t.

Parameters:

  • hidden_states (Tensor)

  • target_hidden_states (Tensor)

  • logits (Tensor)

  • teacher_lm_head_weight (Tensor)

  • targets (Tensor)

  • temperature (float)

  • ignore_index (int)

  • chunk_size (int)

  • calc_on_cpu (bool)

Return type:

dict

calculate_losses(model, dataloader, target_probs=None, return_probs=False, checkpoint_manager=None)

Do model forward on each batch and calculate LM loss.

Works on lit-llama models (single GPU) and HuggingFace models (can be multi-GPU). Does not support data-parallel.

Note

Anything related to probs and hidden states is not supported currently.

Returns:

Tuple of (outputs, None). outputs is a dict:

{
    "lm_loss": [float, ...],
    "token_accuracy_top_1": [float, ...],
    "token_accuracy_top_5": [float, ...],
    "token_accuracy_top_10": [float, ...],
}

Parameters:

  • model (Module)

  • dataloader (DataLoader)

  • target_probs (None)

  • return_probs (bool)

Return type:

tuple[dict[str, dict], None] | tuple[None, None]

cosine_embedding_loss(hidden_states, target_hidden_states)
Parameters:

  • hidden_states (Tensor)

  • target_hidden_states (Tensor)

Return type:

list[float]

kl_div(logits, target_probs, clip_epsilon=ClipEpsilon.NO_CLIP, epsilon_factor=1.0)

Kullback-Leibler Divergence for a single sample. logits: [tokens, vocab] target_probs: [tokens, vocab]

Parameters:

  • logits (Tensor)

  • target_probs (Tensor)

  • clip_epsilon (ClipEpsilon)

  • epsilon_factor (float)

Return type:

float

mse_loss(hidden_states, target_hidden_states)
Parameters:

  • hidden_states (Tensor)

  • target_hidden_states (Tensor)

Return type:

list[float]

normalized_mse_loss(hidden_states, target_hidden_states)
Parameters:

  • hidden_states (Tensor)

  • target_hidden_states (Tensor)

Return type:

list[float]