Repository navigation
SFT: device-agnostic chunked linear+CE loss and --gradient-checkpointing forwarding - #330
Merged
Merged
Conversation
michaellchung
force-pushed
the
resume-version-account
branch
from
November 3, 2026 02:05
dc0309e to
a9212e7
Compare
michaellchung
force-pushed
the
sft-chunked-ce
branch
from
November 3, 2026 02:05
94414e6 to
848aaac
Compare
michaellchung
force-pushed
the
resume-version-account
branch
from
November 3, 2026 03:29
a9212e7 to
f424f86
Compare
michaellchung
force-pushed
the
sft-chunked-ce
branch
from
November 3, 2026 03:29
848aaac to
f96f591
Compare
michaellchung
force-pushed
the
resume-version-account
branch
from
November 3, 2026 04:09
f424f86 to
38eee6a
Compare
michaellchung
force-pushed
the
sft-chunked-ce
branch
from
November 3, 2026 04:09
f96f591 to
5b677b5
Compare
--kernel-backend chunked runs the causal LM to its final hidden states and computes the cross-entropy over --loss-chunk-tokens rows at a time under activation checkpointing, so full-sequence logits are never materialized in forward or backward. Semantics match the native sft_loss. Works on any device family; the launcher forwards the chunk size to learners. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
The flag defaults to auto and is passed only when set explicitly, so default learner commands are unchanged. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
michaellchung
force-pushed
the
sft-chunked-ce
branch
from
November 3, 2026 04:53
5b677b5 to
128e1ef
Compare
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
--kernel-backend chunked: a device-agnostic linear+cross-entropy SFT loss. The causal LM runs to its final hidden states and the loss computeshidden @ lm_head.weight^Tplus CE over--loss-chunk-tokenstarget rows at a time under non-reentrant activation checkpointing, so full-sequence[seq, vocab]logits are never materialized in forward or backward.sft_loss(label shift, assistant mask, fractional token weights, FP32 CE, token-sum reduction, target count). Tied embeddings are supported; soft-capped/scaled heads, LoRA onlm_head, quantized bases and non-DDP sharding are refused.--loss-chunk-tokensand, when set explicitly,--gradient-checkpointingto torch learners; default learner commands are unchanged.docs/A100_KERNELS.mdand a long-context section indocs/ASCEND.md.Tests
tests/test_chunked_ce.py: loss and gradient parity with the native loss (masks, fractional weights, tied embeddings), refusal cases.tests/test_island_backend.py: launcher forwarding of the new flags.🤖 Generated with Claude Code