Repository navigation

Fix the Top-K KD loss running out of memory at long sequences with large vocabularies - #3080

Merged

kevalmorabia97 merged 3 commits into

mainfrom

kmorabia/topk-kd-loss-memory

Oct 28, 2026

Merged

Fix the Top-K KD loss running out of memory at long sequences with large vocabularies#3080
kevalmorabia97 merged 3 commits into
mainfrom
kmorabia/topk-kd-loss-memory

Commits on Oct 28, 2026
- Compute the Top-K KD loss's full-vocab normalizers without fp32 copies of the logits
  
  kevalmorabia97
  and
  claude
  committed
  4c1a801
- Keep the Top-K KD loss exact for logits with a large common offset
  
  kevalmorabia97
  and
  claude
  committed
  c6d9e53
- Keep the Top-K KD loss working for empty batch or sequence inputs
  
  kevalmorabia97
  and
  claude
  committed
  685afb4