Add fused CE + TV distance kernel - #1384
Open
kashif wants to merge 2 commits into
Open
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Adds
LigerFusedCETVDFunction: per-token cross-entropy and total variation distance against a teacher, computed in one streaming pass over the vocabulary.Distillation losses that pair a hard-label term with a distributional distance currently materialize two
(BT, V)float32 softmax temporaries. This fuses both terms and keeps onlyO(BT)state for backward, rebuilding the distributions from logits the caller already holds.ops/tvd.pyis unchanged — it consumes probabilities, so it cannot avoid those temporaries.Both terms are returned unreduced, so the caller keeps ownership of masking, per-token weighting, and normalization.
Testing Done
test/transformers/test_fused_ce_tvd.py: forward and backward parity against eager autograd in fp32 and bf16, over vocab sizes 128–131072 (multi-tile) and non-power-of-two tails, plusignore_index, teacher-receives-no-gradient, and shape validation.Exact
p == qties are covered in the backward pass, including self-distillation: the sign uses a zero subgradient at ties so the kernel matchestorch.absbackward. Swapping it for a two-wayp > qsplit fails both tie tests — the same bug class #1374 fixes inops/tvd.py.At
V=151936bf16, against a chunked and activation-checkpointed eager reference: marginal memory falls from ~0.87 MB/token to ~0 (890 MiB at 1024 tokens), and HBM traffic ~5.9x.make testto ensure correctnessmake checkstyleto ensure code stylemake test-convergenceto ensure convergencemake testreached 2741 passed / 858 skipped before stopping ontest_fused_moe.py::test_correctness, which is unrelated and simply too slow on this integrated GPU; tests ordered after it did not run. The timings above come from that same box, so the traffic ratio is the transferable number rather than the wall-clock. Please run the full suite and convergence in CI.