[megatron] fused linear cross-entropy for 23x memory savings at 262k context#1841
Conversation
…t activation memory SkyRL's Megatron RL/SFT loss turns decoder hidden states into per-token log-probs by materializing the full [B, S, vocab//TP] logits (the output layer) and, in the chunked-logprob backward, a float32 gradient of the same shape. For large vocabularies + long contexts this dominates activation memory (Qwen3.6-35B-A3B, vocab=248320: ~150 GB/rank at 262k) and is the wall that OOMs long-context training. Add a fused-linear log-prob that folds the LM-head projection into the chunked, tensor/context-parallel log-prob/entropy via the GPTModel `output_processor` hook, so the logits and their gradient are never built — the Liger fused-linear-cross-entropy memory technique, generalized to Megatron TP/CP and to per-token log-probs (which Liger's stock FLCE does not support: single-device, no reduction='none' backward). - model_utils: FusedLinearChunkedDistributedLogprob + from_parallel_hidden_to_logprobs[_packed_sequences] + a no-grad fused entropy. Mirrors ChunkedDistributedLogprob's math, with the cross-rank max / sum-exp / chosen-logit all-reduces and bf16/fp32 dtype handling. - MegatronModelWrapper: an output_processor that gathers the sequence-parallel hidden states (as the ColumnParallelLinear output layer would) and exposes the LM-head weight; loss_func uses the fused path under `trainer.fused_lm_head_logprob` (default off; numerically matches the default path). - tests: equivalency (fwd + grad_hidden + grad_weight) vs the materialize- logits path across chunk sizes, out-of-vocab targets, mixed dtypes, edge shapes, and TP=1/2/4 (torchrun); plus a Liger-FLCE forward oracle. - benchmark: skyrl/benchmarks/bench_fused_linear_logprob.py — up to ~38x less LM-head activation memory at 262k (OOM -> fits) at a ~2.5x logprob-compute cost.
…d CE Extends the LM-head benchmark to a 3-way comparison so the "why not the existing Megatron/Apex/TE fused CE?" question is answered with numbers (all three verified to produce the same gradient): baseline : megatron vocab_parallel_cross_entropy (eager) on materialized logits nvidia : megatron fused_vocab_parallel_cross_entropy (@jit_fuser) on logits liger : FusedLinearChunkedDistributedLogprob (no logits) Finding: the NVIDIA fused CE optimizes *compute* (fastest) but not *memory* — its staged @jit_fuser kernels allocate extra buffers, so it uses MORE memory than the eager baseline and OOMs earlier at long context. Liger uses ~23-38x less LM-head memory and is the only path that fits 262k (at a compute cost). The two are complementary: fused CE for short-context throughput, Liger for long-context memory.
There was a problem hiding this comment.
Code Review
This pull request introduces a fused LM-head and distributed token-logprob implementation (FusedLinearChunkedDistributedLogprob) that chunks over the sequence to avoid materializing the full logits tensor, dramatically reducing peak memory for long-context training. It integrates this optimization into the Megatron model wrapper and adds benchmark and unit tests. The review feedback highlights three key improvements: deferring the all_reduce operation in the forward pass to a single call after the chunk loop to reduce synchronization overhead, updating the forward-only inference path in MegatronModelWrapper to support the fused path and prevent OOMs during reference model evaluations, and performing the backward weight gradient matrix multiplication in low precision to leverage Tensor Cores before accumulating in float32.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
…grad matmul, fused inference path Gemini review on NovaSky-AI#1841: - model_utils: defer the per-chunk all_reduce to a single collective after the chunk loop (num_chunks -> 1; 256 -> 1 at S=262k, chunk=1024), removing the per-chunk launch/sync overhead. Numerically identical (SUM is associative across the sequence concat; each rank holds the log-prob for its own vocab shard, 0 elsewhere). - model_utils: run the backward weight-grad matmul in the low-precision compute dtype (Tensor Cores) and only cast up to fp32 for the accumulator, matching ColumnParallelLinear's backward. ~2x faster long-context backward at equal memory (262k per-rank shard: 1697 -> 784 ms); equivalency tests still pass. - megatron_model_wrapper: route the forward-only inference path (e.g. PPO reference/old logprobs) through the fused LM head via output_processor + from_parallel_hidden_to_logprobs(inference_only=True), so it no longer materializes the full [B, S, vocab//TP] logits and OOMs at long context. CI: fix ruff I001 import-sort in the benchmark + black formatting.
|
How does this compare with #1765? |
I don't know, you didn't include any benchmarks in your PR. Please feel free to compare! In terms of simplicity, this one gives you a lot for 1k lines of code where as yours seem to import from verl with 3.5k lines of code. |
The test imports the megatron model_utils and is launched with --extra megatron, so it belongs in the `-m megatron` GPU suite. Without the marker it was collected by neither the megatron job (`-m megatron`) nor the fsdp job (which `--ignore`s the megatron dir), i.e. it never ran in CI.
|
@SumanthRH @erictang000 Could you help review? Much appreciated! |
Adds skyrl/benchmarks/bench_fused_ce_h100.py — forward+backward through the LM-head training signal comparing megatron-core's eager vocab_parallel_cross_entropy and NVIDIA fused_vocab_parallel_cross_entropy (both materialize the [B,S,vocab//TP] logits) against this PR's two fused-linear backends (FusedLinearLogprob, FusedLinearLogprobTriton), which never materialize the logits. Same harness shape as NovaSky-AI#1841's bench_fused_linear_logprob.py; reports peak max_memory_allocated + wall-clock. Used to produce the Qwen3.6-35B-A3B table in the PR description. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
|
hey @dyurk-lila @casper-hansen thanks both for the PRs - so the current overlap is that both of the two PRs #1841 and #1765 are adding a fused + chunked lm head. #1841 (this pr) adds:
#1765 adds:
I'm inclined to merge #1841 first, since it's less LOC and less invasive for now, then we can gradually port over/merge components from #1765, in order:
|
erictang000
left a comment
There was a problem hiding this comment.
very clean PR, lgtm! thanks again @casper-hansen
Adds skyrl/benchmarks/bench_fused_ce_h100.py — forward+backward through the LM-head training signal comparing megatron-core's eager vocab_parallel_cross_entropy and NVIDIA fused_vocab_parallel_cross_entropy (both materialize the [B,S,vocab//TP] logits) against this PR's two fused-linear backends (FusedLinearLogprob, FusedLinearLogprobTriton), which never materialize the logits. Same harness shape as NovaSky-AI#1841's bench_fused_linear_logprob.py; reports peak max_memory_allocated + wall-clock. Used to produce the Qwen3.6-35B-A3B table in the PR description. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Rebase re-scoped this PR to layer only its distinguishing additions (the Triton backend + its GPU tests) on top of the merged NovaSky-AI#1841 fused-linear-cross-entropy. The torch chunked-logprob core is now upstream's FusedLinearChunkedDistributedLogprob; this commit fixes the vendored Triton adapter's docstrings, the benchmark, and the GPU test to reference the upstream class / the reconciled fused_lm_head_logprob_backend selector instead of this PR's now-dropped FusedLinearLogprob / _fused_linear_logprob_apply. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…the non-packed forward path (#1848) # [megatron] Fix TE sliding-window attention crash at micro_batch>1 on the non-packed forward path ## Summary On the non-packed Megatron forward path (`remove_microbatch_padding=False`), `megatron_model_wrapper.py` passes the 2-D `[batch, seq]` keep-mask returned by `remove_left_padding()` straight to `model(...)`. For models whose attention uses a Transformer Engine **sliding-window** mask, TE's `get_full_mask` combines that mask with the SWA mask via `torch.logical_or(swa_mask, attention_mask)`. A 2-D `[batch, seq]` mask is not broadcastable against the `[..., seq, seq]` SWA mask, so at **micro_batch_size > 1** the batch dimension collides with the sequence dimension and the forward crashes. This adds a pure-torch helper `to_te_attention_mask` that converts the 2-D keep-mask to a `[batch, 1, 1, seq]` padding mask (`True` marks padding), and routes the `model(...)` call sites through it. The original 2-D mask is left untouched for the downstream `recover_left_padding()`. ## Reproduction (current `main`, `0a18dd72`) Real model **`Qwen/Qwen3.6-35B-A3B`** (hybrid GatedDeltaNet + MoE; its full-attention layers use TE sliding-window attention), `tp=2, pp=1`, `micro_batch_size=2`, `remove_microbatch_padding=False`, left-padded batch: - **Stock `main` crashes** in the megatron forward: ``` transformer_engine/pytorch/attention/dot_product_attention/utils.py:1363, in get_full_mask attention_mask = torch.logical_or(swa_mask, attention_mask) RuntimeError: The size of tensor a (24) must match the size of tensor b (2) at non-singleton dimension 2 ``` (`24` = sequence length, `2` = `micro_batch_size`.) - **With this patch** the same forward completes and returns logprobs — test passes. - **Regression:** a standard non-SWA model (`Qwen/Qwen3-0.6B`, non-packed, `micro_batch_size=2`, left-padded) passes the HF-vs-megatron parity test **identically on stock and patched** — the helper is a no-op there. Lifting `micro_batch_size` from 1 to 4 on this path (once unblocked) increased training throughput ~3.3× at unchanged reward in our multi-LoRA RL runs. ## Changes - Add `skyrl/backends/skyrl_train/distributed/megatron/mask_utils.py` with `to_te_attention_mask`. - Route both `model(...)` call sites in `megatron_model_wrapper.py` (4 invocations after #1841) through the helper. - Add `tests/backends/skyrl_train/distributed/test_mask_utils.py`. ## Design notes - The reshape lives in a standalone pure-torch helper so it is unit-testable on CPU without the `megatron` extra. - No-op on the packed path: there `new_attention_mask` is `None` and `packed_seq_params` drives masking, so the helper returns the input unchanged. An already higher-rank mask is also returned unchanged, so non-SWA TE attention (which already accepts the 2-D mask) is unaffected. - Not folded into `remove_left_padding()` because that function's 2-D return is also consumed by `recover_left_padding()`, which needs the 2-D form. - `True = padding` matches TE's mask convention. GatedDeltaNet / linear-attention layers ignore `attention_mask`. No config or API change; behavior is identical at `micro_batch_size = 1`. ## Testing - Unit (CPU): `uv run --extra dev pytest tests/backends/skyrl_train/distributed/test_mask_utils.py` — 5 passed. - Integration (GPU, run on 2× B300): `test_megatron_forward[tp2_pp1_policy]` with `Qwen/Qwen3.6-35B-A3B` — **fails on stock** at `get_full_mask`, **passes with this patch**; `Qwen/Qwen3-0.6B` passes on both (regression-safe). ## Checklist - [x] Follows Google Python style; comments describe what the code does. - [x] CPU unit test added; GPU reproduction + fix verified on Qwen3.6-35B-A3B. - [x] No paths/naming changes affecting `.claude/` docs; no example-script or doc changes needed. - [x] `bash format.sh` (ruff 0.11.9 + black 24.10.0 + gitleaks) clean — verified, zero auto-fixes. --------- Co-authored-by: Eric Tang <erictang000@gmail.com>
Adds skyrl/benchmarks/bench_fused_ce_h100.py — forward+backward through the LM-head training signal comparing megatron-core's eager vocab_parallel_cross_entropy and NVIDIA fused_vocab_parallel_cross_entropy (both materialize the [B,S,vocab//TP] logits) against this PR's two fused-linear backends (FusedLinearLogprob, FusedLinearLogprobTriton), which never materialize the logits. Same harness shape as NovaSky-AI#1841's bench_fused_linear_logprob.py; reports peak max_memory_allocated + wall-clock. Used to produce the Qwen3.6-35B-A3B table in the PR description. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Rebase re-scoped this PR to layer only its distinguishing additions (the Triton backend + its GPU tests) on top of the merged NovaSky-AI#1841 fused-linear-cross-entropy. The torch chunked-logprob core is now upstream's FusedLinearChunkedDistributedLogprob; this commit fixes the vendored Triton adapter's docstrings, the benchmark, and the GPU test to reference the upstream class / the reconciled fused_lm_head_logprob_backend selector instead of this PR's now-dropped FusedLinearLogprob / _fused_linear_logprob_apply. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
## Reviewers: Where to Look > **Rebase note (re-scoped onto #1841):** [#1841](#1841) landed the torch fused LM-head log-prob path before this PR. This PR now sits on top of #1841 and only adds the remaining Triton backend pieces. It does **not** implement entropy-with-gradient for fused LM-head; upstream #1841 keeps fused entropy metric-only/no-grad, and this PR preserves that contract. Please review closely: 1. **Triton backend selector + dispatch** - `skyrl/backends/skyrl_train/distributed/megatron/model_utils.py`, `_fused_lm_head_logprob_apply(...)`. This chooses between upstream `FusedLinearChunkedDistributedLogprob` and the vendored Triton adapter. Both backends now share the same logprob-only apply contract and return TP-combined `[B, S]` log-probs. 2. **Vendored Triton kernel adapter** - `skyrl/backends/skyrl_train/distributed/megatron/fused_linear_logprob_triton.py`, `FusedLinearLogprobTriton`. Review the SkyRL adapter around the vendored verl/NVIDIA kernel: arbitrary shard-layout target re-keying, fully-OOV forward masking/backward pass-through, hidden-to-weight dtype casting, and import-safe CUDA/Triton availability checks. 3. **Backend config wiring** - `skyrl/train/config/config.py` and `skyrl/backends/skyrl_train/workers/megatron/megatron_model_wrapper.py`. This adds `trainer.fused_lm_head_logprob_backend: "torch" | "triton"` as a backend selector behind the existing `trainer.fused_lm_head_logprob` on/off flag. 4. **Triton GPU tests** - `tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_fused_linear_logprob_triton.py`. These drive `FusedLinearLogprobTriton.apply` directly against the stock materialized-logits path for logprob forward, grad-hidden, and grad-weight across TP1/TP2, chunked/single-chunk, fp32/bf16, and OOV targets. Needs less scrutiny: the vendored Triton kernel internals and the H100 benchmark harness. The top-level `NOTICE` and standalone `triton` extra were removed after review; attribution remains in the vendored file header and the Megatron/FSDP dependency graph already resolves Triton through the pinned torch stack. ## Summary Adds an opt-in **Triton backend** for SkyRL/Megatron fused LM-head log-prob, layered on top of upstream #1841. The default `torch` backend remains upstream `FusedLinearChunkedDistributedLogprob`; `triton` uses the vendored verl/NVIDIA online-softmax kernel to avoid materializing `[B, S, vocab//TP]` logits while improving long-context throughput on H100. The runtime contract is intentionally narrow: ```yaml trainer.fused_lm_head_logprob: false # existing on/off switch trainer.fused_lm_head_logprob_backend: torch # "torch" | "triton" ``` If `triton` is selected but Triton/CUDA is unavailable, dispatch warns and falls back to the upstream torch backend. Entropy under fused LM-head remains metric-only/no-grad, matching #1841; `use_entropy_loss=True` on the fused path still raises until entropy gradients are wired and validated separately. ## Benchmark Single H100 80GB, `Qwen/Qwen3.6-35B-A3B`, hidden=2048, chunk=1024, bf16, batch 1, mean of 3 reps. Compared against materialized-logits eager CE and NVIDIA fused CE; both materialize `[B, S, vocab//TP]` logits, while torch/triton fused-linear paths do not. ### Per-rank vocab shard = 62,080 | seq_len | baseline | nvidia | fused-torch | fused-triton | torch mem vs baseline | triton mem vs baseline | |---:|---:|---:|---:|---:|---:|---:| | 8,192 | 3,249 MB / 19 ms | 5,189 MB / 18 ms | 2,323 MB / 116 ms | **842 MB / 27 ms** | 1.4x | 3.9x | | 16,384 | 6,191 MB / 37 ms | 10,071 MB / 35 ms | 2,387 MB / 231 ms | **1,134 MB / 54 ms** | 2.6x | 5.5x | | 32,768 | 12,076 MB / 74 ms | 19,835 MB / 69 ms | 2,514 MB / 462 ms | **1,716 MB / 108 ms** | 4.8x | 7.0x | | 65,536 | 23,845 MB / 152 ms | 39,364 MB / 177 ms | 2,772 MB / 922 ms | **2,883 MB / 216 ms** | 8.6x | 8.3x | | 131,072 | 47,383 MB / 387 ms | 78,422 MB / 414 ms | 3,285 MB / 1,851 ms | **5,218 MB / 436 ms** | 14.4x | 9.1x | | 262,144 | **OOM** | **OOM** | 4,311 MB / 3,675 ms | **9,886 MB / 880 ms** | n/a | n/a | ## Test plan - [x] `git diff --check` - [x] Review grep: no stale `return_entropy` / dual-output / top-level NOTICE / old config-name claims remain in `skyrl`, `tests`, or `pyproject.toml` - [x] `PYTHONPYCACHEPREFIX=/tmp/codex-pyc python3 -m py_compile skyrl/backends/skyrl_train/distributed/megatron/model_utils.py skyrl/backends/skyrl_train/workers/megatron/megatron_model_wrapper.py tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_fused_linear_logprob_triton.py` - [ ] GPU: `pytest -s tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_fused_linear_logprob_triton.py -m megatron` on CUDA/H100 Generated with [Codex](https://openai.com/codex) --------- Signed-off-by: dyurk-lila <dyurk@lila.ai> Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> Co-authored-by: lila-sync-bot <lila-sync-bot@users.noreply.github.com>

I wanted to run Qwen/Qwen3.6-35B-A3B on 8xB300 using SkyRL and found that it cannot handle full context without OOMing on the cross-entropy step in the trainer. To solve this, I implemented a fused cross-entropy that saves a lot of memory for the Megatron backend while handling the parallelism.
Why not the existing Megatron/Apex/TE fused cross-entropy?
fused_vocab_parallel_cross_entropyoptimizes the CE compute, not memory: it materializes the[S, vocab//TP]logits and its staged fused kernels allocate extra buffers, so it uses ~1.7x more memory than eager CE and OOMs earlier at long context.Per-rank vocab = 62,080 (the real TP=4 shard)
Full vocab = 248,320
Convergence
I did a small SFT test to overfit on GSM8K with Qwen/Qwen2.5-1.5B-Instruct. Convergence is the same, just with slight variations as the computation has some natural variance when it's chunked.