[megatron] Fused LM-head log-prob + entropy (avoid full [*, seq, vocab] logit materialization)#1765
Open
dyurk-lila wants to merge 9 commits into
Open
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Code Review
This pull request introduces a fused LM-head log-prob optimization for the Megatron backend, featuring both a pure-PyTorch chunked implementation and a high-performance Triton-based backend vendored from verl. By fusing the LM-head projection into the chunked log-prob and entropy computation, the PR avoids materializing the massive full-vocab logits tensor, significantly reducing peak memory usage. The feedback recommends enhancing the robustness of the model wrapper by using *args and **kwargs to guard against future Megatron-core signature changes, broadening the exception handling when importing the Triton backend to catch driver or library mismatches, and checking CUDA availability alongside Triton to prevent cryptic errors in CPU-only environments.
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 on lines
+84
to
+101
| def fused_output_layer_forward(input_, weight=None, runtime_gather_output=None): | ||
| w = weight if weight is not None else getattr(layer, "weight", None) | ||
| if w is None: | ||
| # Tied-embedding stage with no allocated head weight AND none passed in: | ||
| # we cannot fuse safely — defer to the original projection. | ||
| return original(input_, weight=weight, runtime_gather_output=runtime_gather_output) | ||
| hidden = input_ | ||
| if getattr(layer, "sequence_parallel", False): | ||
| # Gather the sequence-parallel-scattered hidden across TP. Its backward is a | ||
| # reduce-scatter (tensor_parallel_output_grad=True), which reduces the fused | ||
| # function's per-shard grad_hidden exactly as ColumnParallelLinear would. | ||
| hidden = gather_from_sequence_parallel_region( | ||
| hidden, tensor_parallel_output_grad=True, group=tpg | ||
| ) | ||
| holder["weight"] = w | ||
| # Return hidden in the logits position (+ no bias). The model's _scale_logits is | ||
| # identity here (MuP excluded above) and its transpose yields [b, s, h]. | ||
| return hidden, None |
Contributor
There was a problem hiding this comment.
Using explicit keyword arguments like weight=None and runtime_gather_output=None can lead to TypeError if Megatron-core updates its ParallelLinear.forward signature or if other wrappers pass unexpected arguments. Using *args and **kwargs makes the wrapper future-proof and robust against signature changes.
def fused_output_layer_forward(input_, *args, **kwargs):
w = kwargs.get("weight", None)
if w is None and len(args) > 0:
w = args[0]
w = w if w is not None else getattr(layer, "weight", None)
if w is None:
# Tied-embedding stage with no allocated head weight AND none passed in:
# we cannot fuse safely — defer to the original projection.
return original(input_, *args, **kwargs)
hidden = input_
if getattr(layer, "sequence_parallel", False):
# Gather the sequence-parallel-scattered hidden across TP. Its backward is a
# reduce-scatter (tensor_parallel_output_grad=True), which reduces the fused
# function's per-shard grad_hidden exactly as ColumnParallelLinear would.
hidden = gather_from_sequence_parallel_region(
hidden, tensor_parallel_output_grad=True, group=tpg
)
holder["weight"] = w
# Return hidden in the logits position (+ no bias). The model's _scale_logits is
# identity here (MuP excluded above) and its transpose yields [b, s, h].
return hidden, None
dyurk-lila
commented
Jun 9, 2026
Author
There was a problem hiding this comment.
Claude suggested that this is a proper copyright notice but this should definitely be checked by someone more familiar with the repo's conventions around this
Collaborator
There was a problem hiding this comment.
hmm yeah can we leave this off for now and just stick to the repo wide convention of providing notice at the top of the vendored file?
dyurk-lila
added a commit
to dyurk-lila/SkyRL
that referenced
this pull request
Jun 16, 2026
…ted buffer (lower peak memory)
## Summary
Refactor `ChunkedDistributedLogprob.backward` (the vocab-parallel chunked-logprob autograd function used by the Megatron worker for SFT cross-entropy and RL policy/ref losses) to stream each chunk's gradient into a single preallocated fp32 buffer instead of appending to a Python list and concatenating with `torch.cat` at the end. This lowers peak activation memory on the chunked backward path with **no change to numerics**.
## What changed
`skyrl/backends/skyrl_train/distributed/megatron/model_utils.py`, `ChunkedDistributedLogprob.backward`:
- Removed `all_grad_input = []` and the final `grad_input = torch.cat(all_grad_input, dim=1)`.
- Preallocate once before the loop:
`grad_input = torch.empty((batch_size, seq_size, partition_vocab_size), dtype=torch.float32, device=vocab_parallel_logits.device)`.
- Renamed the per-chunk grad to `chunk_grad_input` and, after the `scatter_add_`, write it into its sequence slice: `grad_input[:, chunk_start:chunk_end, :] = chunk_grad_input`.
The change is **unconditional** (no flag) because it is numerically byte-identical. It only engages on the chunked dispatch path (i.e. when `chunk_size < seq_len_local`). The non-chunked `DistributedLogprob`, `forward()`, and the vendored Triton fused-LCE path are untouched. `forward()` retains the list+`cat` form intentionally: its accumulator holds per-chunk `[batch_size, chunk_len]` log-prob tensors (tiny), not the `[batch_size, chunk_len, V//TP]` fp32 grads that make the backward `cat` expensive, so streaming it would not meaningfully lower peak.
The old list-then-`cat` form kept every per-chunk `[B, chunk_len, V//TP]` fp32 grad alive **and** allocated the full concatenated output at the cat moment, so peak was ~2x the full `[B, seq, V//TP]` fp32 grad. Streaming drops peak to full-buffer + one live chunk = ~`(1 + 1/num_chunks)`x of the full grad. The win scales with chunk count; it is **not** a flat halving.
## Numerical equivalence / safety
Byte-identical by construction:
- The per-chunk math is unchanged: same `_compute_distributed_log_softmax` on the same fp32 slice, same `.exp()`, same `neg_`/`mul_`/`scatter_add_` formulation. `chunk_grad_input` is a fresh fp32 tensor, so the slice write is a same-dtype copy with no cast.
- The chunks tile `[0, seq_size)` exactly and contiguously: chunk `i` covers `[i*chunk_size, min(seq_size, (i+1)*chunk_size))`, consecutive chunks meet with no gap/overlap, and `num_chunks = ceil(seq_size/chunk_size)`. `torch.cat(dim=1)` placed chunk `i`'s columns at those same `[chunk_start:chunk_end]` offsets, so values land at identical positions and each slice is written exactly once.
- The new buffer is contiguous fp32 of shape `[B, seq_size, V//TP]` — identical shape/dtype/contiguity to the previous `torch.cat` output — and is fully overwritten by the full-coverage tiling, so no uninitialized `torch.empty` memory survives. Full-coverage tiling is a load-bearing invariant for the `torch.empty` buffer (a partial write would leak garbage silently); the prime-length coverage test below is the regression guard.
A deliberate choice of a **separate fp32 buffer** (rather than writing in place into the autograd-saved `vocab_parallel_logits`): the separate buffer keeps the gradient in fp32 regardless of the saved logits' dtype, does not mutate an autograd-saved tensor (avoiding version-counter / double-backward hazards), and trades only ~`(1/num_chunks)`x extra memory for that safety.
## Test plan
> Written to SkyRL CI conventions; `python -m py_compile` passes and lint is clean (`ruff`: all checks passed; `black --line-length 120`: unchanged) on all three files.
- `tests/backends/skyrl_train/distributed/test_chunked_logprob_backward_streaming.py` (CPU lane): stubs `megatron.core.parallel_state` into `sys.modules` via a module-scoped autouse save/restore fixture (mirroring `test_preprocess_packed_seqs_cp.py`), uses a gloo `world_size=1` TP group (every `all_reduce` is the identity), and asserts `torch.equal` between the chunked-streamed grad and the single-shot `DistributedLogprob` grad across `chunk_size in {1,3,7,16,32,64}` x with/without OOV targets, plus edge cases (seq_len=1, all-in/all-out mask, a prime/ragged tiny-vocab config). This is the **byte-identity gate for the storage refactor**: the log-softmax + scatter-add backward math is purely per-position (reductions only over the vocab dim), so chunk boundaries cannot change any value, and the world_size=1 `torch.equal` fully validates the device/dtype-agnostic slice-write. The prime-length (`seq_len=17`, `chunk_size=5`) coverage test additionally pins that every sequence slice of the preallocated buffer is written (an unwritten `torch.empty` slice cannot coincidentally match the reference). Collected by the SkyRL-Train-CPU lane (`pytest tests/backends/skyrl_train/ --ignore=.../gpu`).
Run with:
```
uv run --isolated --extra dev -- pytest -s \
tests/backends/skyrl_train/distributed/test_chunked_logprob_backward_streaming.py
```
- `tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_chunked_logprob_backward_tp.py` (`@pytest.mark.megatron`): spawns TP NCCL ranks (parametrized `tp_size in {2, 4}`, guarded by a `device_count()` skip), runs `mpu.initialize_model_parallel`, shards the vocab across ranks, runs fwd+bwd on each rank's slice, and asserts (a) the rank vocab slices tile `[0, vocab)` exactly once via `torch.equal` on a coverage counter, and (b) each rank's local streamed grad matches the full-vocab single-process autograd reference columns **within fp32 tolerance** via `torch.testing.assert_close(atol=1e-5, rtol=1e-4)` — a tolerance is correct here because the cross-rank all-reduce reorders the fp32 reduction vs. the single-tensor reference (matching the existing GPU test's grad tolerance), while the no-overlap tiling check uses exact equality. Spawned ranks set the conftest-mandated NCCL env (`NCCL_CUMEM_ENABLE=0`, etc.) before `init_process_group`, since `mp.spawn` children do not inherit the runtime env set by the CI conftest.
Run with (>=2 free GPUs; the `tp_size=4` case is skipped unless 4 are present):
```
uv run --isolated --extra dev --extra megatron -- \
pytest -s tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_chunked_logprob_backward_tp.py
```
## Scope & follow-ups
- **Backends covered:** Megatron only, by nature — `ChunkedDistributedLogprob` is the vocab-parallel (TP/CP) chunked-logprob used on the Megatron worker, reached through the single `forward_backward_mini_batch` train-step chokepoint, so all Megatron-backed training pathways (SFT, sync/async/fully-async RL, full-context) inherit the optimization with no per-pathway edit. FSDP has no vocab-parallel chunked-logprob equivalent, so it is intentionally out of scope.
- **Scope:** Limited to the single `backward` method; `DistributedLogprob`, `forward()`, and the fused-LCE path are unchanged.
- **Deferred:** A direct streaming-vs-old-`cat` comparison at TP>1 was intentionally not added, since it would require shipping the pre-refactor `cat`-based backward as dead reference code. The byte-identity gate for the refactor itself lives on the CPU `world_size=1` lane (`torch.equal`); the GPU lane validates the distributed math against the reference (with tolerance) plus exact no-overlap tiling.
## Related PRs (merge ordering)
- **NovaSky-AI#1765** (fused LM-head log-prob + entropy) edits the same `ChunkedDistributedLogprob.backward` and **keeps** the `all_grad_input = []` / `torch.cat` form. It refactors the per-chunk chosen-token scatter-add into a shared `_add_chosen_token_grad` helper — touching the exact lines this PR renames to `chunk_grad_input` and streams — so the two **will textually collide on this hunk**. NovaSky-AI#1765 does **not** subsume this PR's peak-memory win (it preserves the list+`cat`), and NovaSky-AI#1765 is otherwise complementary in intent. Whichever lands first forces a rebase of the other. **Recommended ordering:** land NovaSky-AI#1765 first, then rebase this streaming change on top of it (the preallocated-buffer write replaces the `append`+`cat` NovaSky-AI#1765 keeps); if the two are reviewed together they can be folded into one change.
- **NovaSky-AI#1543** (WIP) independently removes the same `torch.cat` 2x-peak, but does so **in place** into the autograd-saved logits with no extra buffer. This PR deliberately does **not** take that approach: the in-place variant downcasts the fp32 grad to the saved logits' dtype (breaking byte-identity) and mutates an autograd-saved tensor (the version-counter / double-backward hazard). This PR's separate fp32 buffer preserves numeric fidelity and autograd-safety at the cost of ~`(1/num_chunks)`x extra memory. NovaSky-AI#1543 is also built on pre-merge code and currently conflicts with main, so it is not a viable supersession.
dyurk-lila
added a commit
to dyurk-lila/SkyRL
that referenced
this pull request
Jun 18, 2026
…ted buffer (lower peak memory)
## Summary
Refactor `ChunkedDistributedLogprob.backward` (the vocab-parallel chunked-logprob autograd function used by the Megatron worker for SFT cross-entropy and RL policy/ref losses) to stream each chunk's gradient into a single preallocated fp32 buffer instead of appending to a Python list and concatenating with `torch.cat` at the end. This lowers peak activation memory on the chunked backward path with **no change to numerics**.
## What changed
`skyrl/backends/skyrl_train/distributed/megatron/model_utils.py`, `ChunkedDistributedLogprob.backward`:
- Removed `all_grad_input = []` and the final `grad_input = torch.cat(all_grad_input, dim=1)`.
- Preallocate once before the loop:
`grad_input = torch.empty((batch_size, seq_size, partition_vocab_size), dtype=torch.float32, device=vocab_parallel_logits.device)`.
- Renamed the per-chunk grad to `chunk_grad_input` and, after the `scatter_add_`, write it into its sequence slice: `grad_input[:, chunk_start:chunk_end, :] = chunk_grad_input`.
The change is **unconditional** (no flag) because it is numerically byte-identical. It only engages on the chunked dispatch path (i.e. when `chunk_size < seq_len_local`). The non-chunked `DistributedLogprob`, `forward()`, and the vendored Triton fused-LCE path are untouched. `forward()` retains the list+`cat` form intentionally: its accumulator holds per-chunk `[batch_size, chunk_len]` log-prob tensors (tiny), not the `[batch_size, chunk_len, V//TP]` fp32 grads that make the backward `cat` expensive, so streaming it would not meaningfully lower peak.
The old list-then-`cat` form kept every per-chunk `[B, chunk_len, V//TP]` fp32 grad alive **and** allocated the full concatenated output at the cat moment, so peak was ~2x the full `[B, seq, V//TP]` fp32 grad. Streaming drops peak to full-buffer + one live chunk = ~`(1 + 1/num_chunks)`x of the full grad. The win scales with chunk count; it is **not** a flat halving.
## Numerical equivalence / safety
Byte-identical by construction:
- The per-chunk math is unchanged: same `_compute_distributed_log_softmax` on the same fp32 slice, same `.exp()`, same `neg_`/`mul_`/`scatter_add_` formulation. `chunk_grad_input` is a fresh fp32 tensor, so the slice write is a same-dtype copy with no cast.
- The chunks tile `[0, seq_size)` exactly and contiguously: chunk `i` covers `[i*chunk_size, min(seq_size, (i+1)*chunk_size))`, consecutive chunks meet with no gap/overlap, and `num_chunks = ceil(seq_size/chunk_size)`. `torch.cat(dim=1)` placed chunk `i`'s columns at those same `[chunk_start:chunk_end]` offsets, so values land at identical positions and each slice is written exactly once.
- The new buffer is contiguous fp32 of shape `[B, seq_size, V//TP]` — identical shape/dtype/contiguity to the previous `torch.cat` output — and is fully overwritten by the full-coverage tiling, so no uninitialized `torch.empty` memory survives. Full-coverage tiling is a load-bearing invariant for the `torch.empty` buffer (a partial write would leak garbage silently); the prime-length coverage test below is the regression guard.
A deliberate choice of a **separate fp32 buffer** (rather than writing in place into the autograd-saved `vocab_parallel_logits`): the separate buffer keeps the gradient in fp32 regardless of the saved logits' dtype, does not mutate an autograd-saved tensor (avoiding version-counter / double-backward hazards), and trades only ~`(1/num_chunks)`x extra memory for that safety.
## Test plan
> Written to SkyRL CI conventions; `python -m py_compile` passes and lint is clean (`ruff`: all checks passed; `black --line-length 120`: unchanged) on all three files.
- `tests/backends/skyrl_train/distributed/test_chunked_logprob_backward_streaming.py` (CPU lane): stubs `megatron.core.parallel_state` into `sys.modules` via a module-scoped autouse save/restore fixture (mirroring `test_preprocess_packed_seqs_cp.py`), uses a gloo `world_size=1` TP group (every `all_reduce` is the identity), and asserts `torch.equal` between the chunked-streamed grad and the single-shot `DistributedLogprob` grad across `chunk_size in {1,3,7,16,32,64}` x with/without OOV targets, plus edge cases (seq_len=1, all-in/all-out mask, a prime/ragged tiny-vocab config). This is the **byte-identity gate for the storage refactor**: the log-softmax + scatter-add backward math is purely per-position (reductions only over the vocab dim), so chunk boundaries cannot change any value, and the world_size=1 `torch.equal` fully validates the device/dtype-agnostic slice-write. The prime-length (`seq_len=17`, `chunk_size=5`) coverage test additionally pins that every sequence slice of the preallocated buffer is written (an unwritten `torch.empty` slice cannot coincidentally match the reference). Collected by the SkyRL-Train-CPU lane (`pytest tests/backends/skyrl_train/ --ignore=.../gpu`).
Run with:
```
uv run --isolated --extra dev -- pytest -s \
tests/backends/skyrl_train/distributed/test_chunked_logprob_backward_streaming.py
```
- `tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_chunked_logprob_backward_tp.py` (`@pytest.mark.megatron`): spawns TP NCCL ranks (parametrized `tp_size in {2, 4}`, guarded by a `device_count()` skip), runs `mpu.initialize_model_parallel`, shards the vocab across ranks, runs fwd+bwd on each rank's slice, and asserts (a) the rank vocab slices tile `[0, vocab)` exactly once via `torch.equal` on a coverage counter, and (b) each rank's local streamed grad matches the full-vocab single-process autograd reference columns **within fp32 tolerance** via `torch.testing.assert_close(atol=1e-5, rtol=1e-4)` — a tolerance is correct here because the cross-rank all-reduce reorders the fp32 reduction vs. the single-tensor reference (matching the existing GPU test's grad tolerance), while the no-overlap tiling check uses exact equality. Spawned ranks set the conftest-mandated NCCL env (`NCCL_CUMEM_ENABLE=0`, etc.) before `init_process_group`, since `mp.spawn` children do not inherit the runtime env set by the CI conftest.
Run with (>=2 free GPUs; the `tp_size=4` case is skipped unless 4 are present):
```
uv run --isolated --extra dev --extra megatron -- \
pytest -s tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_chunked_logprob_backward_tp.py
```
## Scope & follow-ups
- **Backends covered:** Megatron only, by nature — `ChunkedDistributedLogprob` is the vocab-parallel (TP/CP) chunked-logprob used on the Megatron worker, reached through the single `forward_backward_mini_batch` train-step chokepoint, so all Megatron-backed training pathways (SFT, sync/async/fully-async RL, full-context) inherit the optimization with no per-pathway edit. FSDP has no vocab-parallel chunked-logprob equivalent, so it is intentionally out of scope.
- **Scope:** Limited to the single `backward` method; `DistributedLogprob`, `forward()`, and the fused-LCE path are unchanged.
- **Deferred:** A direct streaming-vs-old-`cat` comparison at TP>1 was intentionally not added, since it would require shipping the pre-refactor `cat`-based backward as dead reference code. The byte-identity gate for the refactor itself lives on the CPU `world_size=1` lane (`torch.equal`); the GPU lane validates the distributed math against the reference (with tolerance) plus exact no-overlap tiling.
## Related PRs (merge ordering)
- **NovaSky-AI#1765** (fused LM-head log-prob + entropy) edits the same `ChunkedDistributedLogprob.backward` and **keeps** the `all_grad_input = []` / `torch.cat` form. It refactors the per-chunk chosen-token scatter-add into a shared `_add_chosen_token_grad` helper — touching the exact lines this PR renames to `chunk_grad_input` and streams — so the two **will textually collide on this hunk**. NovaSky-AI#1765 does **not** subsume this PR's peak-memory win (it preserves the list+`cat`), and NovaSky-AI#1765 is otherwise complementary in intent. Whichever lands first forces a rebase of the other. **Recommended ordering:** land NovaSky-AI#1765 first, then rebase this streaming change on top of it (the preallocated-buffer write replaces the `append`+`cat` NovaSky-AI#1765 keeps); if the two are reviewed together they can be folded into one change.
- **NovaSky-AI#1543** (WIP) independently removes the same `torch.cat` 2x-peak, but does so **in place** into the autograd-saved logits with no extra buffer. This PR deliberately does **not** take that approach: the in-place variant downcasts the fp32 grad to the saved logits' dtype (breaking byte-identity) and mutates an autograd-saved tensor (the version-counter / double-backward hazard). This PR's separate fp32 buffer preserves numeric fidelity and autograd-safety at the cost of ~`(1/num_chunks)`x extra memory. NovaSky-AI#1543 is also built on pre-merge code and currently conflicts with main, so it is not a viable supersession.
dyurk-lila
added a commit
to dyurk-lila/SkyRL
that referenced
this pull request
Jul 1, 2026
…ted buffer (lower peak memory)
## Summary
Refactor `ChunkedDistributedLogprob.backward` (the vocab-parallel chunked-logprob autograd function used by the Megatron worker for SFT cross-entropy and RL policy/ref losses) to stream each chunk's gradient into a single preallocated fp32 buffer instead of appending to a Python list and concatenating with `torch.cat` at the end. This lowers peak activation memory on the chunked backward path with **no change to numerics**.
## What changed
`skyrl/backends/skyrl_train/distributed/megatron/model_utils.py`, `ChunkedDistributedLogprob.backward`:
- Removed `all_grad_input = []` and the final `grad_input = torch.cat(all_grad_input, dim=1)`.
- Preallocate once before the loop:
`grad_input = torch.empty((batch_size, seq_size, partition_vocab_size), dtype=torch.float32, device=vocab_parallel_logits.device)`.
- Renamed the per-chunk grad to `chunk_grad_input` and, after the `scatter_add_`, write it into its sequence slice: `grad_input[:, chunk_start:chunk_end, :] = chunk_grad_input`.
The change is **unconditional** (no flag) because it is numerically byte-identical. It only engages on the chunked dispatch path (i.e. when `chunk_size < seq_len_local`). The non-chunked `DistributedLogprob`, `forward()`, and the vendored Triton fused-LCE path are untouched. `forward()` retains the list+`cat` form intentionally: its accumulator holds per-chunk `[batch_size, chunk_len]` log-prob tensors (tiny), not the `[batch_size, chunk_len, V//TP]` fp32 grads that make the backward `cat` expensive, so streaming it would not meaningfully lower peak.
The old list-then-`cat` form kept every per-chunk `[B, chunk_len, V//TP]` fp32 grad alive **and** allocated the full concatenated output at the cat moment, so peak was ~2x the full `[B, seq, V//TP]` fp32 grad. Streaming drops peak to full-buffer + one live chunk = ~`(1 + 1/num_chunks)`x of the full grad. The win scales with chunk count; it is **not** a flat halving.
## Numerical equivalence / safety
Byte-identical by construction:
- The per-chunk math is unchanged: same `_compute_distributed_log_softmax` on the same fp32 slice, same `.exp()`, same `neg_`/`mul_`/`scatter_add_` formulation. `chunk_grad_input` is a fresh fp32 tensor, so the slice write is a same-dtype copy with no cast.
- The chunks tile `[0, seq_size)` exactly and contiguously: chunk `i` covers `[i*chunk_size, min(seq_size, (i+1)*chunk_size))`, consecutive chunks meet with no gap/overlap, and `num_chunks = ceil(seq_size/chunk_size)`. `torch.cat(dim=1)` placed chunk `i`'s columns at those same `[chunk_start:chunk_end]` offsets, so values land at identical positions and each slice is written exactly once.
- The new buffer is contiguous fp32 of shape `[B, seq_size, V//TP]` — identical shape/dtype/contiguity to the previous `torch.cat` output — and is fully overwritten by the full-coverage tiling, so no uninitialized `torch.empty` memory survives. Full-coverage tiling is a load-bearing invariant for the `torch.empty` buffer (a partial write would leak garbage silently); the prime-length coverage test below is the regression guard.
A deliberate choice of a **separate fp32 buffer** (rather than writing in place into the autograd-saved `vocab_parallel_logits`): the separate buffer keeps the gradient in fp32 regardless of the saved logits' dtype, does not mutate an autograd-saved tensor (avoiding version-counter / double-backward hazards), and trades only ~`(1/num_chunks)`x extra memory for that safety.
## Test plan
> Written to SkyRL CI conventions; `python -m py_compile` passes and lint is clean (`ruff`: all checks passed; `black --line-length 120`: unchanged) on all three files.
- `tests/backends/skyrl_train/distributed/test_chunked_logprob_backward_streaming.py` (CPU lane): stubs `megatron.core.parallel_state` into `sys.modules` via a module-scoped autouse save/restore fixture (mirroring `test_preprocess_packed_seqs_cp.py`), uses a gloo `world_size=1` TP group (every `all_reduce` is the identity), and asserts `torch.equal` between the chunked-streamed grad and the single-shot `DistributedLogprob` grad across `chunk_size in {1,3,7,16,32,64}` x with/without OOV targets, plus edge cases (seq_len=1, all-in/all-out mask, a prime/ragged tiny-vocab config). This is the **byte-identity gate for the storage refactor**: the log-softmax + scatter-add backward math is purely per-position (reductions only over the vocab dim), so chunk boundaries cannot change any value, and the world_size=1 `torch.equal` fully validates the device/dtype-agnostic slice-write. The prime-length (`seq_len=17`, `chunk_size=5`) coverage test additionally pins that every sequence slice of the preallocated buffer is written (an unwritten `torch.empty` slice cannot coincidentally match the reference). Collected by the SkyRL-Train-CPU lane (`pytest tests/backends/skyrl_train/ --ignore=.../gpu`).
Run with:
```
uv run --isolated --extra dev -- pytest -s \
tests/backends/skyrl_train/distributed/test_chunked_logprob_backward_streaming.py
```
- `tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_chunked_logprob_backward_tp.py` (`@pytest.mark.megatron`): spawns TP NCCL ranks (parametrized `tp_size in {2, 4}`, guarded by a `device_count()` skip), runs `mpu.initialize_model_parallel`, shards the vocab across ranks, runs fwd+bwd on each rank's slice, and asserts (a) the rank vocab slices tile `[0, vocab)` exactly once via `torch.equal` on a coverage counter, and (b) each rank's local streamed grad matches the full-vocab single-process autograd reference columns **within fp32 tolerance** via `torch.testing.assert_close(atol=1e-5, rtol=1e-4)` — a tolerance is correct here because the cross-rank all-reduce reorders the fp32 reduction vs. the single-tensor reference (matching the existing GPU test's grad tolerance), while the no-overlap tiling check uses exact equality. Spawned ranks set the conftest-mandated NCCL env (`NCCL_CUMEM_ENABLE=0`, etc.) before `init_process_group`, since `mp.spawn` children do not inherit the runtime env set by the CI conftest.
Run with (>=2 free GPUs; the `tp_size=4` case is skipped unless 4 are present):
```
uv run --isolated --extra dev --extra megatron -- \
pytest -s tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_chunked_logprob_backward_tp.py
```
## Scope & follow-ups
- **Backends covered:** Megatron only, by nature — `ChunkedDistributedLogprob` is the vocab-parallel (TP/CP) chunked-logprob used on the Megatron worker, reached through the single `forward_backward_mini_batch` train-step chokepoint, so all Megatron-backed training pathways (SFT, sync/async/fully-async RL, full-context) inherit the optimization with no per-pathway edit. FSDP has no vocab-parallel chunked-logprob equivalent, so it is intentionally out of scope.
- **Scope:** Limited to the single `backward` method; `DistributedLogprob`, `forward()`, and the fused-LCE path are unchanged.
- **Deferred:** A direct streaming-vs-old-`cat` comparison at TP>1 was intentionally not added, since it would require shipping the pre-refactor `cat`-based backward as dead reference code. The byte-identity gate for the refactor itself lives on the CPU `world_size=1` lane (`torch.equal`); the GPU lane validates the distributed math against the reference (with tolerance) plus exact no-overlap tiling.
## Related PRs (merge ordering)
- **NovaSky-AI#1765** (fused LM-head log-prob + entropy) edits the same `ChunkedDistributedLogprob.backward` and **keeps** the `all_grad_input = []` / `torch.cat` form. It refactors the per-chunk chosen-token scatter-add into a shared `_add_chosen_token_grad` helper — touching the exact lines this PR renames to `chunk_grad_input` and streams — so the two **will textually collide on this hunk**. NovaSky-AI#1765 does **not** subsume this PR's peak-memory win (it preserves the list+`cat`), and NovaSky-AI#1765 is otherwise complementary in intent. Whichever lands first forces a rebase of the other. **Recommended ordering:** land NovaSky-AI#1765 first, then rebase this streaming change on top of it (the preallocated-buffer write replaces the `append`+`cat` NovaSky-AI#1765 keeps); if the two are reviewed together they can be folded into one change.
- **NovaSky-AI#1543** (WIP) independently removes the same `torch.cat` 2x-peak, but does so **in place** into the autograd-saved logits with no extra buffer. This PR deliberately does **not** take that approach: the in-place variant downcasts the fp32 grad to the saved logits' dtype (breaking byte-identity) and mutates an autograd-saved tensor (the version-counter / double-backward hazard). This PR's separate fp32 buffer preserves numeric fidelity and autograd-safety at the cost of ~`(1/num_chunks)`x extra memory. NovaSky-AI#1543 is also built on pre-merge code and currently conflicts with main, so it is not a viable supersession.
…rializing [*, seq, vocab] logits
Adds an opt-in fused LM-head log-prob path for the Megatron backend that folds the
vocab projection into the chunked log-prob so the full [*, seq, vocab // TP] logits
tensor is never materialized — the dominant activation transient on the SFT/logprob
path at large vocab / long context. Numerically equivalent to the stock logits path.
Design ported from verl's fused linear-cross-entropy. Two backends behind one flag:
* "torch" (default): a pure-PyTorch vocab-parallel FusedLinearLogprob autograd
Function. Chunks over the sequence dim, recomputes the per-chunk logits in
backward, fp32 softmax math, masks out-of-shard targets. Runs anywhere (CPU/GPU),
no extra deps. Verified bit-exact (fwd + grad-hidden) and fp32-rounding-exact
(grad-weight) vs the stock logits path on a real gloo TP=1/TP=2 group.
* "triton": vendored flash-style Triton kernel (from volcengine/verl, Apache-2.0)
that tiles over the vocab dim so per-chunk logits never materialize (lower memory,
faster). GPU + triton required; falls back to "torch" with a warning otherwise.
Integration (model-source-free; no megatron-core edits):
* MegatronModelWrapper wraps output_layer.forward to return the pre-projection
hidden state (skipping the vocab GEMM) and capture the LM-head weight shard. The
weight is resolved from the weight= kwarg the model passes — i.e.
shared_embedding_or_output_weight() — so tied embeddings work with no None-weight
crash. Sequence-parallel hidden is gathered with tensor_parallel_output_grad=True,
whose backward reduce-scatters grad_hidden across TP exactly as ColumnParallelLinear
does, so the fused function's per-shard grad_hidden is reduced correctly.
* Both the training loss_func and the forward-only collection_func (eval / reference
logprobs) route through the fused path; the vocab range is derived from the captured
weight shard, never from hidden.shape[-1].
* Falls back to the stock logits path automatically when the model uses MuP output
scaling (the fused path bypasses _scale_logits) and is restricted to the
cross_entropy loss (the RL entropy/KL terms still need full logits).
Config: trainer.fused_linear_logprob (bool, default False -> byte-identical to today)
and trainer.fused_linear_logprob_backend ("torch" | "triton"). Megatron-only.
Tests: tests/.../cpu/megatron/test_fused_linear_logprob.py (CPU gloo TP=1/TP=2,
with/without out-of-shard targets, chunked + single-chunk) asserts the fused path
matches the stock logits path for forward, grad-hidden, and grad-weight. A GPU test
covers the Triton backend.
The vendored Triton kernel preserves verl's out-of-bounds-vocab -inf masking fix
(verl-project/verl#2656) and reproduces the NVIDIA + Bytedance Apache-2.0 attribution
(see file header + new top-level NOTICE).
Signed-off-by: dyurk-lila <dyurk@lila.ai>
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…ision policy
Validated on 2xH100 (buzz): all 16 GPU cases pass (TP1/TP2 × chunk {8,1000} ×
{fp32,bf16} × with/without out-of-shard targets).
- Fix an out-of-vocab gradient bug in the Triton adapter backward: it was zeroing
grad_output at fully-OOV positions, but the reference (stock DistributedLogprob /
the pure-torch FusedLinearLogprob) produces d_logits = -softmax * grad_output
there (no owned chosen column => the onehot term drops, the -softmax term remains).
Zeroing diverged the gradient (~0.04 abs). Now grad_output passes through unchanged;
the forward still masks the OOV logprob value to 0. Matches the torch backend exactly.
- Precision policy: production uses the fast default matmul precision (TF32 on Hopper) —
torch/triton's default, and the model trains in bf16 where it's empirically fine
(bf16 cases match the stock path within ~3e-3). A module-level FORCE_FP32_IEEE_PRECISION
flag (default False) lets the test pin IEEE for fp32 cases as a reproducible
"kernel math is exact" sanity check. Measured at TP=1 fp32: default precision fwd
error 1.27e-3 (TF32 rounding), IEEE 9.5e-7 (exact) — confirming the only default-path
gap is tensor-core rounding, not a math error. apply() signature unchanged.
Signed-off-by: dyurk-lila <dyurk@lila.ai>
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…gprob adapter The Triton kernel's tl.dot(hidden, weight.T) requires both operands to share a dtype. On a real hybrid model the pre-head hidden state is fp32 (final norm) while the LM-head weight is bf16, so the adapter crashed with "AssertionError: Both operands must be same dtype. Got fp32 and bf16" on the first real training step (the fp32/fp32 and bf16/bf16 test cases never exposed it). Mirror the pure-torch FusedLinearLogprob (hidden.to(weight.dtype) @ weight.t()) and ColumnParallelLinear (which casts its input to the weight dtype before the bf16 GEMM): cast hidden to weight.dtype before the kernel, and cast the returned d_hidden back to the original hidden dtype so autograd accumulates it into the fp32 hidden leaf correctly. No-op when dtypes already match. Signed-off-by: dyurk-lila <dyurk@lila.ai> Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…ropy GPU test The new entropy parity test crashed in its REFERENCE (vocab_parallel_entropy), which all-reduces over mpu.get_tensor_model_parallel_group() — uninitialized because the worker only called dist.init_process_group, never mpu.initialize_model_parallel. (The logprob test sidesteps this by referencing dist.group.WORLD directly.) Initialize model-parallel state with tensor_model_parallel_size=world_size (TP group == WORLD membership, so both the stock reference and the fused path reduce over identical ranks) and destroy it in the finally. Harness-only fix; the Triton kernel/adapter were never reached by the failing assertion. Signed-off-by: dyurk-lila <dyurk@lila.ai> Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Polish the fused linear log-prob feature so it reads and behaves as a native
SkyRL feature, with no internal/development-history references and CI-correct
tests. No functional change to the fused math (CPU gloo parity re-verified at
TP=1/2, fwd bit-exact, grads at fp32 round-off).
CI correctness:
* GPU Triton test: add pytest.mark.megatron (the megatron GPU job selects by
`-m megatron`, so the test was collected but never run); combine with the
existing skipif in a module-level pytestmark list.
* CPU test: stub megatron.core.parallel_state at module import scope when
megatron-core is absent, so collection no longer crashes the CPU CI job
(which installs no --extra megatron) and the pure-torch path actually runs.
The stub is import-scope so it also reaches the mp.spawn workers. Mark the
module pytest.mark.megatron for selection where megatron is installed.
Cleanup:
* Scrub development-history / out-of-band comments ("original downstream
monkey-patch", "accompanying review message", "the original bug", "BUG 2").
* Remove dead code: unused GPU-test helpers (_stock_logprobs/_fused_logprobs)
and their now-unused import; unused partition_vocab_size in
FusedLinearLogprob.forward; drop fully_oov from the Triton save_for_backward
(it was saved but never read in backward).
* DRY: extract the shared chosen-token scatter-add used by both
ChunkedDistributedLogprob.backward and FusedLinearLogprob.backward into
_add_chosen_token_grad so the OOV/index convention stays in lockstep.
* Cache exp(shifted) once in _distributed_log_softmax_and_entropy.
* Fix stale/elided test-path doc pointers and the MuP-fallback pointer; add the
missing lm_head_weight/fused_backend/return_entropy Args to the packed-seq
log-prob docstring.
Formatting: ruff 0.11.9 + black 24.10.0 clean.
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>
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>
dyurk-lila
force-pushed
the
feat/fused-linear-cross-entropy
branch
from
ef0bf7e to
118af6f
Compare
erictang000
reviewed
Jul 3, 2026
erictang000
left a comment
erictang000
left a comment
Collaborator
There was a problem hiding this comment.
hey @dyurk-lila, thanks for rebasing on the other PR! Looking a lot cleaner
just had some questions about some test failures/entropy computation and some potential bugs post-rebase
| chunk_size, | ||
| tp_group, | ||
| inference_only, | ||
| False, # entropy is computed by the no-grad fused entropy helpers |
Collaborator
There was a problem hiding this comment.
does this argument exist? i don't see it in FusedLinearLogprobTriton I'm assuming it should be something like don't return entropy, but FusedLinearLogprobTriton seems to unconditionally return entropy
| @@ -0,0 +1,299 @@ | |||
| """GPU parity tests for ``FusedLinearLogprobTriton``. | |||
Collaborator
There was a problem hiding this comment.
getting these errors when running locally:
=========================================================================== short test summary info ===========================================================================
FAILED tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_fused_linear_logprob_triton.py::test_fused_triton_entropy_matches_stock[False-1-fp32] - torch.multiprocessing.spawn.ProcessRaisedException:
FAILED tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_fused_linear_logprob_triton.py::test_fused_triton_entropy_matches_stock[False-1-bf16] - torch.multiprocessing.spawn.ProcessRaisedException:
FAILED tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_fused_linear_logprob_triton.py::test_fused_triton_entropy_matches_stock[False-2-fp32] - torch.multiprocessing.spawn.ProcessRaisedException:
FAILED tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_fused_linear_logprob_triton.py::test_fused_triton_entropy_matches_stock[False-2-bf16] - torch.multiprocessing.spawn.ProcessRaisedException:
FAILED tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_fused_linear_logprob_triton.py::test_fused_triton_entropy_matches_stock[True-1-fp32] - torch.multiprocessing.spawn.ProcessRaisedException:
FAILED tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_fused_linear_logprob_triton.py::test_fused_triton_entropy_matches_stock[True-1-bf16] - torch.multiprocessing.spawn.ProcessRaisedException:
FAILED tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_fused_linear_logprob_triton.py::test_fused_triton_entropy_matches_stock[True-2-fp32] - torch.multiprocessing.spawn.ProcessRaisedException:
FAILED tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_fused_linear_logprob_triton.py::test_fused_triton_entropy_matches_stock[True-2-bf16] - torch.multiprocessing.spawn.ProcessRaisedException:
============================================================ 8 failed, 16 passed, 5 warnings in 237.01s (0:03:57) =====
> raise ProcessRaisedException(msg, error_index, failed_process.pid)
E torch.multiprocessing.spawn.ProcessRaisedException:
E
E -- Process 1 terminated with the following error:
E Traceback (most recent call last):
E File "/home/etang/.cache/uv/builds-v0/.tmpmFwpCn/lib/python3.12/site-packages/torch/multiprocessing/spawn.py", line 87, in _wrap
E fn(i, *args)
E File "/home/etang/SkyRL/tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_fused_linear_logprob_triton.py", line 243, in _entropy_worker
E ent_fused, gh_fused, gw_fused = _direct_fused_entropy(
E ^^^^^^^^^^^^^^^^^^^^^^
E File "/home/etang/SkyRL/tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_fused_linear_logprob_triton.py", line 187, in _direct_fused_entropy
E lp, ent = FusedLinearLogprobTriton.apply(
E ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
E File "/home/etang/.cache/uv/builds-v0/.tmpmFwpCn/lib/python3.12/site-packages/torch/autograd/function.py", line 596, in apply
E return super().apply(*args, **kwargs) # type: ignore[misc]
E ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
E TypeError: FusedLinearLogprobTriton.forward() takes from 8 to 9 positional arguments but 10 were given
../.cache/uv/builds-v0/.tmpmFwpCn/lib/python3.12/site-packages/torch/multiprocessing/spawn.py:211: ProcessRaisedExcepti
per the other comment I think there's some left over incompatibility with the previous implementation you had which returned both logprobs and entropy from and the FusedLinearChunkedDistributedLogprob which only returns logprobs, can you take a look?
I think leaving entropy with fused lm head for both torch and triton as no_grad only is fine for now, we should just track it as an issue and leave a comment
Collaborator
There was a problem hiding this comment.
hmm yeah can we leave this off for now and just stick to the repo wide convention of providing notice at the top of the vendored file?
| "nvidia-modelopt; sys_platform == 'linux'", | ||
| ] | ||
|
|
||
| # Optional backend for trainer.fused_lm_head_logprob_backend="triton". |
Collaborator
There was a problem hiding this comment.
i think triton should be included by default by both the fsdp and megatron backends if i'm not mistaken, could you double check if this change is needed?
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
3 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.
Reviewers: Where to Look
Please review closely (the genuinely-new logic):
Triton backend selector + dispatch —
skyrl/backends/skyrl_train/distributed/megatron/model_utils.py, new_fused_lm_head_logprob_apply(...). This is the one seam that chooses between upstream's torchFusedLinearChunkedDistributedLogproband the vendored Triton kernel, with a safe fallback (warns + uses torch) when Triton/CUDA is absent. It is threaded via a newfused_backendkwarg intofrom_parallel_hidden_to_logprobsandfrom_parallel_hidden_to_logprobs_packed_sequences. Confirm the fallback path and that the dispatcher's forward contract matches upstream's Function exactly (same 8 args, single-tensor logprob return).Vendored Triton kernel adapter —
skyrl/backends/skyrl_train/distributed/megatron/fused_linear_logprob_triton.py, classFusedLinearLogprobTriton(the SkyRL adapter, ~from line 1905). Review the adapter (the SkyRL-shaped autograd Function): the coordinate re-keying for arbitrary shard layouts, the fully-OOV forward masking / backward pass-through (#2656), the hidden→weight-dtype cast, and thereturn_entropy/set_materialize_grads(False)handling. The vendored verl/NVIDIA kernel internals below it need less scrutiny (attribution in the file header +NOTICE).Backend config wiring —
skyrl/train/config/config.py(fused_lm_head_logprob_backendselector + validation) andskyrl/backends/skyrl_train/workers/megatron/megatron_model_wrapper.py(self._fused_lm_head_backend, passed to the four fused call sites). This reuses [megatron] fused linear cross-entropy for 23x memory savings at 262k context #1841'sfused_lm_head_logprobon/off flag and adds a composing backend selector — not a second on/off switch. Confirm it composes cleanly with [megatron] fused linear cross-entropy for 23x memory savings at 262k context #1841.Triton kernel GPU tests —
tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_fused_linear_logprob_triton.py. DrivesFusedLinearLogprobTriton.applydirectly against the stock (ChunkedDistributedLogprob/DistributedLogprob) path for logprob (fwd + both grads) and the Triton kernel'sreturn_entropy=Trueentropy path (fwd +dentropybackward) vs stockvocab_parallel_entropy. Markedpytest.mark.megatron.Needs less scrutiny: the vendored Triton kernel internals (Apache-2.0, attributed),
NOTICE,pyproject.toml(the optionaltritonextra), andskyrl/benchmarks/bench_fused_ce_h100.py.Deliberately deferred (documented): wiring entropy-with-gradient through the torch backend / the RL loss (
use_entropy_loss=Trueon the fused path) is not included here. Upstream #1841's fused path keeps entropy as a no-grad metric and this PR preserves that restriction, rather than porting untested grad-entropy math onto #1841's differently-structured Function. The Triton kernel's entropy dual-output is validated by the GPU test above, so a future PR can wire it through the loss with confidence.Benchmark (single H100 80GB,
Qwen/Qwen3.6-35B-A3B)Forward+backward through the LM-head training signal — same harness as #1841's
bench_fused_linear_logprob.py, adapted to this PR's two backends.hidden=2048,chunk=1024, bf16, batch 1; peaktorch.cuda.max_memory_allocated, mean of 3 reps. Compared against megatron-core's eager (vocab_parallel_cross_entropy) and NVIDIA fused (fused_vocab_parallel_cross_entropy) CE, both of which materialize the[B, S, vocab//TP]logits.hidden @ weightᵀ→ eagervocab_parallel_cross_entropy(materializes logits + fp32 grad)fused_vocab_parallel_cross_entropy(@jit_fuser; fuses the CE stages, still materializes logits)FusedLinearLogprob(logits never materialized)FusedLinearLogprobTritonCorrectness (vs eager baseline, max abs
grad_hiddenerror): nvidia3.05e-05, fused-torch4.88e-04, fused-triton9.77e-04— all within bf16/fp32 rounding.Per-rank vocab shard = 62,080 (the real TP=4 shard)
Full vocab = 248,320
Takeaways, consistent with #1841: the NVIDIA fused CE is a compute win but a memory loss (~1.7× the eager baseline → OOMs earlier at long context). Both fused-linear backends are memory wins that fit contexts the logits-materializing paths cannot (baseline/nvidia OOM at 256k even on the TP=4 shard, and at ≥64k full-vocab). On H100 the Triton backend is the clear pick — it's both the lowest-memory and fastest fused option (e.g. 32k full-vocab: 3.0 GB / 418 ms vs the torch path's 9.1 GB / 1,798 ms, and ~comparable wall-clock to baseline where baseline still fits). The pure-torch backend keeps the lowest memory floor and remains the portable, dependency-free fallback.
Summary
Adds an opt-in fused LM-head path for the Megatron backend that folds the vocab projection into the chunked log-prob (and entropy), so the full
[*, seq, vocab // TP]logits tensor is never materialized. That tensor is the dominant activation transient on the log-prob path at large vocab + long context (e.g. ~4 GiB fp32 at micro-bs 2 / 32k ctx / 131k vocab / TP4); eliminating it unlocks larger micro-batches / no-recompute without OOM.Serves SFT and all RL losses (cross_entropy, ppo, grpo, …): the fused kernel returns
(log_probs, entropy)— verl's established dual-output — so RL's entropy term is satisfied from the fused output instead of needing the full logits. Numerically equivalent to the stock logits path; default OFF → byte-identical to today.Design ported from verl's fused linear-cross-entropy (
volcengine/verl, Apache-2.0). Two backends behind one flag:torch(default): pure-PyTorch, vocab-parallelFusedLinearLogprob. Chunks over the seq dim, recomputes per-chunk logits in backward, fp32 softmax/entropy math, masks out-of-shard targets. Runs anywhere (CPU/GPU), no extra deps.triton: vendored flash-style Triton kernel — tiles over the vocab dim so per-chunk logits never materialize (lower memory floor + faster). GPU +tritonrequired; falls back totorchwith a warning if unavailable.Config
Megatron-only. Uses the existing
logprobs_chunk_sizefor the chunk width.How it integrates (no megatron-core edits)
MegatronModelWrapperwrapsoutput_layer.forwardat model-build to return the pre-projection hidden state (skipping the vocab GEMM) and capture the LM-head weight shard. Correctness points:cross_entropy) requests logprob-only (skips the entropy all-reduce); RL requestsreturn_entropy=True. The fused per-token entropy =logsumexp(logits) − Σ(softmax·logits), algebraically identical to stock_VocabParallelEntropyand verl; backward adds the matchingdentropyterm. RL entropy is sliced to the action window (non-packed) or fed throughvocab_parallel_entropy_packed_sequencesvia a newprecomputed_entropy_tokensarg (packed action-weighting + CP all-reduce unchanged — verified a bit-identical substitution). KL keeps using logprobs.hidden→ LM-headweight.dtypebefore the GEMM (the common hybrid case is fp32 pre-head hidden + bf16 head), mirroringColumnParallelLinear; grad cast back on the way out.weight=kwarg the model passes (shared_embedding_or_output_weight()) → nooutput_layer.weight is Nonecrash (cf. verl#3730).tensor_parallel_output_grad=True; its reduce-scatter backward performs the TP reduction ofgrad_hiddenexactly as stockColumnParallelLinear. Vocab range derived from the captured weight shard, nothidden.shape[-1].loss_funcand forward-onlycollection_funcroute through the fused path. Auto-fallback to stock under MuP output scaling; default fused-off is byte-identical.Verification
tests/backends/skyrl_train/cpu/megatron/test_fused_linear_logprob.py): gloo TP1/TP2, chunked + single, with/without OOV. Logprob fwd + grad-hidden bit-exact, grad-weight to fp32 rounding vs the stock logits path; entropy fwd + backward parity vs stockvocab_parallel_entropy(~5e-7), entropy target-independence, packedprecomputed_entropy_tokensbit-identical substitution, slice-alignment guard. Runs on CPU CI (no GPU / no megatron-core needed —megatron.core.parallel_stateis stubbed when absent).tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_fused_linear_logprob_triton.py), 2×H100: 24/24 pass. Logprob 16/16 (incl. fp32+bf16 dtype-mix) + entropy 8/8 (fp32 ~5e-7 / ~1e-8 grads under IEEE; bf16 within tolerance; OOV-invariant). Marked@pytest.mark.megatronso the megatron GPU CI job (-m megatron) selects it. The fp32 cases pin IEEE (FORCE_FP32_IEEE_PRECISION) as a math-is-exact sanity check (default-precision fwd error 1.27e-3 TF32 vs 9.5e-7 IEEE); production uses the fast default (TF32) precision — the model trains in bf16 where it matches stock within ~3e-3.Licensing / attribution
verl is Apache-2.0 (same as SkyRL); SkyRL already vendors verl in ~8 files with the
# This code is adapted from VERL … original copyright reproduced below:convention. The vendored files (verl/utils/kernel/{kernels.py,linear_cross_entropy.py}) carry a dual copyright —Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATESandCopyright 2024 Bytedance Ltd. and/or its affiliates(verl is a Bytedance/volcengine project; the NVIDIA line is an NVIDIA-contributed portion). Both are reproduced verbatim in the vendored file header with aChanges:note (§4(b)). Adds a top-levelNOTICE(SkyRL had none) reproducing verl'sNotice.txt. Preserves verl's out-of-bounds-vocab-infmasking fix (verl#2656).Test plan
uv run --isolated --extra dev -- pytest -s tests/backends/skyrl_train/cpu/megatron/test_fused_linear_logprob.py(CPU, no megatron-core)pytest -s tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_fused_linear_logprob_triton.py -m megatronon 2×H100 — 24/24ruff+blackclean (pre-commit)fused_linear_logprob: falsepath byte-identical to current behaviour🤖 Generated with Claude Code