Skip to content

feat(model): add generic chunked linear cross entropy - #177

Open
maomaocun wants to merge 1 commit into
modelscope:mainfrom
maomaocun:codex/chunked-linear-ce-h100
Open

feat(model): add generic chunked linear cross entropy#177
maomaocun wants to merge 1 commit into
modelscope:mainfrom
maomaocun:codex/chunked-linear-ce-h100

Conversation

@maomaocun

Copy link
Copy Markdown

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 one
chunk at a time. The feature is disabled by default.

Scope

  • Add ModelConfig.chunked_linear_ce_chunk_size; zero keeps the native path.
  • Add the canonical CHUNKED_LINEAR_CE_* environment names, while accepting the
    historical LINEAR_CE_* names as compatibility aliases.
  • Integrate the path into the common GPTModel output/loss processing, including
    its MTP loss path. Non-causal tasks, inference, gathered output, and MuP output
    scaling remain on the existing path or fail with an explicit guard.
  • Handle TP vocab shards, sequence parallel input gradients, ignored labels,
    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_compile for all changed Python files.
  • 4 CPU unittest cases: configuration precedence, native forward/backward parity,
    legacy alias, and all-ignored-label backward.
  • 8-process NCCL TP8 forward/backward probe with the real Qwen3.6-27B
    [248320, 5120] lm_head shard.
  • Real Qwen3.6 output-head A/B shows the expected memory-bound behavior: native
    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.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant