feat(loss): chunked masked CE, packing-mask preservation, scoped MoE AC#2996
Draft
yuhezhang-ai wants to merge 8 commits into
Draft
feat(loss): chunked masked CE, packing-mask preservation, scoped MoE AC#2996yuhezhang-ai wants to merge 8 commits into
yuhezhang-ai wants to merge 8 commits into
Conversation
Computing masked cross-entropy on the last pipeline stage upcasts the full [N, V] logits to fp32 and additionally saves cross_entropy's fp32 log-softmax for backward, which dominates the loss-side memory peak for large-vocabulary models at long context. Add an opt-in chunk_size to MaskedCrossEntropy that computes the identical sum-reduced loss one fp32 [chunk_size, V] slice at a time: - forward: per-chunk fp32 logsumexp + label-logit gather accumulated into a scalar; only the original-dtype logits and labels are saved for backward. - backward: recomputes the softmax per chunk and, with inplace_grad=True (default for the chunked path), writes the gradient into the logits storage, removing a further full [N, V] allocation at the backward peak. The default chunk_size=None keeps the existing full-tensor behavior unchanged. CPU unit tests assert loss and gradient parity against the full-tensor path and F.cross_entropy, input validation, and the zero num_label_tokens guard; a CUDA test covers bf16 logits and grad/logits storage aliasing under inplace_grad. Signed-off-by: Yuhe Zhang <yuhez@nvidia.com>
Newer Transformers (5.x) preprocesses 2D attention masks before dispatching attention and can coerce the integer indexed masks produced by neat packing (1, 2, ... per packed document, 0 = padding) into bool masks. That erases the document boundaries the patched _get_unpad_data needs to build per-document cu_seqlens, silently collapsing packed samples into one attention span. Patch transformers.masking_utils._preprocess_mask_arguments from configure_packing so non-bool 2D masks are returned as-is for flash_attention_2/flash_attention_3, probing the installed Transformers return arity so the shim tracks minor signature changes. Also add qwen3 to the list of model modules whose create_causal_mask is bypassed for packing. Unit test asserts the indexed mask is passed through untouched with the correct tuple arity, with and without position_ids. Signed-off-by: Yuhe Zhang <yuhez@nvidia.com>
…E MLP Full-block activation checkpointing on MoE models recomputes the expert dispatch (permute/all-to-all/grouped GEMM) in backward, which is the most expensive part of the block to recompute while the MoE MLP is also the memory consumer that least needs checkpointing once its dispatch is saved. Add two scoped modes to distributed.activation_checkpointing, consumed by the MoE parallelizer (parallelize_model -> apply_ac): - "non_moe": checkpoint attention, layer norms, and dense MLPs per submodule, leaving the MoE MLP (router + expert dispatch) uncheckpointed. - "non_moe_no_attn": additionally leave self-attention uncheckpointed. Submodule wrapping is transparent to block forwards since dispatch on wrapped modules goes through unwrap_checkpoint_wrapper (#2955). Scoped modes never checkpoint the MoE MLP, so the router-recompute warning is skipped for them. The strings survive _normalize_activation_checkpointing like "selective" does, and existing boolean/"selective" behavior is unchanged. Unit tests cover parser normalization (including hyphen/case variants), submodule wrapping and MoE/self-attention exclusion, and the router-warning suppression. Signed-off-by: Yuhe Zhang <yuhez@nvidia.com>
akoumpa
reviewed
Jul 9, 2026
| return chunk_size | ||
|
|
||
|
|
||
| class _ChunkedMaskedCESum(torch.autograd.Function): |
Contributor
There was a problem hiding this comment.
Hi @yuhezhang-ai we already have a chunked ce, maybe you want to update it instead?
yuhezhang-ai
changed the title
Route ChunkedCrossEntropy's default reduction="sum" path through the memory-efficient _ChunkedMaskedCESum kernel shared with MaskedCrossEntropy(chunk_size=...), so chunking now also bounds backward memory: only the original-dtype logits are saved for backward (the legacy per-chunk compiled cross_entropy loop additionally saved a full fp32 copy of the logits), the softmax is recomputed one fp32 chunk at a time, and inplace_grad=True reuses the logits storage as the grad buffer. On [16384, 128256] bf16 logits (3.9 GiB) peak memory beyond the logits drops from 7.8 GiB to 1.5 GiB for forward+backward, with grads matching the legacy path to ~2e-6. Non-"sum" reductions keep the legacy loop unchanged. Also expose chunk_size/inplace_grad on MaskedCrossEntropyConfig so the typed-config and registry-name paths can enable the chunked kernel. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Yuhe Zhang <yuhez@nvidia.com>
…ped AC The scoped activation-checkpointing modes wrap every non-MoE block.mlp in a CheckpointWrapper, but the _mlp helpers of glm4_moe, glm4_moe_lite, glm_moe_dsa, ling_v2, hy_v3, and hy_mt2 dispatched on isinstance(self.mlp, MLP/MoE) without unwrapping, so dense-MLP layers (e.g. GLM first_k_dense_replace) crashed with AssertionError at the first forward. Apply the unwrap_checkpoint_wrapper pattern already used by qwen3_moe and qwen3_next (#2955), and add a CPU regression test that builds REAL dense blocks for all six families, applies the production apply_ac scoped wrapping, and runs a forward+backward through the wrapped _mlp (the DummyBlock-based parallelizer tests mock both the blocks and the wrapper and cannot catch this). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Yuhe Zhang <yuhez@nvidia.com>
The scoped modes (non_moe/non_moe_no_attn) are implemented only by the MoE parallelizer, which is selected when ep_size > 1. Without EP the mode string reached the dense FSDP2 path, where any truthy value was coerced to full per-block checkpointing - silently checkpointing the MoE MLP the user asked to exempt. parse_distributed_section now raises a clear ValueError when a scoped mode is configured with ep_size <= 1, and FSDP2Manager.parallelize rejects string modes other than 'selective' instead of coercing them to True. Also document that non_moe_no_attn only exempts self_attn; linear attention (linear_attn, e.g. Qwen3-Next GatedDeltaNet) is still checkpointed. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Yuhe Zhang <yuhez@nvidia.com>
inplace_grad=True reuses the logits storage as the gradient buffer, destroying the logits values in backward - a silent behavior change for the pre-existing public ChunkedCrossEntropy (whose sum path never mutated the logits) and a version-counter crash if the logits feed another autograd branch. Make it keyword-only and default False on ChunkedCrossEntropy, MaskedCrossEntropy, and MaskedCrossEntropyConfig so the mutation is strictly opt-in via config. Mark _ChunkedMaskedCESum.backward @once_differentiable so double-backward (create_graph=True) fails loudly instead of returning wrong higher-order gradients, document the retain_graph re-backward behavior under inplace_grad=True, and add CPU tests for the non-mutating defaults and the double-backward error. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Yuhe Zhang <yuhez@nvidia.com>
When transformers.masking_utils or _preprocess_mask_arguments is unavailable, or the return-arity probe fails, the packing shim silently no-opped - and on future Transformers drift packed samples would silently attend across document boundaries. Log a warning including the installed transformers version on each of those paths (behavior otherwise unchanged), document the probe's tensor layouts in a Google-style Args docstring, and add a unit test asserting the warning fires with the version string. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Yuhe Zhang <yuhez@nvidia.com>
Contributor
Author
|
/ok to test c7af888 |
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
2 participants
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
Three independent training-efficiency/correctness changes, each with CPU-runnable unit tests:
feat(loss): chunked fp32 path onMaskedCrossEntropy(chunk_size) — computes theidentical sum-reduced masked CE one fp32
[chunk_size, V]slice at a time instead ofupcasting the full
[N, V]logits, and (withinplace_grad=True) writes the backwardgradient into the logits storage instead of allocating a second
[N, V]buffer.fix(models): preserve indexed packing masks for flash attention — newer Transformers(5.x) mask preprocessing coerces the integer indexed masks produced by neat packing
(
1, 2, ...per packed document,0= padding) to bool before dispatch, erasing thedocument boundaries
_get_unpad_dataneeds to build per-documentcu_seqlens. Packedsamples then silently attend across document boundaries.
feat(moe): scoped activation-checkpointing modes —distributed.activation_checkpointing: non_moe | non_moe_no_attncheckpoint per submodule on the MoE parallelizer path, leavingthe MoE MLP (router + expert dispatch) — and optionally self-attention — uncheckpointed.
Motivation
MaskedCrossEntropymaterializes a fullfp32 copy of the logits plus
cross_entropy's saved fp32 log-softmax — 2-3xS * V * 4Bon top of the lm_head output and optimizer state. For large vocabularies (V > 1.5e5) at
long context this dominates the last-stage memory peak (multi-GiB per microbatch) and
scales linearly with context length. Fused linear CE avoids this but is not usable under
pipeline parallelism, where the schedule calls
loss(output, target)on alreadymaterialized logits. The chunked path bounds the fp32 transient at
chunk_size * V * 4Bregardless of sequence length;
inplace_gradadditionally removes a full[N, V]allocation at the backward peak. Distinct from the existing
ChunkedCrossEntropy, whichchunks the forward but still lets autograd save per-chunk fp32 activations for backward.
flash_attention_2/flash_attention_3with current Transformers; without it, packingsilently degrades to one attention span per row.
(permute / all-to-all / grouped GEMM) in backward — the most expensive recompute in the
block. The scoped modes keep AC on attention/norms/dense MLPs while trading the MoE MLP's
activation memory for skipping the dispatch recompute;
non_moe_no_attnalso keepsFA activations live. Complements fix(distributed): fix and scope vlm activation checkpointing #2840 (which scopes the generic FSDP2/VLM path; the MoE
parallelizer path had no scoped modes) and builds on fix(moe): preserve MLP dispatch through checkpoint wrappers #2955 (submodule checkpoint wrappers
are transparent to block-forward dispatch via
unwrap_checkpoint_wrapper).What's included
nemo_automodel/components/loss/masked_ce.py_ChunkedMaskedCESumautograd function: per-chunk fp32 logsumexp + label-gather forward(saves only the original-dtype logits + labels), per-chunk softmax recompute backward,
optional in-place gradient into the logits storage.
MaskedCrossEntropy(chunk_size=None, inplace_grad=True): opt-in;chunk_size=Nonekeeps the existing behavior bit-for-bit. Requires
reduction="sum"andfp32_upcast=True; validates shapes and re-validates a mutatedchunk_sizeat forward.nemo_automodel/components/models/common/packing.py_patch_preprocess_mask_arguments_for_packing(), installed fromconfigure_packing:returns non-bool 2D masks as-is for FA2/FA3, probing the installed Transformers return
arity so the shim tracks signature changes; no-op when
masking_utilsis absent (olderTransformers). Also adds
qwen3to thecreate_causal_maskpacking-patch module list.nemo_automodel/components/moe/parallelizer.py+nemo_automodel/recipes/_dist_utils.pyapply_ac(..., activation_checkpointing=...)with scoped modesnon_moe/non_moe_no_attn(wrapsself_attn/linear_attn/norms/densemlpper submodule,skips
MoEmodules and, forno_attn,self_attn); weight-tied MTP blocks stayexcluded; router-recompute warning suppressed for scoped modes (the router is never
checkpointed there).
_normalize_activation_checkpointingaccepts the two mode strings (case/hyphentolerant), like
"selective".Proof
tests/unit_tests/loss/test_masked_ce.pyasserts loss andgradient parity of the chunked path against both the stock full-tensor path and
F.cross_entropy(reduction="sum")atatol=1e-7(fp32, withignore_indexholes and amask), for
inplace_gradin {False, True}, plus input-validation andnum_label_tokens=0zero-gradient tests. A CUDA test covers bf16 logits and asserts thegradient aliases the logits storage under
inplace_grad.tests/unit_tests/models/common/test_packing.pyasserts the indexed mask ispassed through unmodified with the correct tuple arity, with and without
position_ids.tests/unit_tests/moe/test_parallelizer.pyasserts submodule wrapping,MoE-MLP and self-attention exclusion, no block re-registration, and no router warning;
tests/unit_tests/recipes/test_dist_utils.pycovers parser normalization variants.including long-context large-vocabulary pipeline-parallel training (chunked CE removed
last-stage loss OOMs) and packed-sequence SFT with FA2.
Part of #2985
🤖 Generated with Claude Code