[megatron] Triton backend for fused LM-head log-prob#1765
Merged
erictang000 merged 11 commits intoJul 7, 2026
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
Contributor
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.
dyurk-lila
force-pushed
the
feat/fused-linear-cross-entropy
branch
from
July 1, 2026 20:53
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
Contributor
Author
There was a problem hiding this comment.
Done, reverted the entropy stuff to match the current pattern
| @@ -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
Contributor
Author
There was a problem hiding this comment.
Done, reverted the entropy stuff to match the current pattern
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?
Contributor
Author
There was a problem hiding this comment.
done, removed this
dyurk-lila
changed the title
…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
July 7, 2026 17:02
7d7b96d to
b89a5ee
Compare
erictang000
approved these changes
Jul 7, 2026
erictang000
left a comment
erictang000
left a comment
Collaborator
There was a problem hiding this comment.
looking good, thanks @dyurk-lila!
dyurk-lila
added a commit
to dyurk-lila/SkyRL
that referenced
this pull request
Jul 13, 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.
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:
Triton backend selector + dispatch -
skyrl/backends/skyrl_train/distributed/megatron/model_utils.py,_fused_lm_head_logprob_apply(...). This chooses between upstreamFusedLinearChunkedDistributedLogproband the vendored Triton adapter. Both backends now share the same logprob-only apply contract and return TP-combined[B, S]log-probs.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.Backend config wiring -
skyrl/train/config/config.pyandskyrl/backends/skyrl_train/workers/megatron/megatron_model_wrapper.py. This addstrainer.fused_lm_head_logprob_backend: "torch" | "triton"as a backend selector behind the existingtrainer.fused_lm_head_logprobon/off flag.Triton GPU tests -
tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_fused_linear_logprob_triton.py. These driveFusedLinearLogprobTriton.applydirectly 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
NOTICEand standalonetritonextra 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
torchbackend remains upstreamFusedLinearChunkedDistributedLogprob;tritonuses 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:
If
tritonis 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=Trueon 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
Test plan
git diff --checkreturn_entropy/ dual-output / top-level NOTICE / old config-name claims remain inskyrl,tests, orpyproject.tomlPYTHONPYCACHEPREFIX=/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.pypytest -s tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_fused_linear_logprob_triton.py -m megatronon CUDA/H100Generated with Codex