flash_kd

Memory-efficient (vocab-chunked) cross-entropy + knowledge-distillation losses.

This is a simplified, forward-only distillation of the flash-KD kernel: it computes per-token cross-entropy (student vs. ground-truth labels) and KL divergence KL(P_teacher || P_student) without ever materializing the full [num_tokens, vocab_size] logits/probability tensors in fp32.

Motivation: for large-vocab models (e.g. Qwen3.5 with vocab ~248k) the naive similarity-metric path materializes several full-vocab fp32 tensors per sample (softmax/log_softmax run in fp32 under autocast), which OOMs at long sequence lengths. Here we stream over the vocab dimension in chunks and keep only O(num_tokens) running statistics, using the online-softmax (running max + rescale) trick so the result is numerically identical to the full computation.

Only the forward pass is implemented (scoring runs under torch.no_grad); there are no gradients, no Triton, and no tensor-parallel sharding — the loss is computed on a single rank that holds the full vocab.

Functions

flash_ce_kd_loss

Compute per-token CE and KD losses by streaming over the vocab dimension.
flash_ce_kd_loss(student_logits, teacher_hidden, teacher_lm_head_weight, labels, *, temperature=1.0, ignore_index=-1, chunk_size=16384, upcast=True)

Compute per-token CE and KD losses by streaming over the vocab dimension.

The student side uses already-materialized student_logits (sliced per chunk); the teacher side recomputes its logits chunk-by-chunk as teacher_hidden @ teacher_lm_head_weight[chunk].T so the full teacher logits tensor is never materialized.

Parameters:

  • student_logits (Tensor) – [N, V] student logits (N = number of tokens).

  • teacher_hidden (Tensor) – [N, D] teacher hidden states feeding its LM head.

  • teacher_lm_head_weight (Tensor) – [V, D] teacher LM-head weight (no bias).

  • labels (Tensor) – [N] ground-truth next-token ids; ignore_index entries get a CE of 0 (matching F.cross_entropy(..., reduction="none")).

  • temperature (float) – KD softmax temperature (1.0 reproduces the plain kl_div scoring metric: mean-per-token KL(P_T || P_S)).

  • ignore_index (int) – label value to exclude from cross-entropy.

  • chunk_size (int) – number of vocab columns processed per step (memory knob).

  • upcast (bool) – accumulate in fp32 (recommended; matches the autocast path which runs softmax/log_softmax in fp32).

Returns:

(ce_per_token, kd_per_token), each shape [N] in the accumulation dtype. ce_per_token is the standard cross-entropy; kd_per_token is KL(P_T || P_S) per token (NOT multiplied by temperature**2).

Return type:

tuple[Tensor, Tensor]