diff --git a/skyrl/backends/skyrl_train/workers/megatron/megatron_model_wrapper.py b/skyrl/backends/skyrl_train/workers/megatron/megatron_model_wrapper.py index 8672ec55b9..473442e335 100644 --- a/skyrl/backends/skyrl_train/workers/megatron/megatron_model_wrapper.py +++ b/skyrl/backends/skyrl_train/workers/megatron/megatron_model_wrapper.py @@ -40,6 +40,7 @@ from skyrl.backends.skyrl_train.utils.torch_utils import masked_mean from skyrl.backends.skyrl_train.workers.worker_utils import ( compute_minibatch_rollout_logprob_diff_metrics, + pop_return_per_token_outputs, ) from skyrl.train.config import TrainerConfig @@ -422,6 +423,8 @@ def forward_backward_mini_batch( loss_fn: Optional loss function name (e.g., "cross_entropy", "ppo"). If provided, overrides the config's policy_loss_type. loss_fn_config: Optional config overrides for the loss function. + May include reserved key ``return_per_token_outputs`` to skip + per-token ``loss_fn_outputs`` when callers read only ``metrics``. forward_only: If True, run the forward pass without backward (no gradients). Useful for evaluation / loss-only inference paths (e.g., SFT ``forward(loss_fn=...)`` codepath). @@ -440,10 +443,12 @@ def forward_backward_mini_batch( else: current_loss_fn = self.policy_loss_fn + # Consume the reserved gate before merging AlgorithmConfig overrides. + loss_fn_config, return_per_token_outputs = pop_return_per_token_outputs(loss_fn_config) + # Build config for loss function, applying any overrides loss_config = self.cfg.algorithm - if loss_fn_config is not None: - + if loss_fn_config: new_loss_config = OmegaConf.merge(OmegaConf.create(asdict(loss_config)), OmegaConf.create(loss_fn_config)) # NOTE: users can provide a custom loss config class, so we need to use the same class after applying overrides loss_config = type(loss_config).from_dict_config(new_loss_config) @@ -564,41 +569,43 @@ def loss_func(logits, data): if resolved_loss_name == "cross_entropy": loss = policy_loss - # Compute elementwise loss for Tinker API (per-token NLL) - with torch.no_grad(): - elementwise_loss = -action_log_probs - if loss_mask is not None: - elementwise_loss = elementwise_loss * loss_mask - - # Build per-sequence loss_fn_outputs. - # Compute valid_lens vectorized on GPU, then move tensors to CPU - # exactly once before iterating in Python — avoids ~3N GPU->CPU - # syncs per micro-batch (item()/cpu()/tolist() inside the loop). - batch_size = action_log_probs.shape[0] - seq_len = action_log_probs.shape[1] - if action_mask is not None: - valid_lens_t = action_mask.sum(dim=-1).long() - elif loss_mask is not None: - valid_lens_t = loss_mask.sum(dim=-1).long() + # Only build per-token outputs for callers that consume them. + if return_per_token_outputs: + # Tinker consumes per-token NLL. + with torch.no_grad(): + elementwise_loss = -action_log_probs + if loss_mask is not None: + elementwise_loss = elementwise_loss * loss_mask + + # Trim each sample with one CPU transfer per tensor. + batch_size = action_log_probs.shape[0] + seq_len = action_log_probs.shape[1] + if action_mask is not None: + valid_lens_t = action_mask.sum(dim=-1).long() + elif loss_mask is not None: + valid_lens_t = loss_mask.sum(dim=-1).long() + else: + valid_lens_t = torch.full( + (batch_size,), seq_len, device=action_log_probs.device, dtype=torch.long + ) + + action_log_probs_cpu = action_log_probs.detach().cpu() + elementwise_loss_cpu = elementwise_loss.detach().cpu() + valid_lens = valid_lens_t.cpu().tolist() + + loss_fn_outputs = [] + for i in range(batch_size): + valid_len = valid_lens[i] + loss_fn_outputs.append( + { + "logprobs": (action_log_probs_cpu[i, -valid_len:].tolist() if valid_len > 0 else []), + "elementwise_loss": ( + elementwise_loss_cpu[i, -valid_len:].tolist() if valid_len > 0 else [] + ), + } + ) else: - valid_lens_t = torch.full((batch_size,), seq_len, device=action_log_probs.device, dtype=torch.long) - - # Bulk GPU->CPU sync: one transfer for logprobs, elementwise_loss, and valid_lens. - action_log_probs_cpu = action_log_probs.detach().cpu() - elementwise_loss_cpu = elementwise_loss.detach().cpu() - valid_lens = valid_lens_t.cpu().tolist() - - loss_fn_outputs = [] - for i in range(batch_size): - valid_len = valid_lens[i] - loss_fn_outputs.append( - { - "logprobs": (action_log_probs_cpu[i, -valid_len:].tolist() if valid_len > 0 else []), - "elementwise_loss": ( - elementwise_loss_cpu[i, -valid_len:].tolist() if valid_len > 0 else [] - ), - } - ) + loss_fn_outputs = [{} for _ in range(action_log_probs.shape[0])] metrics = { "loss": loss.item(), diff --git a/skyrl/backends/skyrl_train/workers/worker.py b/skyrl/backends/skyrl_train/workers/worker.py index 76deba8a5a..b15466ea3a 100644 --- a/skyrl/backends/skyrl_train/workers/worker.py +++ b/skyrl/backends/skyrl_train/workers/worker.py @@ -53,6 +53,7 @@ all_reduce_metrics, compute_minibatch_rollout_logprob_diff_metrics, get_microbatch_iterator, + pop_return_per_token_outputs, reduce_metrics, ) from skyrl.env_vars import ( @@ -838,7 +839,9 @@ def _forward_backward_micro( loss_fn: Optional train loss function name to use instead of config default. Public Tinker aliases such as ``ppo`` should be normalized by the backend before reaching the worker. - loss_fn_config: Optional config overrides for the resolved train loss function + loss_fn_config: Optional config overrides for the resolved train loss function. + May include reserved key ``return_per_token_outputs`` to skip + per-token ``loss_fn_outputs`` when callers read only ``metrics``. Returns: Metrics dict for the worker's local micro batch @@ -868,9 +871,12 @@ def _forward_backward_micro( # Fall back to config default current_loss_fn = self.policy_loss_fn + # Consume the reserved gate before merging AlgorithmConfig overrides. + loss_fn_config, return_per_token_outputs = pop_return_per_token_outputs(loss_fn_config) + # Build config for loss function, applying any overrides loss_config = self.cfg.algorithm - if loss_fn_config is not None: + if loss_fn_config: # Create a copy of the config and apply overrides # TODO: Fix nested overrides from dataclasses import asdict @@ -910,40 +916,41 @@ def _forward_backward_micro( loss = unscaled_loss * microbatch_weight self.strategy.backward(loss, self.model, self.optimizer) - # Compute elementwise loss for Tinker API (per-token NLL) - with torch.no_grad(): - elementwise_loss = -action_log_probs - if loss_mask is not None: - elementwise_loss = elementwise_loss * loss_mask - - # Build per-sequence loss_fn_outputs (matches Tinker's ForwardBackwardOutput structure) - # Trim to actual response length per sample (Tinker expects variable-length arrays - # that align with the input weights, not padded to batch max). - # Compute valid_lens vectorized on GPU, then move tensors to CPU exactly - # once before iterating in Python — avoids ~3N GPU->CPU syncs per micro-batch. - batch_size = action_log_probs.shape[0] - seq_len = action_log_probs.shape[1] - if action_mask is not None: - valid_lens_t = action_mask.sum(dim=-1).long() - elif loss_mask is not None: - valid_lens_t = loss_mask.sum(dim=-1).long() + # Only build per-token outputs for callers that consume them. + if return_per_token_outputs: + # Tinker consumes per-token NLL. + with torch.no_grad(): + elementwise_loss = -action_log_probs + if loss_mask is not None: + elementwise_loss = elementwise_loss * loss_mask + + # Trim each sample with one CPU transfer per tensor. + batch_size = action_log_probs.shape[0] + seq_len = action_log_probs.shape[1] + if action_mask is not None: + valid_lens_t = action_mask.sum(dim=-1).long() + elif loss_mask is not None: + valid_lens_t = loss_mask.sum(dim=-1).long() + else: + valid_lens_t = torch.full((batch_size,), seq_len, device=action_log_probs.device, dtype=torch.long) + + action_log_probs_cpu = action_log_probs.detach().cpu() + elementwise_loss_cpu = elementwise_loss.detach().cpu() + valid_lens = valid_lens_t.cpu().tolist() + + loss_fn_outputs = [] + for i in range(batch_size): + valid_len = valid_lens[i] + loss_fn_outputs.append( + { + "logprobs": action_log_probs_cpu[i, -valid_len:].tolist() if valid_len > 0 else [], + "elementwise_loss": ( + elementwise_loss_cpu[i, -valid_len:].tolist() if valid_len > 0 else [] + ), + } + ) else: - valid_lens_t = torch.full((batch_size,), seq_len, device=action_log_probs.device, dtype=torch.long) - - # Bulk GPU->CPU sync: one transfer for logprobs, elementwise_loss, and valid_lens. - action_log_probs_cpu = action_log_probs.detach().cpu() - elementwise_loss_cpu = elementwise_loss.detach().cpu() - valid_lens = valid_lens_t.cpu().tolist() - - loss_fn_outputs = [] - for i in range(batch_size): - valid_len = valid_lens[i] - loss_fn_outputs.append( - { - "logprobs": action_log_probs_cpu[i, -valid_len:].tolist() if valid_len > 0 else [], - "elementwise_loss": (elementwise_loss_cpu[i, -valid_len:].tolist() if valid_len > 0 else []), - } - ) + loss_fn_outputs = [{} for _ in range(action_log_probs.shape[0])] status = { "loss": loss.item(), @@ -1109,6 +1116,13 @@ def _forward_micro_with_loss( Runs the model + loss under ``torch.no_grad()`` (no backward, no KL/entropy terms), and returns the same metrics shape as the SFT branch of ``_forward_backward_micro``, minus ``lr`` (no optimizer state involved). + + Args: + experience: Experience object for one micro batch. + loss_fn: Eval loss function name (e.g., "cross_entropy"). + loss_fn_config: Optional config overrides for the resolved loss function. + May include reserved key ``return_per_token_outputs`` to skip + per-token ``loss_fn_outputs`` when callers read only ``metrics``. """ self.model.eval() experience.to_device(torch.cuda.current_device()) @@ -1124,9 +1138,12 @@ def _forward_micro_with_loss( current_loss_fn = PolicyLossRegistry.get(loss_fn) + # Consume the reserved gate before merging AlgorithmConfig overrides. + loss_fn_config, return_per_token_outputs = pop_return_per_token_outputs(loss_fn_config) + # Build config for loss function, applying any overrides loss_config = self.cfg.algorithm - if loss_fn_config is not None: + if loss_fn_config: from dataclasses import asdict new_loss_config = OmegaConf.merge(OmegaConf.create(asdict(loss_config)), OmegaConf.create(loss_fn_config)) @@ -1153,36 +1170,39 @@ def _forward_micro_with_loss( rollout_logprobs=rollout_action_logprobs, ) - elementwise_loss = -action_log_probs - if loss_mask is not None: - elementwise_loss = elementwise_loss * loss_mask + # Only build per-token outputs for callers that consume them. + if return_per_token_outputs: + elementwise_loss = -action_log_probs + if loss_mask is not None: + elementwise_loss = elementwise_loss * loss_mask - # Compute valid_lens vectorized on GPU, then move tensors to CPU - # exactly once before iterating in Python. Avoids ~3N GPU->CPU syncs - # per micro-batch (item()/cpu()/tolist() inside the per-sample loop). - batch_size = action_log_probs.shape[0] - seq_len = action_log_probs.shape[1] - if action_mask is not None: - valid_lens_t = action_mask.sum(dim=-1).long() - elif loss_mask is not None: - valid_lens_t = loss_mask.sum(dim=-1).long() + # Trim each sample with one CPU transfer per tensor. + batch_size = action_log_probs.shape[0] + seq_len = action_log_probs.shape[1] + if action_mask is not None: + valid_lens_t = action_mask.sum(dim=-1).long() + elif loss_mask is not None: + valid_lens_t = loss_mask.sum(dim=-1).long() + else: + valid_lens_t = torch.full((batch_size,), seq_len, device=action_log_probs.device, dtype=torch.long) + + action_log_probs_cpu = action_log_probs.detach().cpu() + elementwise_loss_cpu = elementwise_loss.detach().cpu() + valid_lens = valid_lens_t.cpu().tolist() + + loss_fn_outputs = [] + for i in range(batch_size): + valid_len = valid_lens[i] + loss_fn_outputs.append( + { + "logprobs": action_log_probs_cpu[i, -valid_len:].tolist() if valid_len > 0 else [], + "elementwise_loss": ( + elementwise_loss_cpu[i, -valid_len:].tolist() if valid_len > 0 else [] + ), + } + ) else: - valid_lens_t = torch.full((batch_size,), seq_len, device=action_log_probs.device, dtype=torch.long) - - # Bulk GPU->CPU sync: one transfer for logprobs, elementwise_loss, and valid_lens. - action_log_probs_cpu = action_log_probs.detach().cpu() - elementwise_loss_cpu = elementwise_loss.detach().cpu() - valid_lens = valid_lens_t.cpu().tolist() - - loss_fn_outputs = [] - for i in range(batch_size): - valid_len = valid_lens[i] - loss_fn_outputs.append( - { - "logprobs": action_log_probs_cpu[i, -valid_len:].tolist() if valid_len > 0 else [], - "elementwise_loss": (elementwise_loss_cpu[i, -valid_len:].tolist() if valid_len > 0 else []), - } - ) + loss_fn_outputs = [{} for _ in range(action_log_probs.shape[0])] return { "loss": policy_loss.item(), diff --git a/skyrl/backends/skyrl_train/workers/worker_utils.py b/skyrl/backends/skyrl_train/workers/worker_utils.py index 31aa86cecc..69e515f68d 100644 --- a/skyrl/backends/skyrl_train/workers/worker_utils.py +++ b/skyrl/backends/skyrl_train/workers/worker_utils.py @@ -1,5 +1,5 @@ import math -from typing import Dict, Iterator, List, Optional +from typing import Any, Dict, Iterator, List, Optional, Tuple import torch import torch.distributed as dist @@ -49,6 +49,20 @@ def compute_minibatch_rollout_logprob_diff_metrics( } +# Reserved ``loss_fn_config`` key consumed before AlgorithmConfig validation. +RETURN_PER_TOKEN_OUTPUTS_KEY = "return_per_token_outputs" + + +def pop_return_per_token_outputs( + loss_fn_config: Optional[Dict[str, Any]], +) -> Tuple[Optional[Dict[str, Any]], bool]: + """Return a copied config and whether per-token outputs should be built.""" + if loss_fn_config is None: + return None, True + loss_fn_config = dict(loss_fn_config) + return loss_fn_config, loss_fn_config.pop(RETURN_PER_TOKEN_OUTPUTS_KEY, True) + + def reduce_metrics(metrics: Dict[str, List[float]], sum_loss_metrics: bool = False) -> Dict[str, float]: """Reduce scalar metrics from a list of entries per key with the appropriate reduction. diff --git a/skyrl/train/sft_trainer.py b/skyrl/train/sft_trainer.py index 08b955e54e..419bcb5a52 100644 --- a/skyrl/train/sft_trainer.py +++ b/skyrl/train/sft_trainer.py @@ -44,6 +44,7 @@ from skyrl.backends.skyrl_train.utils.io import io from skyrl.backends.skyrl_train.workers.worker import PPORayActorGroup from skyrl.backends.skyrl_train.workers.worker_dispatch import WorkerDispatch +from skyrl.backends.skyrl_train.workers.worker_utils import RETURN_PER_TOKEN_OUTPUTS_KEY from skyrl.env_vars import SKYRL_RAY_PG_TIMEOUT_IN_S from skyrl.train.config import SkyRLTrainConfig from skyrl.train.config.sft_config import ( @@ -1546,11 +1547,12 @@ def run_eval(self) -> tuple[dict, int]: # was 0/1 before scaling. Recover the count from the batch by counting positive entries. # Padded rows have loss_mask=0 so they are excluded here. nonpad_tokens = int((batch["loss_mask"] > 0).sum().item()) + # Eval consumes metrics only; skip per-token loss_fn_outputs. output = self.dispatch.forward( "policy", batch, loss_fn="cross_entropy", - loss_fn_config=None, + loss_fn_config={RETURN_PER_TOKEN_OUTPUTS_KEY: False}, ) batch_loss = float(output.metrics.get("loss", float("nan"))) total_loss_weighted += batch_loss * nonpad_tokens @@ -1571,7 +1573,13 @@ def train_step(self, batch: TrainingInputBatch, step: int) -> dict: """ timings: dict[str, float] = {} with Timer("forward_backward", timings): - output = self.dispatch.forward_backward("policy", batch, loss_fn="cross_entropy") + # SFT consumes metrics only; skip per-token loss_fn_outputs. + output = self.dispatch.forward_backward( + "policy", + batch, + loss_fn="cross_entropy", + loss_fn_config={RETURN_PER_TOKEN_OUTPUTS_KEY: False}, + ) with Timer("optim_step", timings): grad_norm = self.dispatch.optim_step("policy") diff --git a/tests/backends/skyrl_train/gpu/gpu_ci/test_training_step.py b/tests/backends/skyrl_train/gpu/gpu_ci/test_training_step.py index a1d9c22bcd..ed7a861d5c 100644 --- a/tests/backends/skyrl_train/gpu/gpu_ci/test_training_step.py +++ b/tests/backends/skyrl_train/gpu/gpu_ci/test_training_step.py @@ -9,6 +9,7 @@ import ray from skyrl.backends.skyrl_train.workers.worker_dispatch import WorkerDispatch +from skyrl.backends.skyrl_train.workers.worker_utils import RETURN_PER_TOKEN_OUTPUTS_KEY from skyrl.train.config import SkyRLTrainConfig from skyrl.train.utils.utils import print_mem, validate_cfg from tests.backends.skyrl_train.gpu.utils import ( @@ -333,3 +334,125 @@ async def test_sft_forward_with_cross_entropy(ray_init_fixture, cfg, strategy): for out in result.loss_fn_outputs: assert "logprobs" in out and len(out["logprobs"]) == num_actions assert "elementwise_loss" in out and len(out["elementwise_loss"]) == num_actions + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("strategy"), + ["fsdp", pytest.param("megatron", marks=pytest.mark.megatron)], + ids=["fsdp", "megatron"], +) +async def test_sft_forward_backward_return_per_token_outputs_gate(ray_init_fixture, cfg, strategy): + """The gate affects per-token outputs, not loss or consumed metrics.""" + cfg.trainer.remove_microbatch_padding = False + cfg.trainer.strategy = strategy + if strategy == "megatron": + cfg.trainer.policy.megatron_config.tensor_model_parallel_size = 1 + cfg.trainer.policy.megatron_config.pipeline_model_parallel_size = 1 + cfg.trainer.placement.policy_num_gpus_per_node = 2 + validate_cfg(cfg) + + try: + actor_group = init_worker_with_type( + "policy", + shared_pg=None, + colocate_all=False, + num_gpus_per_node=cfg.trainer.placement.policy_num_gpus_per_node, + cfg=cfg, + ) + + dp_size = actor_group.actor_infos[0].rank.dp_size + batch_size = dp_size * 2 + num_actions = 4 + + def _run(loss_fn_config): + batch = make_dummy_training_batch(batch_size=batch_size, num_actions=num_actions) + refs = actor_group.async_run_ray_method( + "mesh", + "forward_backward", + data=batch, + loss_fn="cross_entropy", + loss_fn_config=loss_fn_config, + ) + return ray.get(refs) + + kept_results = _run(None) + skipped_results = _run({RETURN_PER_TOKEN_OUTPUTS_KEY: False}) + + kept_outputs = [] + skipped_outputs = [] + for kept, skipped in zip(kept_results, skipped_results): + # Consumed metrics are unaffected by the gate. + assert kept.metrics["loss"] == skipped.metrics["loss"] + assert kept.metrics["response_length"] == skipped.metrics["response_length"] + assert kept.loss_fn_output_type == skipped.loss_fn_output_type == "scalar" + kept_outputs.extend(kept.loss_fn_outputs) + skipped_outputs.extend(skipped.loss_fn_outputs) + + assert len(kept_outputs) == batch_size + assert len(skipped_outputs) == batch_size + for out in kept_outputs: + assert len(out["logprobs"]) == num_actions + assert len(out["elementwise_loss"]) == num_actions + for out in skipped_outputs: + assert out == {} + + finally: + ray.shutdown() + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "strategy", + ["fsdp", pytest.param("megatron", marks=pytest.mark.megatron)], + ids=["fsdp", "megatron"], +) +async def test_sft_forward_return_per_token_outputs_gate(ray_init_fixture, cfg, strategy): + """Eval uses the same per-token-output gate.""" + cfg.trainer.remove_microbatch_padding = False + cfg.trainer.strategy = strategy + if strategy == "megatron": + cfg.trainer.policy.megatron_config.tensor_model_parallel_size = 1 + cfg.trainer.policy.megatron_config.pipeline_model_parallel_size = 1 + cfg.trainer.placement.policy_num_gpus_per_node = 2 + validate_cfg(cfg) + + try: + actor_group = init_worker_with_type( + "policy", + shared_pg=None, + colocate_all=False, + num_gpus_per_node=cfg.trainer.placement.policy_num_gpus_per_node, + cfg=cfg, + ) + + dp_size = actor_group.actor_infos[0].rank.dp_size + batch_size = dp_size * 2 + num_actions = 4 + dispatch = WorkerDispatch(cfg, policy_actor_group=actor_group) + + kept = dispatch.forward( + "policy", + make_dummy_training_batch(batch_size=batch_size, num_actions=num_actions), + loss_fn="cross_entropy", + loss_fn_config=None, + ) + skipped = dispatch.forward( + "policy", + make_dummy_training_batch(batch_size=batch_size, num_actions=num_actions), + loss_fn="cross_entropy", + loss_fn_config={RETURN_PER_TOKEN_OUTPUTS_KEY: False}, + ) + + assert kept.metrics["loss"] == skipped.metrics["loss"] + assert kept.loss_fn_output_type == skipped.loss_fn_output_type == "scalar" + assert len(kept.loss_fn_outputs) == batch_size + assert len(skipped.loss_fn_outputs) == batch_size + for out in kept.loss_fn_outputs: + assert len(out["logprobs"]) == num_actions + assert len(out["elementwise_loss"]) == num_actions + for out in skipped.loss_fn_outputs: + assert out == {} + + finally: + ray.shutdown() diff --git a/tests/backends/skyrl_train/workers/test_sft_loss_fn_outputs_gate.py b/tests/backends/skyrl_train/workers/test_sft_loss_fn_outputs_gate.py new file mode 100644 index 0000000000..04172dd964 --- /dev/null +++ b/tests/backends/skyrl_train/workers/test_sft_loss_fn_outputs_gate.py @@ -0,0 +1,226 @@ +"""CPU coverage for the per-request ``return_per_token_outputs`` gate. + +These tests drive worker loss-build methods with CPU mocks and pin default, +skip, config-pop, and RL-ungated behavior. GPU parity lives in +``tests/backends/skyrl_train/gpu/gpu_ci/test_training_step.py``. +""" + +from unittest.mock import MagicMock, patch + +import pytest +import torch + +from skyrl.backends.skyrl_train.utils.ppo_utils import PolicyLossRegistry +from skyrl.backends.skyrl_train.workers.worker import PolicyWorkerBase +from skyrl.backends.skyrl_train.workers.worker_utils import RETURN_PER_TOKEN_OUTPUTS_KEY +from skyrl.train.config import SkyRLTrainConfig +from skyrl.train.dataset.replay_buffer import Experience + +NUM_ACTIONS = 4 +BATCH_SIZE = 2 +SEQ_LEN = 6 + + +@pytest.fixture(scope="module", autouse=True) +def _repopulate_policy_loss_registry() -> None: + """Restore defaults after registry tests reset global state.""" + PolicyLossRegistry.repopulate_registry() + + +def _make_cpu_policy_worker() -> PolicyWorkerBase: + """Construct a CPU PolicyWorkerBase with mocked distributed deps.""" + cfg = SkyRLTrainConfig() + cfg.trainer.algorithm.policy_loss_type = "cross_entropy" + cfg.generator.sampling_params.temperature = 1.0 + cfg.trainer.algorithm.temperature = 1.0 + + worker = PolicyWorkerBase( + cfg=cfg.trainer, + world_size=1, + rank=0, + local_rank=0, + master_addr="localhost", + master_port=12345, + sequence_parallel_size=1, + ) + worker.strategy = MagicMock() + worker.scheduler = MagicMock() + worker.scheduler.get_last_lr.return_value = [1e-4] + worker.optimizer = MagicMock() + return worker + + +def _patch_model(worker: PolicyWorkerBase, action_log_probs: torch.Tensor) -> None: + """Make ``worker.model(...)`` return canned logprobs.""" + model = MagicMock() + model.return_value = (action_log_probs, {"entropy": torch.zeros_like(action_log_probs)}) + worker.model = model + + +def _make_experience() -> Experience: + """Small CPU Experience with valid_len == NUM_ACTIONS.""" + return Experience( + sequences=torch.randint(0, 100, (BATCH_SIZE, SEQ_LEN)), + action_log_probs=None, + base_action_log_probs=None, + rollout_logprobs=None, + values=None, + returns=None, + advantages=torch.zeros(BATCH_SIZE, NUM_ACTIONS), + attention_mask=torch.ones(BATCH_SIZE, SEQ_LEN, dtype=torch.long), + loss_mask=torch.ones(BATCH_SIZE, NUM_ACTIONS, dtype=torch.long), + action_mask=None, + rollout_expert_indices=None, + num_actions=NUM_ACTIONS, + info={}, + ) + + +def _run_forward_backward_micro(loss_fn_config): + """Drive the train cross_entropy build on CPU.""" + worker = _make_cpu_policy_worker() + action_log_probs = torch.full((BATCH_SIZE, NUM_ACTIONS), -0.5) + _patch_model(worker, action_log_probs) + experience = _make_experience() + with patch("torch.cuda.current_device", return_value="cpu"), patch("torch.autocast", MagicMock()): + return worker._forward_backward_micro( + experience, + microbatch_weight=1.0, + loss_fn="cross_entropy", + loss_fn_config=loss_fn_config, + ) + + +def _run_forward_micro_with_loss(loss_fn_config): + """Drive the eval cross_entropy build on CPU.""" + worker = _make_cpu_policy_worker() + action_log_probs = torch.full((BATCH_SIZE, NUM_ACTIONS), -0.5) + _patch_model(worker, action_log_probs) + experience = _make_experience() + with patch("torch.cuda.current_device", return_value="cpu"), patch("torch.autocast", MagicMock()): + return worker._forward_micro_with_loss( + experience, + loss_fn="cross_entropy", + loss_fn_config=loss_fn_config, + ) + + +def _run_forward_backward_micro_rl(loss_fn_config): + """Drive the PPO path, which must ignore the SFT-only gate.""" + worker = _make_cpu_policy_worker() + worker.cfg.algorithm.policy_loss_type = "regular" + # Disable extra PPO terms so this test isolates loss_fn_outputs. + worker.cfg.algorithm.use_kl_loss = False + worker.cfg.algorithm.use_entropy_loss = False + worker.mesh_rank = MagicMock() + worker.mesh_rank.dp_size = 1 + + action_log_probs = torch.full((BATCH_SIZE, NUM_ACTIONS), -0.5) + model = MagicMock() + # The RL branch slices entropy from the full sequence. + model.return_value = (action_log_probs, {"entropy": torch.zeros(BATCH_SIZE, SEQ_LEN)}) + worker.model = model + + experience = _make_experience() + experience.action_log_probs = torch.full((BATCH_SIZE, NUM_ACTIONS), -0.5) + experience.advantages = torch.ones(BATCH_SIZE, NUM_ACTIONS) + with patch("torch.cuda.current_device", return_value="cpu"), patch("torch.autocast", MagicMock()): + return worker._forward_backward_micro( + experience, + microbatch_weight=1.0, + loss_fn="regular", + loss_fn_config=loss_fn_config, + ) + + +class TestForwardBackwardMicroGate: + def test_default_keeps_per_token_outputs(self): + """Default: every sequence carries logprobs + NLL.""" + status = _run_forward_backward_micro(loss_fn_config=None) + outputs = status["loss_fn_outputs"] + assert len(outputs) == BATCH_SIZE + for out in outputs: + assert len(out["logprobs"]) == NUM_ACTIONS + assert len(out["elementwise_loss"]) == NUM_ACTIONS + + def test_explicit_true_keeps_per_token_outputs(self): + status = _run_forward_backward_micro(loss_fn_config={RETURN_PER_TOKEN_OUTPUTS_KEY: True}) + outputs = status["loss_fn_outputs"] + assert len(outputs) == BATCH_SIZE + for out in outputs: + assert len(out["logprobs"]) == NUM_ACTIONS + assert len(out["elementwise_loss"]) == NUM_ACTIONS + + def test_false_skips_per_token_outputs(self): + """False: one empty dict per sequence.""" + status = _run_forward_backward_micro(loss_fn_config={RETURN_PER_TOKEN_OUTPUTS_KEY: False}) + outputs = status["loss_fn_outputs"] + assert len(outputs) == BATCH_SIZE + for out in outputs: + assert out == {} + + def test_loss_and_metrics_identical_across_flag(self): + """Skipping per-token outputs must not perturb consumed metrics.""" + kept = _run_forward_backward_micro(loss_fn_config={RETURN_PER_TOKEN_OUTPUTS_KEY: True}) + skipped = _run_forward_backward_micro(loss_fn_config={RETURN_PER_TOKEN_OUTPUTS_KEY: False}) + assert kept["loss"] == skipped["loss"] + assert kept["response_length"] == skipped["response_length"] + assert kept["lr"] == skipped["lr"] + + def test_flag_popped_before_algorithm_config_merge(self): + """The reserved flag is removed before AlgorithmConfig validation.""" + # A real override key confirms legitimate config merge still runs. + status = _run_forward_backward_micro(loss_fn_config={RETURN_PER_TOKEN_OUTPUTS_KEY: False, "eps_clip_low": 0.1}) + assert status["loss_fn_outputs"] == [{} for _ in range(BATCH_SIZE)] + + def test_caller_loss_fn_config_not_mutated(self): + """Popping the flag must not mutate the caller-provided dict.""" + cfg_dict = {RETURN_PER_TOKEN_OUTPUTS_KEY: False} + _run_forward_backward_micro(loss_fn_config=cfg_dict) + assert cfg_dict == {RETURN_PER_TOKEN_OUTPUTS_KEY: False} + + +class TestForwardMicroWithLossGate: + def test_default_keeps_per_token_outputs(self): + status = _run_forward_micro_with_loss(loss_fn_config=None) + outputs = status["loss_fn_outputs"] + assert len(outputs) == BATCH_SIZE + for out in outputs: + assert len(out["logprobs"]) == NUM_ACTIONS + assert len(out["elementwise_loss"]) == NUM_ACTIONS + + def test_false_skips_per_token_outputs(self): + status = _run_forward_micro_with_loss(loss_fn_config={RETURN_PER_TOKEN_OUTPUTS_KEY: False}) + outputs = status["loss_fn_outputs"] + assert len(outputs) == BATCH_SIZE + for out in outputs: + assert out == {} + + def test_loss_identical_across_flag(self): + kept = _run_forward_micro_with_loss(loss_fn_config={RETURN_PER_TOKEN_OUTPUTS_KEY: True}) + skipped = _run_forward_micro_with_loss(loss_fn_config={RETURN_PER_TOKEN_OUTPUTS_KEY: False}) + assert kept["loss"] == skipped["loss"] + assert kept["response_length"] == skipped["response_length"] + + +class TestForwardBackwardMicroRLPathUngated: + def test_rl_builds_logprobs_with_flag_true(self): + status = _run_forward_backward_micro_rl(loss_fn_config={RETURN_PER_TOKEN_OUTPUTS_KEY: True}) + outputs = status["loss_fn_outputs"] + assert len(outputs) == BATCH_SIZE + for out in outputs: + assert len(out["logprobs"]) == NUM_ACTIONS + + def test_rl_builds_logprobs_even_when_flag_false(self): + """The SFT-only gate must not empty PPO outputs.""" + status = _run_forward_backward_micro_rl(loss_fn_config={RETURN_PER_TOKEN_OUTPUTS_KEY: False}) + outputs = status["loss_fn_outputs"] + assert len(outputs) == BATCH_SIZE + for out in outputs: + assert len(out["logprobs"]) == NUM_ACTIONS + + def test_rl_loss_fn_outputs_identical_across_flag(self): + kept = _run_forward_backward_micro_rl(loss_fn_config={RETURN_PER_TOKEN_OUTPUTS_KEY: True}) + skipped = _run_forward_backward_micro_rl(loss_fn_config={RETURN_PER_TOKEN_OUTPUTS_KEY: False}) + assert kept["loss_fn_outputs"] == skipped["loss_fn_outputs"] + assert kept["final_loss"] == skipped["final_loss"] diff --git a/tests/backends/skyrl_train/workers/test_worker_utils.py b/tests/backends/skyrl_train/workers/test_worker_utils.py index 2c01499faa..3b43c4e403 100644 --- a/tests/backends/skyrl_train/workers/test_worker_utils.py +++ b/tests/backends/skyrl_train/workers/test_worker_utils.py @@ -9,7 +9,9 @@ import pytest from skyrl.backends.skyrl_train.workers.worker_utils import ( + RETURN_PER_TOKEN_OUTPUTS_KEY, all_reduce_metrics, + pop_return_per_token_outputs, reduce_metrics, ) @@ -197,3 +199,34 @@ def mock_all_reduce(d, op, group=None): assert result["is_ratio_min"] == 0.05 # 0.1 / 2 (min op) assert result["policy_loss"] == 6.0 # sum op assert result["entropy"] == 1.0 # 0.5 * 2 (mean op) + + +class TestPopReturnPerTokenOutputs: + """Tests for the reserved loss_fn_config gate.""" + + def test_none_defaults_to_true(self): + cfg, flag = pop_return_per_token_outputs(None) + assert cfg is None + assert flag is True + + def test_absent_flag_defaults_to_true(self): + cfg, flag = pop_return_per_token_outputs({"eps_clip_low": 0.1}) + assert flag is True + assert cfg == {"eps_clip_low": 0.1} + + def test_pops_explicit_false(self): + cfg, flag = pop_return_per_token_outputs({RETURN_PER_TOKEN_OUTPUTS_KEY: False, "eps_clip_low": 0.1}) + assert flag is False + assert cfg == {"eps_clip_low": 0.1} + + def test_pops_explicit_true(self): + cfg, flag = pop_return_per_token_outputs({RETURN_PER_TOKEN_OUTPUTS_KEY: True}) + assert flag is True + assert cfg == {} + + def test_does_not_mutate_caller_dict(self): + original = {RETURN_PER_TOKEN_OUTPUTS_KEY: False} + cfg, flag = pop_return_per_token_outputs(original) + assert original == {RETURN_PER_TOKEN_OUTPUTS_KEY: False} + assert cfg == {} + assert flag is False diff --git a/tests/train/test_sft_callbacks.py b/tests/train/test_sft_callbacks.py index 1a7d1c2d6d..8532b8c728 100644 --- a/tests/train/test_sft_callbacks.py +++ b/tests/train/test_sft_callbacks.py @@ -20,6 +20,9 @@ from unittest.mock import MagicMock +import pytest + +from skyrl.backends.skyrl_train.workers.worker_utils import RETURN_PER_TOKEN_OUTPUTS_KEY from skyrl.train.config.sft_config import ( SFTConfig, SFTPlacementConfig, @@ -34,6 +37,35 @@ _FAKE_CKPT_PATH = "/fake/sft-callback-test/global_step_2" +def _make_mock_dispatch() -> MagicMock: + """Worker-dispatch mock for SFT orchestration tests.""" + step_output = MagicMock() + step_output.metrics = {"loss": 0.42, "final_loss": 0.42} + eval_output = MagicMock() + eval_output.metrics = {"loss": 0.31} + + dispatch_mock = MagicMock() + dispatch_mock.forward_backward = MagicMock(return_value=step_output) + dispatch_mock.optim_step = MagicMock(return_value=1.0) + dispatch_mock.forward = MagicMock(return_value=eval_output) + dispatch_mock.dp_size = MagicMock(return_value=1) + return dispatch_mock + + +def _attach_mock_sft_deps(trainer: SFTTrainer, dispatch_mock: MagicMock) -> None: + """Wire mocked setup outputs onto the trainer.""" + tokenizer = MagicMock() + tokenizer.pad_token_id = 0 + trainer.tokenizer = tokenizer + trainer.collator = trainer._build_collator(tokenizer) + trainer.dispatch = dispatch_mock + + +@pytest.fixture +def mock_dispatch() -> MagicMock: + return _make_mock_dispatch() + + class RecorderCallback(TrainingCallback): """Spy: records every event with a snapshot of the relevant CallbackInput fields.""" @@ -158,7 +190,7 @@ def _dummy_tokenized() -> list[dict]: return [example, example] -def test_callbacks_fire_during_sft_training(monkeypatch): +def test_callbacks_fire_during_sft_training(monkeypatch, mock_dispatch): """A 2-step SFT run fires every relevant event, in order, with the right payloads.""" cfg = _build_test_sft_config() skyrl_cfg = build_skyrl_config_for_sft(cfg) @@ -168,30 +200,10 @@ def test_callbacks_fire_during_sft_training(monkeypatch): force_save = ForceSaveAtStep(step=2) trainer = SFTTrainer(cfg, skyrl_cfg=skyrl_cfg, callbacks=[recorder, force_eval, force_save]) - # Skip setup() (which would load the model + spin up workers). Replace - # what setup() would have set with mocks. - tokenizer = MagicMock() - tokenizer.pad_token_id = 0 - trainer.tokenizer = tokenizer - # setup() also builds the collator once the tokenizer is available. - trainer.collator = trainer._build_collator(tokenizer) + # Skip setup() by wiring the deps it normally creates. + _attach_mock_sft_deps(trainer, mock_dispatch) trainer.tracker = MagicMock() - # Mock the worker dispatch — the only thing train_step / run_eval touch - # that requires real GPU workers. forward_backward returns an object with - # ``.metrics`` (loss); optim_step returns a grad_norm; forward (eval path) - # returns ``.metrics`` with a per-batch loss. - step_output = MagicMock() - step_output.metrics = {"loss": 0.42, "final_loss": 0.42} - eval_output = MagicMock() - eval_output.metrics = {"loss": 0.31} - dispatch_mock = MagicMock() - dispatch_mock.forward_backward = MagicMock(return_value=step_output) - dispatch_mock.optim_step = MagicMock(return_value=1.0) - dispatch_mock.forward = MagicMock(return_value=eval_output) - dispatch_mock.dp_size = MagicMock(return_value=1) - trainer.dispatch = dispatch_mock - # Bypass HF network fetch + tokenization: both load_dataset() and # load_eval_dataset() funnel through _load_and_tokenize. monkeypatch.setattr(trainer, "_load_and_tokenize", lambda *_args, **_kw: _dummy_tokenized()) @@ -285,3 +297,39 @@ def _record_load_checkpoint(): continue assert snap["total_steps"] == 2, f"{name}: total_steps={snap['total_steps']}" assert snap["steps_per_epoch"] == 2, f"{name}: steps_per_epoch={snap['steps_per_epoch']}" + + +def _build_minimal_trainer(dispatch_mock: MagicMock) -> SFTTrainer: + """Build an SFTTrainer with mocked dispatch.""" + cfg = _build_test_sft_config() + skyrl_cfg = build_skyrl_config_for_sft(cfg) + trainer = SFTTrainer(cfg, skyrl_cfg=skyrl_cfg) + _attach_mock_sft_deps(trainer, dispatch_mock) + return trainer + + +def test_sft_train_step_opts_out_of_per_token_outputs(mock_dispatch): + """train_step opts out of unused per-token outputs.""" + trainer = _build_minimal_trainer(mock_dispatch) + batch = trainer.collator(_dummy_tokenized(), batch_size=1) + + trainer.train_step(batch, step=1) + + mock_dispatch.forward_backward.assert_called_once() + call = mock_dispatch.forward_backward.call_args + assert call.kwargs["loss_fn"] == "cross_entropy" + assert call.kwargs["loss_fn_config"] == {RETURN_PER_TOKEN_OUTPUTS_KEY: False} + + +def test_sft_run_eval_opts_out_of_per_token_outputs(mock_dispatch): + """run_eval reads only ``output.metrics["loss"]``; it skips per-token outputs.""" + trainer = _build_minimal_trainer(mock_dispatch) + trainer.eval_dataloader = trainer.build_eval_dataloader(_dummy_tokenized()) + + metrics, _ = trainer.run_eval() + + assert "eval_loss" in metrics + mock_dispatch.forward.assert_called() + for call in mock_dispatch.forward.call_args_list: + assert call.kwargs["loss_fn"] == "cross_entropy" + assert call.kwargs["loss_fn_config"] == {RETURN_PER_TOKEN_OUTPUTS_KEY: False}