diff --git a/skyrl/backends/skyrl_train/workers/worker.py b/skyrl/backends/skyrl_train/workers/worker.py index b56d56b518..ecc8ab0f8b 100644 --- a/skyrl/backends/skyrl_train/workers/worker.py +++ b/skyrl/backends/skyrl_train/workers/worker.py @@ -54,6 +54,7 @@ compute_minibatch_rollout_logprob_diff_metrics, get_microbatch_iterator, reduce_metrics, + restore_microbatch_response_padding, ) from skyrl.env_vars import ( SKYRL_RAY_PG_TIMEOUT_IN_S, @@ -870,6 +871,7 @@ def forward_backward( data, micro_batch_size=self.cfg.micro_train_batch_size_per_gpu, max_tokens_per_microbatch=self.cfg.max_tokens_per_microbatch, + trim_padding=self.cfg.strategy == "fsdp" and not self.cfg.remove_microbatch_padding, ) all_metrics = defaultdict(list) all_loss_fn_outputs = [] # Handle separately from scalar metrics @@ -1174,6 +1176,7 @@ def forward( data, micro_batch_size=self.cfg.micro_forward_batch_size_per_gpu, max_tokens_per_microbatch=self.cfg.max_tokens_per_microbatch, + trim_padding=self.cfg.strategy == "fsdp" and not self.cfg.remove_microbatch_padding, ) outputs = [self._forward_micro_batch(micro_batch) for micro_batch in microbatch_iterator] output = microbatch_iterator.reorder_and_combine_batches(outputs) @@ -1187,7 +1190,12 @@ def forward( all_metrics = defaultdict(list) all_loss_fn_outputs: List[Dict[str, Any]] = [] - for micro_batch in BatchIterator(data, micro_batch_size, drop_last=False): + for micro_batch in BatchIterator( + data, + micro_batch_size, + drop_last=False, + trim_padding=self.cfg.strategy == "fsdp" and not self.cfg.remove_microbatch_padding, + ): metrics = self._forward_micro_with_loss( micro_batch, loss_fn=loss_fn, @@ -1331,6 +1339,7 @@ def _forward_micro_batch(self, micro_batch: TrainingInputBatch) -> TrainingOutpu image_grid_thw=image_grid_thw, ) policy_logprob = policy_logprob.to("cpu") + policy_logprob = restore_microbatch_response_padding(policy_logprob, micro_batch.metadata) output = TrainingOutputBatch( {"output": policy_logprob}, ) @@ -1426,6 +1435,7 @@ def forward_backward(self, data: TrainingInputBatch) -> WorkerOutput: data, micro_batch_size=self.cfg.micro_train_batch_size_per_gpu, max_tokens_per_microbatch=self.cfg.max_tokens_per_microbatch, + trim_padding=self.cfg.strategy == "fsdp" and not self.cfg.remove_microbatch_padding, ) all_metrics = defaultdict(list) @@ -1555,6 +1565,7 @@ def _forward_micro_batch( ) self.model.train() # reset model state value = value.to("cpu") + value = restore_microbatch_response_padding(value, micro_batch.metadata) output = TrainingOutputBatch( {"output": value}, ) @@ -1574,6 +1585,7 @@ def forward(self, data: TrainingInputBatch) -> WorkerOutput: data, micro_batch_size=self.cfg.micro_forward_batch_size_per_gpu, max_tokens_per_microbatch=self.cfg.max_tokens_per_microbatch, + trim_padding=self.cfg.strategy == "fsdp" and not self.cfg.remove_microbatch_padding, ) outputs = [self._forward_micro_batch(micro_batch) for micro_batch in microbatch_iterator] output = microbatch_iterator.reorder_and_combine_batches(outputs) @@ -1605,6 +1617,7 @@ def forward(self, data: TrainingInputBatch) -> WorkerOutput: data, micro_batch_size=self.cfg.micro_forward_batch_size_per_gpu, max_tokens_per_microbatch=self.cfg.max_tokens_per_microbatch, + trim_padding=self.cfg.strategy == "fsdp" and not self.cfg.remove_microbatch_padding, ) outputs = [self._forward_micro_batch(micro_batch) for micro_batch in microbatch_iterator] output = microbatch_iterator.reorder_and_combine_batches(outputs) @@ -1632,6 +1645,7 @@ def _forward_micro_batch(self, micro_batch: TrainingInputBatch) -> TrainingOutpu image_grid_thw=image_grid_thw, ) log_probs = log_probs.to("cpu") + log_probs = restore_microbatch_response_padding(log_probs, micro_batch.metadata) output = TrainingOutputBatch( {"output": log_probs}, ) diff --git a/skyrl/backends/skyrl_train/workers/worker_utils.py b/skyrl/backends/skyrl_train/workers/worker_utils.py index 2efa37ad6c..909e2ae986 100644 --- a/skyrl/backends/skyrl_train/workers/worker_utils.py +++ b/skyrl/backends/skyrl_train/workers/worker_utils.py @@ -11,6 +11,30 @@ from skyrl.train.dataset.bin_packing import make_seq_packer from skyrl.train.dataset.replay_buffer import Experience +_SEQUENCE_FIELDS = frozenset( + { + "sequences", + "attention_mask", + "rollout_expert_indices", + "router_padding_mask", + } +) +_RESPONSE_FIELDS = frozenset( + { + "action_log_probs", + "base_action_log_probs", + "values", + "returns", + "advantages", + "kl", + "rewards", + "loss_mask", + "response_mask", + "rollout_logprobs", + } +) +_PADDED_RESPONSE_LENGTH = "response_length_before_microbatch_trim" + # Metrics that end in `_loss` but are plain per-token MEANS, not pre-scaled minibatch sums. # The `sum_loss_metrics` convention sums every `_loss` key because the *policy* losses are # pre-scaled (by num_microbatches * dp_size) so that summing recovers the correct minibatch @@ -130,11 +154,87 @@ def all_reduce_metrics( return status_mean +def trim_microbatch_padding(batch: TrainingInputBatch) -> TrainingInputBatch: + """Project a left-padded batch to the final microbatch's local widths. + + The controller keeps one rectangular representation for advantage calculation, + dispatch, and replay metadata. Dense FSDP forwards do not remove padding inside + the model, so they should receive a view padded only to the longest sequence and + response in the final microbatch. + + Packed rows and batches without an explicit ``response_mask`` retain their + original representation because their response boundary cannot be recovered + from ``loss_mask`` (which may contain semantic zeros). + """ + attention_mask = batch.get("attention_mask") + response_mask = batch.get("response_mask") + if ( + attention_mask is None + or response_mask is None + or len(batch) == 0 + or batch.get("sub_seq_lengths") is not None + or (batch.metadata or {}).get("is_padding_batch", False) + ): + return batch + if attention_mask.ndim != 2 or response_mask.ndim != 2: + raise ValueError( + "Expected 2D attention_mask and response_mask for microbatch trimming, " + f"got {attention_mask.shape} and {response_mask.shape}" + ) + + sequence_length = int(attention_mask.sum(dim=1).max().item()) + response_length = int(response_mask.sum(dim=1).max().item()) + if sequence_length <= 0 or response_length <= 0: + return batch + if sequence_length < response_length + 1: + raise ValueError( + f"Microbatch sequence length ({sequence_length}) must exceed response length ({response_length})" + ) + + sequence_width = attention_mask.shape[1] + response_width = response_mask.shape[1] + if sequence_length == sequence_width and response_length == response_width: + return batch + + projected = {} + for key, value in batch.items(): + if isinstance(value, torch.Tensor) and key in _SEQUENCE_FIELDS: + projected[key] = value[:, -sequence_length:] + elif isinstance(value, torch.Tensor) and key in _RESPONSE_FIELDS: + projected[key] = value[:, -response_length:] + else: + projected[key] = value + + microbatch = TrainingInputBatch(projected) + microbatch.metadata = dict(batch.metadata or {}) + microbatch.metadata[_PADDED_RESPONSE_LENGTH] = response_width + microbatch.metadata["response_length"] = response_length + return microbatch + + +def restore_microbatch_response_padding(tensor: torch.Tensor, metadata: Optional[dict]) -> torch.Tensor: + """Restore a microbatch output to the controller batch's response width.""" + if metadata is None or _PADDED_RESPONSE_LENGTH not in metadata: + return tensor + if tensor.ndim != 2: + raise ValueError(f"Expected a [batch, response] output, got shape {tuple(tensor.shape)}") + padded_length = metadata[_PADDED_RESPONSE_LENGTH] + if tensor.shape[-1] > padded_length: + raise ValueError(f"Cannot restore response width {tensor.shape[-1]} to smaller width {padded_length}") + if tensor.shape[-1] == padded_length: + return tensor + return torch.nn.functional.pad(tensor, (padded_length - tensor.shape[-1], 0)) + + class BaseBatchIterator: """Base class for batch iterators that chunk a TrainingInputBatch into microbatches.""" - def __init__(self, data: TrainingInputBatch): + def __init__(self, data: TrainingInputBatch, trim_padding: bool = False): self.data = data + self.trim_padding = trim_padding + + def _project(self, batch: TrainingInputBatch) -> TrainingInputBatch: + return trim_microbatch_padding(batch) if self.trim_padding else batch def __len__(self): raise NotImplementedError @@ -186,8 +286,14 @@ class BatchIterator(BaseBatchIterator): This is the original sample-based iterator. Kept as an alias for SampleBasedBatchIterator. """ - def __init__(self, data: TrainingInputBatch, sample_batch_size: int, drop_last: bool = False): - super().__init__(data) + def __init__( + self, + data: TrainingInputBatch, + sample_batch_size: int, + drop_last: bool = False, + trim_padding: bool = False, + ): + super().__init__(data, trim_padding=trim_padding) self.sample_batch_size = sample_batch_size self.total_batch_size = data.batch_size self.drop_last = drop_last @@ -206,7 +312,7 @@ def __iter__(self): def __next__(self) -> Experience: try: - batch = next(self._iter) + batch = self._project(next(self._iter)) exp = self.batch_to_experience(batch) return exp except StopIteration: @@ -224,8 +330,14 @@ class SampleBasedBatchIterator(BaseBatchIterator): Yields TrainingInputBatch objects (not Experience), unlike the legacy BatchIterator. """ - def __init__(self, data: TrainingInputBatch, sample_batch_size: int, drop_last: bool = False): - super().__init__(data) + def __init__( + self, + data: TrainingInputBatch, + sample_batch_size: int, + drop_last: bool = False, + trim_padding: bool = False, + ): + super().__init__(data, trim_padding=trim_padding) self.sample_batch_size = sample_batch_size self.total_batch_size = data.batch_size self.drop_last = drop_last @@ -238,7 +350,7 @@ def __len__(self): return self.num_micro_batches def __iter__(self) -> Iterator[TrainingInputBatch]: - return iter(self._chunks) + return (self._project(chunk) for chunk in self._chunks) def reorder_and_combine_batches(self, batches: List[TensorBatch]) -> TensorBatch: """Concatenate output batches. No reordering needed for sample-based splitting.""" @@ -257,13 +369,14 @@ def __init__( self, data: TrainingInputBatch, max_tokens_per_microbatch: int, + trim_padding: bool = False, ): """ Args: data: The training input batch to chunk. max_tokens_per_microbatch: Maximum number of tokens per microbatch. """ - super().__init__(data) + super().__init__(data, trim_padding=trim_padding) self._max_tokens_per_microbatch = max_tokens_per_microbatch # Compute token counts per sample using attention_mask @@ -291,7 +404,7 @@ def _create_microbatch_from_indices(self, indices: List[int]) -> TrainingInputBa selected_data[key] = value[indices_tensor] microbatch = TrainingInputBatch(selected_data) microbatch.metadata = self.data.metadata - return microbatch + return self._project(microbatch) def _create_padding_microbatch(self) -> TrainingInputBatch: """Create a padding microbatch with loss_mask=0 so it doesn't affect the loss.""" @@ -426,7 +539,10 @@ def reorder_and_combine_batches(self, batches: List[TensorBatch]) -> TensorBatch def get_microbatch_iterator( - data: TrainingInputBatch, micro_batch_size: int, max_tokens_per_microbatch: int + data: TrainingInputBatch, + micro_batch_size: int, + max_tokens_per_microbatch: int, + trim_padding: bool = False, ) -> BaseBatchIterator: """Factory function to get the appropriate microbatch iterator. @@ -434,11 +550,21 @@ def get_microbatch_iterator( data: The training input batch. micro_batch_size: Number of samples per microbatch (used if max_tokens_per_microbatch <= 0). max_tokens_per_microbatch: Maximum tokens per microbatch. If > 0, uses token-based batching. + trim_padding: Whether to project each microbatch to its local sequence and response widths. Returns: A BaseBatchIterator instance. """ if max_tokens_per_microbatch > 0: - return TokenBasedBatchIterator(data, max_tokens_per_microbatch=max_tokens_per_microbatch) + return TokenBasedBatchIterator( + data, + max_tokens_per_microbatch=max_tokens_per_microbatch, + trim_padding=trim_padding, + ) else: - return SampleBasedBatchIterator(data, sample_batch_size=micro_batch_size, drop_last=False) + return SampleBasedBatchIterator( + data, + sample_batch_size=micro_batch_size, + drop_last=False, + trim_padding=trim_padding, + ) diff --git a/skyrl/train/trainer.py b/skyrl/train/trainer.py index ae5eec4dfc..b8abf9570d 100644 --- a/skyrl/train/trainer.py +++ b/skyrl/train/trainer.py @@ -1323,7 +1323,7 @@ def fwd_logprobs_values_reward( - `["action_log_probs"]`: Float[torch.Tensor, "batch_size response_len"] - `["values"]`: Float[torch.Tensor, "batch_size response_len"] """ - fwd_keys = ["sequences", "attention_mask"] + fwd_keys = ["sequences", "attention_mask", "response_mask"] if training_input.get("rollout_expert_indices") is not None: fwd_keys.append("rollout_expert_indices") if training_input.get("router_padding_mask") is not None: diff --git a/tests/backends/skyrl_train/test_token_based_batching_utils.py b/tests/backends/skyrl_train/test_token_based_batching_utils.py index 4fac8a6679..d02f3e5d01 100644 --- a/tests/backends/skyrl_train/test_token_based_batching_utils.py +++ b/tests/backends/skyrl_train/test_token_based_batching_utils.py @@ -9,12 +9,16 @@ from typing import List +import pytest import torch from skyrl.backends.skyrl_train.training_batch import TensorList, TrainingInputBatch from skyrl.backends.skyrl_train.workers.worker_utils import ( + SampleBasedBatchIterator, TokenBasedBatchIterator, get_microbatch_iterator, + restore_microbatch_response_padding, + trim_microbatch_padding, ) from skyrl.train.dataset.bin_packing import make_seq_packer @@ -168,10 +172,6 @@ def test_get_microbatch_iterator_factory(self): assert isinstance(it, TokenBasedBatchIterator) # Sample-based (disabled) - from skyrl.backends.skyrl_train.workers.worker_utils import ( - SampleBasedBatchIterator, - ) - it = get_microbatch_iterator(batch, micro_batch_size=2, max_tokens_per_microbatch=-1) assert isinstance(it, SampleBasedBatchIterator) @@ -232,3 +232,96 @@ def test_multimodal_tensorlist_microbatching(self): assert len(pv) == microbatch["sequences"].shape[0] total_pv += len(pv) assert total_pv == batch_size # every sample's pixel_values is accounted for + + def test_trim_padding_uses_final_token_microbatch_widths(self): + batch = self._make_left_padded_batch() + iterator = TokenBasedBatchIterator(batch, max_tokens_per_microbatch=6, trim_padding=True) + + microbatches = list(iterator) + short = next(mb for mb in microbatches if mb["attention_mask"].sum() == 2) + + assert short["sequences"].shape == (1, 2) + assert short["sequences"].tolist() == [[11, 12]] + assert short["response_mask"].shape == (1, 1) + assert short["action_log_probs"].shape == (1, 1) + assert short["rollout_expert_indices"].shape == (1, 2, 1, 1) + assert short["router_padding_mask"].shape == (1, 2) + assert short.metadata["response_length"] == 1 + assert short.metadata["response_length_before_microbatch_trim"] == 3 + + # The controller batch remains the canonical, full-width representation. + assert batch["sequences"].shape == (2, 6) + assert batch["response_mask"].shape == (2, 3) + assert batch.metadata == {"response_length": 3} + + def test_trim_padding_applies_to_sample_microbatches(self): + batch = self._make_left_padded_batch() + iterator = SampleBasedBatchIterator(batch, sample_batch_size=1, trim_padding=True) + + short = next(iter(iterator)) + + assert short["sequences"].shape == (1, 2) + assert short["response_mask"].shape == (1, 1) + assert short.metadata["response_length"] == 1 + + def test_restore_response_padding_preserves_output_contract(self): + batch = self._make_left_padded_batch() + iterator = SampleBasedBatchIterator(batch, sample_batch_size=1, trim_padding=True) + short = next(iter(iterator)) + + local_output = torch.tensor([[0.25]]) + restored = restore_microbatch_response_padding(local_output, short.metadata) + + assert restored.tolist() == [[0.0, 0.0, 0.25]] + + def test_trim_padding_empty_batch_is_noop(self): + batch = TrainingInputBatch( + { + "sequences": torch.empty((0, 6), dtype=torch.long), + "attention_mask": torch.empty((0, 6), dtype=torch.long), + "response_mask": torch.empty((0, 3), dtype=torch.long), + } + ) + + assert trim_microbatch_padding(batch) is batch + + def test_restore_response_padding_rejects_non_matrix_output(self): + with pytest.raises(ValueError, match=r"Expected a \[batch, response\] output"): + restore_microbatch_response_padding( + torch.zeros((1, 2, 8)), + {"response_length_before_microbatch_trim": 3}, + ) + + @staticmethod + def _make_left_padded_batch(): + sequences = torch.tensor( + [ + [0, 0, 0, 0, 11, 12], + [21, 22, 23, 24, 25, 26], + ] + ) + attention_mask = torch.tensor( + [ + [0, 0, 0, 0, 1, 1], + [1, 1, 1, 1, 1, 1], + ] + ) + response_mask = torch.tensor( + [ + [0, 0, 1], + [1, 1, 1], + ] + ) + batch = TrainingInputBatch( + { + "sequences": sequences, + "attention_mask": attention_mask, + "response_mask": response_mask, + "loss_mask": response_mask.clone(), + "action_log_probs": torch.tensor([[0.0, 0.0, 0.1], [0.2, 0.3, 0.4]]), + "rollout_expert_indices": torch.arange(12).reshape(2, 6, 1, 1), + "router_padding_mask": ~attention_mask.bool(), + } + ) + batch.metadata = {"response_length": 3} + return batch diff --git a/tests/backends/skyrl_train/workers/test_policy_worker_loss_scaling.py b/tests/backends/skyrl_train/workers/test_policy_worker_loss_scaling.py index 6f3f847c67..4e9880b825 100644 --- a/tests/backends/skyrl_train/workers/test_policy_worker_loss_scaling.py +++ b/tests/backends/skyrl_train/workers/test_policy_worker_loss_scaling.py @@ -103,6 +103,42 @@ def _all_reduce_payload(strategy, op, key): raise AssertionError(f"No all_reduce call found for op={op!r}, key={key!r}") +def test_policy_forward_trims_dense_fsdp_microbatches_and_restores_output_width(): + cfg = _make_loss_scaling_cfg("dual_clip", micro_batch_size=1) + cfg.trainer.micro_forward_batch_size_per_gpu = 1 + cfg.trainer.max_tokens_per_microbatch = -1 + + batch = TrainingInputBatch( + { + "sequences": torch.tensor([[0, 0, 0, 0, 11, 12], [21, 22, 23, 24, 25, 26]]), + "attention_mask": torch.tensor([[0, 0, 0, 0, 1, 1], [1, 1, 1, 1, 1, 1]]), + "response_mask": torch.tensor([[0, 0, 1], [1, 1, 1]]), + } + ) + batch.metadata = {"response_length": 3} + + model_calls = [] + + def model_forward(sequences, num_actions, attention_mask, **kwargs): + model_calls.append((sequences.clone(), num_actions, attention_mask.clone())) + return torch.ones((sequences.shape[0], num_actions)) + + model = MagicMock(side_effect=model_forward) + worker = _make_policy_worker(cfg, model=model) + + with _patch_worker_cuda_for_cpu(): + result = worker.forward(batch) + + assert [(call[0].shape, call[1]) for call in model_calls] == [ + (torch.Size([1, 2]), 1), + (torch.Size([1, 6]), 3), + ] + assert result.loss_fn_outputs == [ + {"logprobs": [0.0, 0.0, 1.0]}, + {"logprobs": [1.0, 1.0, 1.0]}, + ] + + @pytest.mark.parametrize("loss_fn", ["cross_entropy", "dual_clip"]) def test_policy_forward_backward_loss_scaling_with_mocked_ranks(loss_fn): """FSDP forward_backward scales local microbatch losses before DP reduction."""