flash_kld

Exact memory-bounded CE and KLD from hidden states and LM heads.

Classes

TrainingFlashKLD

Compute exact CE and KLD without retaining sequence-by-vocabulary logits.
class TrainingFlashKLD

Bases: Module

Compute exact CE and KLD without retaining sequence-by-vocabulary logits.

The student LM head is applied once per token chunk. CE and KLD share that projection, while a disabled KLD term avoids the teacher projection entirely. Vocabulary-sharded DTensors remain sharded throughout the loss.

Parameters:

  • token_chunk_size – Maximum number of tokens projected at once.

  • temperature – KLD softmax temperature.

  • ignore_index – Label value excluded from both CE and KLD.

  • fp32_upcast – Compute probability reductions in float32.

  • checkpoint_chunks – Recompute projections during backward to bound activation memory.

__init__(*, token_chunk_size=128, temperature=1.0, ignore_index=-100, fp32_upcast=True, checkpoint_chunks=True)
Parameters:

  • token_chunk_size (int)

  • temperature (float)

  • ignore_index (int)

  • fp32_upcast (bool)

  • checkpoint_chunks (bool)

Return type:

None

forward(student_hidden, student_head, labels, *, teacher_hidden=None, teacher_project=None, compute_ce=True, compute_kd=True, num_label_tokens=None)

Return independently normalized exact CE and KLD scalars.

Parameters:

  • student_hidden (Tensor)

  • student_head (Module)

  • labels (Tensor)

  • teacher_hidden (Tensor | None)

  • teacher_project (Callable[[Tensor, Tensor], Tensor] | None)

  • compute_ce (bool)

  • compute_kd (bool)

  • num_label_tokens (int | Tensor | None)

Return type:

tuple[Tensor, Tensor]