diff --git a/skyrl/backends/skyrl_train/distributed/megatron/model_utils.py b/skyrl/backends/skyrl_train/distributed/megatron/model_utils.py index 2d37fbf0d3..626edc1b87 100644 --- a/skyrl/backends/skyrl_train/distributed/megatron/model_utils.py +++ b/skyrl/backends/skyrl_train/distributed/megatron/model_utils.py @@ -214,10 +214,16 @@ def backward( seq_size = int(vocab_parallel_logits.shape[1]) num_chunks = (seq_size + chunk_size - 1) // chunk_size - all_grad_input = [] - batch_size = int(vocab_parallel_logits.shape[0]) + # Stream chunk grads into a preallocated buffer instead of keeping every + # chunk alive until torch.cat. Each chunk owns one contiguous seq slice. + grad_input = torch.empty( + (batch_size, seq_size, partition_vocab_size), + dtype=torch.float32, + device=vocab_parallel_logits.device, + ) + for chunk_idx in range(num_chunks): chunk_start = chunk_idx * chunk_size chunk_end = min(seq_size, (chunk_idx + 1) * chunk_size) @@ -251,15 +257,14 @@ def backward( flat_chosen = flat_idx.masked_select(valid_mask.reshape(-1)) + chunk_masked_target.masked_select(valid_mask) # `neg` is zero-copy; the subsequent mul_ writes in place. - grad_input = softmax_output.neg_() - grad_input.mul_(chunk_grad_output.unsqueeze(-1)) + chunk_grad_input = softmax_output.neg_() + chunk_grad_input.mul_(chunk_grad_output.unsqueeze(-1)) grad_output_selected = chunk_grad_output.masked_select(valid_mask) - grad_input.view(-1).scatter_add_(0, flat_chosen, grad_output_selected) - - all_grad_input.append(grad_input) + chunk_grad_input.view(-1).scatter_add_(0, flat_chosen, grad_output_selected) - grad_input = torch.cat(all_grad_input, dim=1) + # Write this chunk into its non-overlapping sequence slice. + grad_input[:, chunk_start:chunk_end, :] = chunk_grad_input # if you add an argument to the forward method, then you must add a corresponding None here return grad_input, None, None, None, None, None, None diff --git a/tests/backends/skyrl_train/distributed/test_chunked_logprob_backward_streaming.py b/tests/backends/skyrl_train/distributed/test_chunked_logprob_backward_streaming.py new file mode 100644 index 0000000000..f554bae8d5 --- /dev/null +++ b/tests/backends/skyrl_train/distributed/test_chunked_logprob_backward_streaming.py @@ -0,0 +1,116 @@ +"""Targeted CPU regression test for streamed ``ChunkedDistributedLogprob.backward``. + +Run with: + uv run --isolated --extra skyrl-train --extra dev -- pytest -s \ + tests/backends/skyrl_train/distributed/test_chunked_logprob_backward_streaming.py +""" + +import os +import sys +from types import ModuleType + +import pytest +import torch +import torch.distributed as dist + +from skyrl.backends.skyrl_train.distributed.utils import get_free_port + +# Stub megatron so CPU CI can import model_utils without megatron-core. +# The fixture restores prior modules, leaving GPU lanes with real megatron intact. + +_MEGATRON_MODULES = [ + "megatron", + "megatron.core", + "megatron.core.parallel_state", +] + +_mock_modules: dict[str, ModuleType] = {} +for _name in _MEGATRON_MODULES: + _mock_modules[_name] = ModuleType(_name) + +_mock_modules["megatron.core"].parallel_state = _mock_modules["megatron.core.parallel_state"] + + +@pytest.fixture(scope="module", autouse=True) +def _stub_megatron_modules(): + """Install the mock ``megatron`` modules for this module only.""" + saved = {_name: sys.modules.get(_name) for _name in _MEGATRON_MODULES} + for _name in _MEGATRON_MODULES: + sys.modules[_name] = _mock_modules[_name] + try: + yield + finally: + for _name in _MEGATRON_MODULES: + if saved[_name] is None: + sys.modules.pop(_name, None) + else: + sys.modules[_name] = saved[_name] + + +@pytest.fixture(scope="module") +def tp_group(): + """Single-rank gloo TP group; only destroy it if this fixture created it.""" + initialized_here = False + if not dist.is_initialized(): + os.environ["MASTER_ADDR"] = "localhost" + os.environ["MASTER_PORT"] = str(get_free_port()) + os.environ["RANK"] = "0" + os.environ["WORLD_SIZE"] = "1" + dist.init_process_group(backend="gloo", rank=0, world_size=1) + initialized_here = True + yield dist.group.WORLD + if initialized_here and dist.is_initialized(): + dist.destroy_process_group() + + +def _backward_grad(func_cls, logits, target, vocab_start, vocab_end, tp_group, *, chunk_size=None): + """Return the input grad using a non-uniform upstream gradient.""" + leaf = logits.detach().clone().requires_grad_(True) + if chunk_size is None: + out = func_cls.apply(leaf, target, vocab_start, vocab_end, tp_group, False) + else: + out = func_cls.apply(leaf, target, vocab_start, vocab_end, chunk_size, tp_group, False) + grad_seed = torch.linspace(0.5, 1.5, steps=out.numel(), device=out.device, dtype=out.dtype).reshape(out.shape) + out.backward(grad_seed) + return leaf.grad.detach() + + +@pytest.mark.parametrize( + "case", + [ + pytest.param((3, 17, 64, 5, "mixed_oov"), id="ragged_mixed_oov"), + pytest.param((2, 8, 32, 64, "all_oov"), id="single_chunk_all_oov"), + ], +) +def test_streamed_backward_matches_non_chunked_for_targeted_cases(tp_group, case): + """Chunked backward stays bit-identical while writing into the streamed buffer.""" + from skyrl.backends.skyrl_train.distributed.megatron.model_utils import ( + ChunkedDistributedLogprob, + DistributedLogprob, + ) + + batch_size, seq_len, vocab_size, chunk_size, target_mode = case + device = torch.device("cpu") + torch.manual_seed(1) + + logits = torch.randn(batch_size, seq_len, vocab_size, dtype=torch.float32, device=device) * 2.0 + if target_mode == "all_oov": + target = torch.full((batch_size, seq_len), vocab_size + 5, device=device, dtype=torch.long) + else: + target = torch.randint(0, vocab_size, (batch_size, seq_len), device=device, dtype=torch.long) + target[:, ::3] = vocab_size + 5 + + grad_ref = _backward_grad(DistributedLogprob, logits, target, 0, vocab_size, tp_group) + grad_chunk = _backward_grad( + ChunkedDistributedLogprob, + logits, + target, + 0, + vocab_size, + tp_group, + chunk_size=chunk_size, + ) + + assert grad_chunk.shape == grad_ref.shape == logits.shape + assert grad_chunk.dtype == torch.float32 + assert torch.equal(grad_chunk, grad_ref), "streamed chunked grad must be bit-identical to non-chunked grad" diff --git a/tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_chunked_logprob_backward.py b/tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_chunked_logprob_backward.py index cd3fc8630d..845a14aaac 100644 --- a/tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_chunked_logprob_backward.py +++ b/tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_chunked_logprob_backward.py @@ -12,7 +12,7 @@ ChunkedDistributedLogprob, DistributedLogprob, ) -from skyrl.train.utils.utils import get_free_port +from skyrl.backends.skyrl_train.distributed.utils import get_free_port @pytest.fixture(scope="module")