Skip to content

[megatron] fused linear cross-entropy for 23x memory savings at 262k context#1841

Merged
erictang000 merged 4 commits into
NovaSky-AI:mainfrom
casper-hansen:casper/liger-ce
Jun 29, 2026
Merged

[megatron] fused linear cross-entropy for 23x memory savings at 262k context#1841
erictang000 merged 4 commits into
NovaSky-AI:mainfrom
casper-hansen:casper/liger-ce

Conversation

@casper-hansen

@casper-hansen casper-hansen commented Jun 26, 2026

Copy link
Copy Markdown
Contributor

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_entropy optimizes 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)

seq_len baseline (eager CE) nvidia (fused CE) fused linear CE (ours) fused linear mem vs baseline
8,192 3.1 GB / 11 ms 5.0 GB / 8 ms 2.1 GB / 25 ms 1.5×
32,768 11.7 GB / 43 ms 19.3 GB / 29 ms 2.3 GB / 98 ms 5.1×
65,536 23.2 GB / 86 ms 38.4 GB / 58 ms 2.5 GB / 196 ms 9.1×
131,072 46.2 GB / 175 ms 76.5 GB / 120 ms 3.0 GB / 392 ms 15.2×
262,144 92.2 GB / 357 ms 152.8 GB / 245 ms 4.0 GB / 784 ms 22.8×

Full vocab = 248,320

seq_len baseline (eager CE) nvidia (fused CE) fused linear CE (ours)
8,192 12.4 GB / 42 ms 19.9 GB / 34 ms 8.1 GB / 92 ms
32,768 46.6 GB / 177 ms 76.9 GB / 151 ms 8.3 GB / 367 ms
65,536 92.2 GB / 356 ms 152.8 GB / 316 ms 8.6 GB / 734 ms
131,072 183.3 GB / 723 ms OOM 9.1 GB / 1,467 ms
262,144 OOM OOM 10.1 GB / 2,932 ms

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.

image

…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.

@gemini-code-assist gemini-code-assist Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Comment thread skyrl/backends/skyrl_train/distributed/megatron/model_utils.py Outdated
Comment thread skyrl/backends/skyrl_train/distributed/megatron/model_utils.py Outdated
…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.
@dyurk-lila

Copy link
Copy Markdown
Contributor

How does this compare with #1765?

@casper-hansen

Copy link
Copy Markdown
Contributor Author

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.
@casper-hansen

Copy link
Copy Markdown
Contributor Author

@SumanthRH @erictang000 Could you help review? Much appreciated!

@erictang000 erictang000 self-assigned this Jun 26, 2026
dyurk-lila pushed a commit to dyurk-lila/SkyRL that referenced this pull request Jun 26, 2026
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>
@casper-hansen

casper-hansen commented Jun 27, 2026

Copy link
Copy Markdown
Contributor Author

GSM8K RL GRPO run (another convergence test)

image

@erictang000

Copy link
Copy Markdown
Collaborator

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:

  • a fused torch only implementation
  • a fused triton implementation from verl that shows speedups over torch only
  • both of which compute entropy along with logprobs at the same time, with grads, saving an additional forward pass for entropy.

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:

  1. entropy backward for fused torch implementation
  2. triton backend kernel + efficient entropy kernel + backend selection ("torch" vs "triton")

@erictang000 erictang000 left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

very clean PR, lgtm! thanks again @casper-hansen

@erictang000
erictang000 merged commit 0a18dd7 into NovaSky-AI:main Jun 29, 2026
4 of 6 checks passed
dyurk-lila pushed a commit to dyurk-lila/SkyRL that referenced this pull request Jul 1, 2026
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>
dyurk-lila added a commit to dyurk-lila/SkyRL that referenced this pull request Jul 1, 2026
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>
erictang000 added a commit that referenced this pull request Jul 2, 2026
…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>
dyurk-lila pushed a commit to dyurk-lila/SkyRL that referenced this pull request Jul 7, 2026
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>
dyurk-lila added a commit to dyurk-lila/SkyRL that referenced this pull request Jul 7, 2026
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>
erictang000 pushed a commit that referenced this pull request Jul 7, 2026
## 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>
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.

3 participants