feat(model): add generic chunked linear cross entropy - #177
Open
maomaocun wants to merge 1 commit 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
Add a model-independent, memory-bounded Chunked Linear Cross Entropy path for
Megatron GPT models.
The path computes the vocab-parallel output projection and cross entropy in
supervised-token chunks. It does not retain the full
[local_tokens, vocab_partition]logits tensor in the autograd graph; backward recomputes onechunk at a time. The feature is disabled by default.
Scope
ModelConfig.chunked_linear_ce_chunk_size; zero keeps the native path.CHUNKED_LINEAR_CE_*environment names, while accepting thehistorical
LINEAR_CE_*names as compatibility aliases.GPTModeloutput/loss processing, includingits MTP loss path. Non-causal tasks, inference, gathered output, and MuP output
scaling remain on the existing path or fail with an explicit guard.
tail padding, and all-ignored batches without creating a detached loss.
This is a GPTModel-level optimization and is not tied to DeepSeek or any one
model adapter. Context/pipeline/expert parallelism continue to use their existing
groups; TP reductions are performed on the model TP group. MTP acceptance counts
are zero when logits are intentionally not materialized.
Validation
python3 -m py_compilefor all changed Python files.legacy alias, and all-ignored-label backward.
[248320, 5120]lm_head shard.full-logits CE OOMs at 256K tokens on 8xL20Z 80 GiB, while chunked CE completes.
The companion ms-swift documentation/benchmark is tracked in
endless-frontier/ms-swift#23.