[WIP] Add changes needed for FP8 megatron training#1543
Open
pcmoritz wants to merge 2 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 several optimizations and fixes for distributed training, including memory-efficient gradient computation, FP8 alignment adjustments, and enhanced sequence parallelism support. Feedback focuses on correcting a potential typo in the CUDA version for nvidia-nvshmem, increasing the chunk_size for log-probability calculations to mitigate communication overhead, and vectorizing loops used for attention mask reconstruction to improve performance.
| "hf_transfer", | ||
| "cloudpathlib>=0.23.0", | ||
| # "deep-ep", | ||
| "nvidia-nvshmem-cu13>=3.6.5; sys_platform == 'linux'", |
Contributor
There was a problem hiding this comment.
Potential typo: nvidia-nvshmem-cu13 refers to CUDA 13, which is not yet released. You likely intended to use nvidia-nvshmem-cu12 or a version compatible with your current CUDA environment (e.g., cu128 as seen in your index URL).
| inference_only=True, | ||
| cp_group=None, # we handle cp gathering in `postprocess_packed_seqs` | ||
| chunk_size=None, | ||
| chunk_size=16, |
Contributor
There was a problem hiding this comment.
The hardcoded chunk_size=16 is extremely small. Since the log-probability computation involves multiple all_reduce operations per chunk (3 per chunk in the current implementation), such a small value will lead to significant communication overhead. Consider increasing this to a larger value (e.g., 128 or 256) to improve performance while maintaining memory efficiency.
Comment on lines
+179
to
+185
| gather_attention_mask = torch.zeros( | ||
| outputs.shape[0], gathered_seq_len, | ||
| dtype=new_attention_mask.dtype, device=new_attention_mask.device | ||
| ) | ||
| seq_lens = new_attention_mask.sum(dim=1) | ||
| for i in range(outputs.shape[0]): | ||
| gather_attention_mask[i, :seq_lens[i]] = 1 |
Contributor
There was a problem hiding this comment.
This loop for reconstructing the attention mask after gathering can be vectorized to improve performance, especially as the batch size or sequence length increases.
seq_lens = new_attention_mask.sum(dim=1, keepdim=True)
gather_attention_mask = (torch.arange(outputs.shape[1], device=new_attention_mask.device) < seq_lens).to(new_attention_mask.dtype)| inference_only=False, | ||
| cp_group=None, # we handle cp gathering in `postprocess_packed_seqs` | ||
| chunk_size=None, | ||
| chunk_size=16, |
Contributor
There was a problem hiding this comment.
The hardcoded chunk_size=16 is extremely small. Since the log-probability computation involves multiple all_reduce operations per chunk, such a small value will lead to significant communication overhead. Consider increasing this to a larger value (e.g., 128 or 256) to improve performance while maintaining memory efficiency.
Comment on lines
+493
to
+499
| gather_attention_mask = torch.zeros( | ||
| outputs.shape[0], gathered_seq_len, | ||
| dtype=new_attention_mask.dtype, device=new_attention_mask.device | ||
| ) | ||
| seq_lens = new_attention_mask.sum(dim=1) | ||
| for i in range(outputs.shape[0]): | ||
| gather_attention_mask[i, :seq_lens[i]] = 1 |
Contributor
There was a problem hiding this comment.
This loop for reconstructing the attention mask after gathering can be vectorized to improve performance, especially as the batch size or sequence length increases.
seq_lens = new_attention_mask.sum(dim=1, keepdim=True)
gather_attention_mask = (torch.arange(outputs.shape[1], device=new_attention_mask.device) < seq_lens).to(new_attention_mask.dtype)
![devin-ai-integration[bot]](https://avatars.githubusercontent.com/in/811515?s=60&v=4){kind=small}
| all_grad_input.append(grad_input) | ||
|
|
||
| grad_input = torch.cat(all_grad_input, dim=1) | ||
| vocab_parallel_logits[:, chunk_start:chunk_end, :] = chunk_grad.to(vocab_parallel_logits.dtype) |
Contributor
There was a problem hiding this comment.
🔴 In-place overwrite of saved logits tensor in ChunkedDistributedLogprob.backward() crashes when use_entropy_loss=True
The new ChunkedDistributedLogprob.backward() overwrites vocab_parallel_logits in-place with gradient data and returns it. This tensor is the original model output logits, which is also saved by _VocabParallelEntropy (via ctx.save_for_backward at model_utils.py:561) when use_entropy_loss=True. Both autograd functions share the same underlying tensor storage (since action_logits = logits[:, -num_actions - 1 : -1, :] at megatron_model_wrapper.py:349 is a view of logits). During backward, whichever function runs second will encounter a PyTorch RuntimeError: one of the variables needed for gradient computation has been modified by an inplace operation because ctx.saved_tensors detects the version mismatch caused by the first function's in-place modification.
Previously this was not an issue because chunk_size=None was used, which routed through DistributedLogprob — that class saves softmax_output (a derived tensor) rather than the original logits (model_utils.py:101). The PR switches to chunk_size=16 (megatron_model_wrapper.py:292), forcing all training paths through ChunkedDistributedLogprob, which saves the original logits (model_utils.py:197) and now destructively overwrites them in backward.
Prompt for agents
The in-place write `vocab_parallel_logits[:, chunk_start:chunk_end, :] = chunk_grad.to(vocab_parallel_logits.dtype)` at model_utils.py:240 destructively overwrites the saved logits tensor, which conflicts with _VocabParallelEntropy (model_utils.py:561-572) that also saves a view of the same tensor. When both are active (use_entropy_loss=True), PyTorch's saved-tensor version check causes a RuntimeError in backward.
Possible fixes:
1. Revert to the previous approach of collecting chunk gradients into a list and concatenating them (avoids in-place modification of the saved tensor entirely).
2. Allocate a separate output tensor of the same shape as vocab_parallel_logits and write chunk gradients into it, then return that new tensor.
3. Use `vocab_parallel_logits.clone()` at the start of backward to create a separate buffer for writing gradients, leaving the original saved tensor untouched.
Option 2 is closest to the memory-saving intent of this PR while avoiding the version conflict. The key files involved are model_utils.py (ChunkedDistributedLogprob.backward, lines 217-243) and megatron_model_wrapper.py (loss_func which calls both from_parallel_logits_to_logprobs and vocab_parallel_entropy).
Was this helpful? React with 👍 or 👎 to provide feedback.
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
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
1 participant
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.
Uh oh!
There was an error while loading. Please reload this page.