diff --git a/skyrl/backends/skyrl_train/workers/megatron/megatron_worker.py b/skyrl/backends/skyrl_train/workers/megatron/megatron_worker.py index d116213abd..287ae313b0 100644 --- a/skyrl/backends/skyrl_train/workers/megatron/megatron_worker.py +++ b/skyrl/backends/skyrl_train/workers/megatron/megatron_worker.py @@ -1254,11 +1254,13 @@ def forward_backward( torch.cuda.empty_cache() # Aggregate metrics across micro-batches - all_loss_fn_outputs = [] # Handle separately from scalar metrics + loss_fn_output_batches = [] for m_batch, metrics in zip(micro_buffer, metrics_list): # Extract loss_fn_outputs before reduce_metrics (it's not a scalar metric) - if "loss_fn_outputs" in metrics: - all_loss_fn_outputs.extend(metrics.pop("loss_fn_outputs")) + if metrics is None: + loss_fn_output_batches.append([]) + continue + loss_fn_output_batches.append(metrics.pop("loss_fn_outputs", [])) # Skip fully-padding microbatches: their metrics (clip_ratio=0, policy_entropy=0, # ...) are meaningless and would drag down the mean-reduced metrics. Summed # metrics (e.g. policy_loss) are unaffected since padding contributes 0, but @@ -1303,6 +1305,13 @@ def forward_backward( for k, v in moe_metrics.items(): status[k] = v + if not any(loss_fn_output_batches): + all_loss_fn_outputs = [] + elif isinstance(microbatch_iterator, TokenBasedBatchIterator): + all_loss_fn_outputs = microbatch_iterator.reorder_and_combine_items(loss_fn_output_batches) + else: + all_loss_fn_outputs = [item for batch in loss_fn_output_batches for item in batch] + return WorkerOutput(loss_fn_outputs=all_loss_fn_outputs, metrics=status) def optim_step(self) -> Optional[float]: diff --git a/skyrl/backends/skyrl_train/workers/worker_utils.py b/skyrl/backends/skyrl_train/workers/worker_utils.py index 2efa37ad6c..71fb9518dc 100644 --- a/skyrl/backends/skyrl_train/workers/worker_utils.py +++ b/skyrl/backends/skyrl_train/workers/worker_utils.py @@ -424,6 +424,18 @@ def reorder_and_combine_batches(self, batches: List[TensorBatch]) -> TensorBatch reordered_batch.metadata = ref_microbatch.metadata return reordered_batch + def reorder_and_combine_items(self, batches: List[List[dict]]) -> List[dict]: + """Restore per-sample microbatch outputs to input order.""" + ordered = [None] * self.data.batch_size + for original_indices, items in zip(self._microbatches, batches): + if len(items) < len(original_indices): + raise ValueError("Microbatch output has fewer items than input samples") + for original_idx, item in zip(original_indices, items): + ordered[original_idx] = item + if any(item is None for item in ordered): + raise ValueError("Microbatch outputs do not cover every input sample") + return ordered + def get_microbatch_iterator( data: TrainingInputBatch, micro_batch_size: int, max_tokens_per_microbatch: int diff --git a/skyrl/backends/skyrl_train_backend.py b/skyrl/backends/skyrl_train_backend.py index ea08b13d01..2ad126dec1 100644 --- a/skyrl/backends/skyrl_train_backend.py +++ b/skyrl/backends/skyrl_train_backend.py @@ -808,6 +808,21 @@ def _normalize_policy_loss_request( if role == "critic": return loss_fn, loss_fn_config + if loss_fn == "gspo": + normalized_config = dict(loss_fn_config or {}) + clip_low_threshold = normalized_config.pop("clip_low_threshold", None) + clip_high_threshold = normalized_config.pop("clip_high_threshold", None) + if clip_low_threshold is None or clip_high_threshold is None: + raise ValueError("loss_fn='gspo' requires clip_low_threshold and clip_high_threshold") + normalized_config.update( + eps_clip_low=1.0 - clip_low_threshold, + eps_clip_high=clip_high_threshold - 1.0, + loss_reduction="sequence_mean", + use_entropy_loss=False, + use_kl_loss=False, + ) + return loss_fn, normalized_config + if loss_fn == "dppo": # DPPO thresholds live in the nested `algorithm.dppo` sub-config, but # Tinker's loss_fn_config is a flat float dict. Re-nest so the @@ -832,6 +847,47 @@ def _normalize_policy_loss_request( normalized_config["eps_clip_high"] = clip_high_threshold - 1.0 return ("regular" if loss_fn == "ppo" else "gspo"), normalized_config or None + @staticmethod + def _normalize_gspo_batch(batch: TrainingInputBatch) -> None: + """Apply Tinker's sequence-mean GSPO reduction before DP sharding.""" + advantages = batch["advantages"] + loss_mask = advantages.ne(0).to(batch["loss_mask"].dtype) + token_counts = loss_mask.sum(dim=-1, keepdim=True) + trainable_sequences = token_counts.squeeze(-1).gt(0) + num_trainable_sequences = trainable_sequences.sum().clamp(min=1) + batch["advantages"] = advantages / token_counts.clamp(min=1) / num_trainable_sequences + batch["loss_mask"] = loss_mask + + @staticmethod + def _select_trainable_gspo_rows(batch: TrainingInputBatch) -> tuple[TrainingInputBatch, list[int]]: + """Drop inactive rows while preserving all-zero-batch optimizer semantics.""" + trainable_rows = batch["loss_mask"].ne(0).any(dim=-1) + trainable_indices = trainable_rows.nonzero(as_tuple=False).flatten() + if trainable_indices.numel() == 0 or trainable_indices.numel() == batch.batch_size: + return batch, list(range(batch.batch_size)) + + selected = TrainingInputBatch( + {key: None if value is None else value[trainable_indices] for key, value in batch.items()} + ) + selected.metadata = batch.metadata + return selected, trainable_indices.tolist() + + @staticmethod + def _get_gspo_pruning_metrics( + original_token_counts: list[int], submitted_indices: set[int], start_idx: int, end_idx: int + ) -> dict[str, float]: + """Report one SDK chunk's pruning counts so combined futures sum to the batch total.""" + request_indices = range(start_idx, end_idx) + request_submitted_indices = [index for index in request_indices if index in submitted_indices] + submitted_tokens = sum(original_token_counts[index] for index in request_submitted_indices) + original_tokens = sum(original_token_counts[index] for index in request_indices) + return { + "skyrl.ai/submitted_datums:sum": float(len(request_submitted_indices)), + "skyrl.ai/skipped_datums:sum": float(end_idx - start_idx - len(request_submitted_indices)), + "skyrl.ai/submitted_tokens:sum": float(submitted_tokens), + "skyrl.ai/skipped_tokens:sum": float(original_tokens - submitted_tokens), + } + def forward_backward( self, prepared_batch: types.PreparedModelPassBatch, @@ -860,6 +916,12 @@ def _forward_backward_single_model_batch( ): raise ValueError("Critic forward_backward requires values and returns for every response token") batch = self._to_training_batch(prepared_batch, role) + original_batch_size = batch.batch_size + original_token_counts = [int(count) for count in batch["attention_mask"].sum(dim=-1).tolist()] + submitted_indices = list(range(original_batch_size)) + if loss_fn == "gspo": + self._normalize_gspo_batch(batch) + batch, submitted_indices = self._select_trainable_gspo_rows(batch) micro_bs = ( self._cfg.trainer.micro_train_batch_size_per_gpu if self._cfg.trainer.strategy == "megatron" else None ) @@ -897,7 +959,14 @@ def _forward_backward_single_model_batch( if pad_size > 0 and per_sample_outputs: per_sample_outputs = per_sample_outputs[:-pad_size] + if per_sample_outputs and len(submitted_indices) != original_batch_size: + restored_outputs = [{"logprobs": list(logprobs)} for logprobs in prepared_batch.all_sampling_logprobs] + for index, output in zip(submitted_indices, per_sample_outputs, strict=True): + restored_outputs[index] = output + per_sample_outputs = restored_outputs + metrics = self._extract_metrics(data.metrics) + submitted_index_set = set(submitted_indices) results = {} for request_id, _, start_idx, end_idx in prepared_batch.request_batch_slices: @@ -917,10 +986,15 @@ def _forward_backward_single_model_batch( loss_fn_outputs.append(formatted_output) else: loss_fn_outputs = [{} for _ in range(end_idx - start_idx)] + request_metrics = dict(metrics) + if loss_fn == "gspo": + request_metrics.update( + self._get_gspo_pruning_metrics(original_token_counts, submitted_index_set, start_idx, end_idx) + ) results[request_id] = types.ForwardBackwardOutput( loss_fn_output_type=data.loss_fn_output_type, loss_fn_outputs=loss_fn_outputs, - metrics=metrics, + metrics=request_metrics, ) return results diff --git a/skyrl/tinker/api.py b/skyrl/tinker/api.py index e63fcbeaa4..a1e581a9d5 100644 --- a/skyrl/tinker/api.py +++ b/skyrl/tinker/api.py @@ -7,6 +7,7 @@ import signal import threading from contextlib import asynccontextmanager, suppress +from dataclasses import dataclass from datetime import datetime, timedelta, timezone from typing import Annotated, Any, AsyncGenerator, ClassVar, Literal from uuid import uuid4 @@ -78,6 +79,13 @@ TERMINAL_STATUSES = (RequestStatus.COMPLETED, RequestStatus.FAILED) +@dataclass(frozen=True) +class SamplingTarget: + base_model: str | None + model_id: str + checkpoint_id: str + + def raw_json_response(payload: str | None) -> Response: """Return already-serialized JSON without routing it through FastAPI's encoder. @@ -260,6 +268,8 @@ async def lifespan(app: FastAPI): # SkyRL-Train default is colocate_all=True; only opt into forwarding # when the operator explicitly sets it to False. is_colocated = bool(backend_cfg.get("trainer.placement.colocate_all", True)) + app.state.external_sampling_targets = {} + app.state.external_sampling_target_locks = {} if app.state.engine_config.external_inference_url: app.state.external_inference_client = ExternalInferenceClient(app.state.engine_config, app.state.db_engine) logger.info(f"External engine configured: {app.state.engine_config.external_inference_url}") @@ -625,6 +635,8 @@ class ForwardBackwardInput(BaseModel): def validate_loss_fn_config_keys(self): """Validate loss_fn_config keys based on the selected loss function.""" if self.loss_fn_config is None: + if self.loss_fn == "gspo": + raise ValueError("loss_fn='gspo' requires clip_low_threshold and clip_high_threshold.") return self allowed_keys = self._ALLOWED_KEYS_BY_LOSS_FN[self.loss_fn] @@ -638,6 +650,18 @@ def validate_loss_fn_config_keys(self): raise ValueError( f"loss_fn='{self.loss_fn}' does not accept loss_fn_config keys. " f"Received: {invalid_keys}." ) + if self.loss_fn == "gspo": + required_keys = {"clip_low_threshold", "clip_high_threshold"} + missing_keys = sorted(required_keys - self.loss_fn_config.keys()) + if missing_keys: + raise ValueError(f"loss_fn='gspo' is missing required loss_fn_config keys: {missing_keys}.") + clip_low = self.loss_fn_config["clip_low_threshold"] + clip_high = self.loss_fn_config["clip_high_threshold"] + if not 0.0 <= clip_low <= 1.0 <= clip_high: + raise ValueError( + "loss_fn='gspo' requires 0 <= clip_low_threshold <= 1 <= clip_high_threshold; " + f"got {clip_low=} and {clip_high=}." + ) return self def to_types(self) -> types.ForwardBackwardInput: @@ -771,12 +795,14 @@ def validate_model_source(self): class SaveWeightsRequest(BaseModel): model_id: str path: str = Field(..., pattern=ID_PATTERN, max_length=ID_MAX_LENGTH) + seq_id: int | None = None type: Literal["save_weights"] | None = None class LoadWeightsRequest(BaseModel): model_id: str path: str + seq_id: int | None = None type: Literal["load_weights"] | None = None @@ -859,6 +885,7 @@ class SupportedModel(BaseModel): class GetServerCapabilitiesResponse(BaseModel): supported_models: list[SupportedModel] + supported_loss_fns: list[str] class ListCheckpointsResponse(BaseModel): @@ -1185,6 +1212,7 @@ async def load_weights(request: LoadWeightsRequest, req: Request, session: Async request_type=types.RequestType.LOAD_WEIGHTS, model_id=request.model_id, request_data=types.LoadWeightsInput(source_model_id=source_model_id, checkpoint_id=checkpoint_id), + seq_id=request.seq_id, ) await session.commit() @@ -1208,6 +1236,7 @@ async def save_weights(request: SaveWeightsRequest, session: AsyncSession = Depe request_type=types.RequestType.SAVE_WEIGHTS, model_id=request.model_id, request_data=types.SaveWeightsInput(path=request.path), + seq_id=request.seq_id, ) await session.commit() @@ -1253,6 +1282,7 @@ async def save_weights_for_sampler(request: SaveWeightsForSamplerRequest, sessio seq_id=request.seq_id, sampling_session_id=sampling_session_id, ), + seq_id=request.seq_id, ) await session.commit() @@ -1260,7 +1290,7 @@ async def save_weights_for_sampler(request: SaveWeightsForSamplerRequest, sessio return FutureResponse(future_id=str(request_id), status="pending", request_id=str(request_id)) -async def get_sampling_model(request: SampleRequest, session: AsyncSession) -> (str | None, str | None): +async def get_sampling_model(request: SampleRequest, session: AsyncSession) -> tuple[str | None, str | None]: """Return (base_model, model_path) for a sampling request.""" # Resolve model/base from sampling_session_id if provided if request.sampling_session_id is not None: @@ -1271,25 +1301,18 @@ async def get_sampling_model(request: SampleRequest, session: AsyncSession) -> ( return (request.base_model, request.model_path) -@app.post("/api/v1/asample", response_model=FutureResponse) -async def asample(request: SampleRequest, req: Request, session: AsyncSession = Depends(get_session)): - """Generates samples from the model (async version).""" - if request.sampling_session_id is not None and ":" in request.sampling_session_id: - raise HTTPException( - status_code=400, - detail="sampling_session_id must not contain ':' (the routing-key delimiter)", - ) +async def get_sampling_target(request: SampleRequest, req: Request, session: AsyncSession) -> SamplingTarget: + """Resolve and validate the model used by a sample request.""" - base_model, model_path = await get_sampling_model(request, session) + async def resolve() -> SamplingTarget: + base_model, model_path = await get_sampling_model(request, session) + if base_model: + return SamplingTarget(base_model=base_model, model_id="", checkpoint_id="") - if base_model: - model_id = checkpoint_id = "" - else: assert model_path is not None path = types.TinkerPath.parse(model_path) if ( not path - # Accept either tinker://model_id/checkpoint_id or tinker://model_id/sampler_weights/checkpoint_id or path.kind not in ("", "sampler_weights") or not (model_id := path.primary_id) or not (checkpoint_id := path.secondary_id) @@ -1299,21 +1322,51 @@ async def asample(request: SampleRequest, req: Request, session: AsyncSession = detail="model_path must be tinker://model_id/checkpoint_id or tinker://model_id/sampler_weights/checkpoint_id", ) await get_model(session, model_id) - # Validate that the checkpoint exists and is ready await validate_checkpoint(req, model_id, checkpoint_id, types.CheckpointType.SAMPLER, session) + return SamplingTarget(base_model=None, model_id=model_id, checkpoint_id=checkpoint_id) + + sampling_session_id = request.sampling_session_id + if sampling_session_id is None or not req.app.state.external_inference_client: + return await resolve() + + targets: dict[str, SamplingTarget] = req.app.state.external_sampling_targets + target = targets.get(sampling_session_id) + if target is not None: + return target + + locks: dict[str, asyncio.Lock] = req.app.state.external_sampling_target_locks + lock = locks.setdefault(sampling_session_id, asyncio.Lock()) + async with lock: + target = targets.get(sampling_session_id) + if target is None: + target = await resolve() + targets[sampling_session_id] = target + return target + + +@app.post("/api/v1/asample", response_model=FutureResponse) +async def asample(request: SampleRequest, req: Request, session: AsyncSession = Depends(get_session)): + """Generates samples from the model (async version).""" + if request.sampling_session_id is not None and ":" in request.sampling_session_id: + raise HTTPException( + status_code=400, + detail="sampling_session_id must not contain ':' (the routing-key delimiter)", + ) + + target = await get_sampling_target(request, req, session) request_id = await create_future( session=session, request_type=( types.RequestType.EXTERNAL if req.app.state.external_inference_client else types.RequestType.SAMPLE ), - model_id=model_id, + model_id=target.model_id, request_data=types.SampleInput( - base_model=base_model, + base_model=target.base_model, prompt=request.prompt.to_types(), sampling_params=request.sampling_params.to_types(), num_samples=request.num_samples, - checkpoint_id=checkpoint_id, + checkpoint_id=target.checkpoint_id, # A positive topk implies prompt logprobs: both are read off the same # prompt forward pass, so asking for one asks for the other. prompt_logprobs=bool(request.prompt_logprobs) or request.topk_prompt_logprobs > 0, @@ -1328,10 +1381,13 @@ async def asample(request: SampleRequest, req: Request, session: AsyncSession = if req.app.state.external_inference_client: asyncio.create_task( req.app.state.external_inference_client.call_and_store_result( - request_id, request, model_id, checkpoint_id, base_model=base_model + request_id, + request, + target.model_id, + target.checkpoint_id, + base_model=target.base_model, ) ) - return FutureResponse(future_id=str(request_id), status="pending", request_id=str(request_id)) @@ -1341,7 +1397,10 @@ async def get_server_capabilities(request: Request): supported_models = [ SupportedModel(model_name=request.app.state.engine_config.base_model), ] - return GetServerCapabilitiesResponse(supported_models=supported_models) + return GetServerCapabilitiesResponse( + supported_models=supported_models, + supported_loss_fns=sorted(types.SUPPORTED_LOSS_FNS), + ) class RetrieveFutureRequest(BaseModel): diff --git a/skyrl/tinker/engine.py b/skyrl/tinker/engine.py index 25b29e4ee1..537e58e718 100644 --- a/skyrl/tinker/engine.py +++ b/skyrl/tinker/engine.py @@ -5,6 +5,7 @@ from collections import defaultdict from contextlib import contextmanager from datetime import datetime, timedelta, timezone +from itertools import groupby from pathlib import Path from typing import Any, Callable @@ -377,6 +378,53 @@ def _find_destructive_barriers(self, session: Session) -> dict[str, int]: ) return dict(session.exec(query).all()) + def _find_next_sequence_ids(self, session: Session) -> dict[str, int]: + """Derive each model's next SDK sequence from its persisted futures.""" + query = ( + select(FutureDB.model_id, func.max(FutureDB.seq_id)) + .where(FutureDB.seq_id.is_not(None)) + .where(FutureDB.status != RequestStatus.PENDING) + .group_by(FutureDB.model_id) + ) + return {model_id: last_seq_id + 1 for model_id, last_seq_id in session.exec(query).all()} + + def _find_sequenced_requests( + self, session: Session, request_types: set[types.RequestType] + ) -> list[tuple[int, str, types.RequestType]]: + """Find each model's leading contiguous SDK requests while their types are allowed.""" + next_sequence_ids = self._find_next_sequence_ids(session) + pending = session.exec( + select(FutureDB.request_id, FutureDB.model_id, FutureDB.seq_id, FutureDB.request_type) + .where(FutureDB.seq_id.is_not(None)) + .where(FutureDB.status == RequestStatus.PENDING) + .order_by(FutureDB.model_id, FutureDB.seq_id) + ).all() + + batchable = [] + for model_id, requests in groupby(pending, key=lambda row: row.model_id): + requests = list(requests) + expected_seq_id = next_sequence_ids.get(model_id) + if expected_seq_id is None: + expected_seq_id = requests[0].seq_id + for request_id, _, seq_id, pending_type in requests: + if seq_id != expected_seq_id or pending_type not in request_types: + break + batchable.append((request_id, model_id, pending_type)) + expected_seq_id += 1 + return batchable + + def _find_sequenced_single_requests(self, session: Session) -> list[tuple[int, str, types.RequestType]]: + """Find leading contiguous SDK requests that do not run as model-pass batches.""" + return self._find_sequenced_requests( + session, + { + types.RequestType.OPTIM_STEP, + types.RequestType.SAVE_WEIGHTS_FOR_SAMPLER, + types.RequestType.SAVE_WEIGHTS, + types.RequestType.LOAD_WEIGHTS, + }, + ) + def find_batchable_model_passes( self, session: Session, request_type: types.RequestType ) -> dict[str, tuple[str, types.ForwardBackwardInput]]: @@ -399,16 +447,20 @@ def find_batchable_model_passes( select(FutureDB.request_id, FutureDB.model_id) .where(FutureDB.request_type == request_type) .where(FutureDB.status == RequestStatus.PENDING) + .where(FutureDB.seq_id.is_(None)) .order_by(FutureDB.request_id) ) ops = session.exec(query).all() # Filter: only include ops that come before their model's barrier batchable = [ + (request_id, model_id) for request_id, model_id, _ in self._find_sequenced_requests(session, {request_type}) + ] + batchable.extend( (request_id, model_id) for request_id, model_id in ops if model_id not in barriers or request_id < barriers[model_id] - ] + ) return { str(request_id): (model_id, types.ForwardBackwardInput.model_validate(request_data)) @@ -490,6 +542,7 @@ def find_single_requests(self, session: Session) -> dict[str, tuple[str, types.R statement = ( select(FutureDB.request_id, FutureDB.model_id, FutureDB.request_type) .where(FutureDB.status == RequestStatus.PENDING) + .where(FutureDB.seq_id.is_(None)) .where(FutureDB.request_type != types.RequestType.FORWARD_BACKWARD) .where(FutureDB.request_type != types.RequestType.FORWARD) .where(FutureDB.request_type != types.RequestType.SAMPLE) @@ -499,15 +552,18 @@ def find_single_requests(self, session: Session) -> dict[str, tuple[str, types.R other_futures = session.exec(statement).all() # Filter: only include ops that come before the first blocked pass for their model - other_futures = [ + legacy_futures = [ (request_id, model_id, request_type) for request_id, model_id, request_type in other_futures if model_id not in blocked_pass_barriers or request_id < blocked_pass_barriers[model_id] ] + sequenced_futures = self._find_sequenced_single_requests(session) return { str(request_id): (model_id, request_type, request_data) - for request_id, model_id, request_type, request_data in self._load_requests(session, other_futures) + for request_id, model_id, request_type, request_data in self._load_requests( + session, sequenced_futures + legacy_futures + ) } def process_create_model(self, model_id: str, request_data: types.CreateModelInput) -> types.CreateModelOutput: 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..160cbccdc5 100644 --- a/tests/backends/skyrl_train/test_token_based_batching_utils.py +++ b/tests/backends/skyrl_train/test_token_based_batching_utils.py @@ -160,6 +160,19 @@ def test_reorder_and_combine(self): for i in range(batch.batch_size): assert torch.equal(reordered["sequences"][i], batch["sequences"][i]) + def test_reorder_and_combine_items_drops_padding(self): + batch = self._make_batch([10, 3, 8, 5]) + iterator = TokenBasedBatchIterator(batch, max_tokens_per_microbatch=12) + output_batches = [ + [{"sample": index} for index in indices] + [{"sample": "padding"}] for indices in iterator._microbatches + ] + iterator._num_padding_microbatches = 1 + output_batches.append([{"sample": "padding-microbatch"}]) + + reordered = iterator.reorder_and_combine_items(output_batches) + + assert [output["sample"] for output in reordered] == list(range(batch.batch_size)) + def test_get_microbatch_iterator_factory(self): batch = self._make_batch([10, 10, 5, 5]) diff --git a/tests/tinker/skyrl_train/test_loss_normalization.py b/tests/tinker/skyrl_train/test_loss_normalization.py index dff334598a..e802a6254e 100644 --- a/tests/tinker/skyrl_train/test_loss_normalization.py +++ b/tests/tinker/skyrl_train/test_loss_normalization.py @@ -8,6 +8,7 @@ from __future__ import annotations import pytest +import torch # Skip if skyrl_train_backend.py cannot be imported skyrl_train_backend = pytest.importorskip("skyrl.backends.skyrl_train_backend") @@ -21,6 +22,86 @@ def test_ppo_thresholds_map_to_eps_clip(): assert config == pytest.approx({"eps_clip_low": 0.2, "eps_clip_high": 0.28}) +def test_gspo_thresholds_map_to_native_sequence_loss(): + loss_fn, config = _normalize( + None, + "policy", + "gspo", + {"clip_low_threshold": 0.98, "clip_high_threshold": 1.03}, + ) + assert loss_fn == "gspo" + assert config == { + "eps_clip_low": pytest.approx(0.02), + "eps_clip_high": pytest.approx(0.03), + "loss_reduction": "sequence_mean", + "use_entropy_loss": False, + "use_kl_loss": False, + } + + +def test_gspo_batch_normalization_gives_each_trainable_sequence_equal_weight(): + batch = skyrl_train_backend.TrainingInputBatch( + { + "advantages": torch.tensor([[2.0, 2.0, 0.0, 0.0], [-1.0, -1.0, -1.0, 0.0], [0.0] * 4]), + "loss_mask": torch.ones(3, 4), + } + ) + + skyrl_train_backend.SkyRLTrainBackend._normalize_gspo_batch(batch) + + assert batch["loss_mask"].tolist() == [[1.0, 1.0, 0.0, 0.0], [1.0, 1.0, 1.0, 0.0], [0.0] * 4] + assert (batch["advantages"] * batch["loss_mask"]).sum(dim=-1).tolist() == pytest.approx([1.0, -0.5, 0.0]) + + +def test_gspo_selects_only_rows_with_gradient_contributions(): + batch = skyrl_train_backend.TrainingInputBatch( + { + "sequences": torch.tensor([[1, 2], [3, 0], [4, 5]]), + "attention_mask": torch.tensor([[1, 1], [1, 0], [1, 1]]), + "loss_mask": torch.tensor([[1.0, 0.0], [0.0, 0.0], [0.0, 1.0]]), + } + ) + batch.metadata = {"response_length": 2} + + selected, indices = skyrl_train_backend.SkyRLTrainBackend._select_trainable_gspo_rows(batch) + + assert indices == [0, 2] + assert selected["sequences"].tolist() == [[1, 2], [4, 5]] + assert selected.metadata == batch.metadata + + +def test_gspo_keeps_entirely_inactive_batch_for_zero_gradient_backward(): + batch = skyrl_train_backend.TrainingInputBatch( + { + "sequences": torch.tensor([[1, 2], [3, 4]]), + "loss_mask": torch.zeros(2, 2), + } + ) + + selected, indices = skyrl_train_backend.SkyRLTrainBackend._select_trainable_gspo_rows(batch) + + assert selected is batch + assert indices == [0, 1] + + +def test_gspo_pruning_metrics_sum_across_sdk_chunks_without_duplication(): + first = skyrl_train_backend.SkyRLTrainBackend._get_gspo_pruning_metrics([3, 2, 4, 1], {0, 2}, 0, 2) + second = skyrl_train_backend.SkyRLTrainBackend._get_gspo_pruning_metrics([3, 2, 4, 1], {0, 2}, 2, 4) + + assert first == { + "skyrl.ai/submitted_datums:sum": 1.0, + "skyrl.ai/skipped_datums:sum": 1.0, + "skyrl.ai/submitted_tokens:sum": 3.0, + "skyrl.ai/skipped_tokens:sum": 2.0, + } + assert second == { + "skyrl.ai/submitted_datums:sum": 1.0, + "skyrl.ai/skipped_datums:sum": 1.0, + "skyrl.ai/submitted_tokens:sum": 4.0, + "skyrl.ai/skipped_tokens:sum": 1.0, + } + + def test_dppo_deltas_are_nested_under_dppo(): loss_fn, config = _normalize(None, "policy", "dppo", {"delta_low": 0.2, "delta_high": 0.3}) assert loss_fn == "dppo" diff --git a/tests/tinker/test_api_validation.py b/tests/tinker/test_api_validation.py index 3207645704..04436d69cf 100644 --- a/tests/tinker/test_api_validation.py +++ b/tests/tinker/test_api_validation.py @@ -1,4 +1,5 @@ import base64 +from types import SimpleNamespace import pytest from pydantic import TypeAdapter, ValidationError @@ -27,6 +28,42 @@ def test_forward_backward_input_accepts_ppo_threshold_keys(): assert req.loss_fn_config == {"clip_low_threshold": 0.9, "clip_high_threshold": 1.1} +def test_forward_backward_input_accepts_gspo_threshold_keys(): + req = api.ForwardBackwardInput( + data=[_make_datum()], + loss_fn="gspo", + loss_fn_config={"clip_low_threshold": 0.98, "clip_high_threshold": 1.03}, + ) + assert req.to_types().loss_fn == "gspo" + + +@pytest.mark.parametrize( + "loss_fn_config", + [None, {"clip_low_threshold": 0.98}, {"clip_low_threshold": 1.01, "clip_high_threshold": 1.03}], +) +def test_forward_backward_input_rejects_incomplete_or_invalid_gspo_thresholds(loss_fn_config): + with pytest.raises(ValidationError, match="loss_fn='gspo'"): + api.ForwardBackwardInput(data=[_make_datum()], loss_fn="gspo", loss_fn_config=loss_fn_config) + + +@pytest.mark.asyncio +async def test_server_capabilities_advertise_native_gspo(): + request = SimpleNamespace( + app=SimpleNamespace(state=SimpleNamespace(engine_config=SimpleNamespace(base_model="test-model"))) + ) + capabilities = await api.get_server_capabilities(request) + + assert "gspo" in capabilities.supported_loss_fns + + +def test_weight_requests_preserve_sdk_sequence_ids(): + save_request = api.SaveWeightsRequest(model_id="model", path="checkpoint", seq_id=7) + load_request = api.LoadWeightsRequest(model_id="model", path="tinker://model/weights/checkpoint", seq_id=8) + + assert save_request.seq_id == 7 + assert load_request.seq_id == 8 + + def test_forward_backward_input_accepts_ppo_value_clip(): req = api.ForwardBackwardInput( data=[_make_datum()], diff --git a/tests/tinker/test_engine.py b/tests/tinker/test_engine.py index 46af0332b0..c7604c3a8f 100644 --- a/tests/tinker/test_engine.py +++ b/tests/tinker/test_engine.py @@ -109,6 +109,15 @@ def test_cleanup_stale_sessions(): [], id="cispo", ), + pytest.param( + "gspo", + {"clip_low_threshold": 0.98, "clip_high_threshold": 1.03}, + [0.1, 0.2, 0.3], + [-1.1, -1.0, -0.9], + [], + [], + id="gspo", + ), pytest.param("ppo_critic", {"value_clip": 0.2}, [], [], [0.4, 0.5, 0.6], [0.7, 0.8, 0.9], id="ppo_critic"), pytest.param( "dppo", @@ -394,6 +403,104 @@ def test_find_batchable_model_passes_stops_at_barrier(scheduling_engine): assert request_data.data[0].loss_fn_inputs.target_tokens.data == [1, 2, 3] +def test_find_batchable_model_passes_waits_for_next_sdk_sequence(scheduling_engine): + """Later chunks must wait for the first chunk that the SDK deliberately submits last.""" + engine = scheduling_engine + payload = forward_backward_payload() + with Session(engine.db_engine) as session: + session.add( + FutureDB( + request_type=types.RequestType.SAVE_WEIGHTS_FOR_SAMPLER, + model_id="model_a", + seq_id=1, + request_data={}, + status=RequestStatus.COMPLETED, + ) + ) + later_chunks = [ + FutureDB( + request_type=types.RequestType.FORWARD_BACKWARD, + model_id="model_a", + seq_id=seq_id, + request_data=payload, + status=RequestStatus.PENDING, + ) + for seq_id in [3, 4] + ] + session.add_all(later_chunks) + session.commit() + later_request_ids = [row.request_id for row in later_chunks] + + with Session(engine.db_engine) as session: + assert engine.find_batchable_model_passes(session, types.RequestType.FORWARD_BACKWARD) == {} + + with Session(engine.db_engine) as session: + first_chunk = FutureDB( + request_type=types.RequestType.FORWARD_BACKWARD, + model_id="model_a", + seq_id=2, + request_data=payload, + status=RequestStatus.PENDING, + ) + session.add(first_chunk) + session.commit() + first_request_id = first_chunk.request_id + + with Session(engine.db_engine) as session: + batchable = engine.find_batchable_model_passes(session, types.RequestType.FORWARD_BACKWARD) + + assert list(batchable) == [str(first_request_id), *(str(request_id) for request_id in later_request_ids)] + + +@pytest.mark.parametrize("initial_seq_id", [0, 1]) +def test_find_single_requests_accepts_sdk_sequence_bootstraps(scheduling_engine, initial_seq_id): + engine = scheduling_engine + with Session(engine.db_engine) as session: + initial_request = FutureDB( + request_type=types.RequestType.SAVE_WEIGHTS_FOR_SAMPLER, + model_id="model_a", + seq_id=initial_seq_id, + request_data={}, + status=RequestStatus.PENDING, + ) + session.add(initial_request) + session.commit() + initial_request_id = initial_request.request_id + + with Session(engine.db_engine) as session: + assert engine.find_single_requests(session) == { + str(initial_request_id): ( + "model_a", + types.RequestType.SAVE_WEIGHTS_FOR_SAMPLER, + {}, + ) + } + + +def test_find_single_requests_bootstraps_resumed_sdk_sequence(scheduling_engine): + engine = scheduling_engine + with Session(engine.db_engine) as session: + resumed_request = FutureDB( + request_type=types.RequestType.OPTIM_STEP, + model_id="model_a", + seq_id=100, + request_data={}, + status=RequestStatus.PENDING, + ) + session.add(resumed_request) + session.commit() + resumed_request_id = resumed_request.request_id + + with Session(engine.db_engine) as session: + assert engine.find_single_requests(session) == { + str(resumed_request_id): ( + "model_a", + types.RequestType.OPTIM_STEP, + {}, + ) + } + + def sample_payload(checkpoint_id: str) -> dict: return types.SampleInput( prompt=types.ModelInput(chunks=[types.EncodedTextChunk(tokens=[1, 2])]), diff --git a/tests/tinker/test_sampling_target_cache.py b/tests/tinker/test_sampling_target_cache.py new file mode 100644 index 0000000000..a46acfd0e0 --- /dev/null +++ b/tests/tinker/test_sampling_target_cache.py @@ -0,0 +1,107 @@ +import asyncio +from types import SimpleNamespace +from unittest.mock import AsyncMock + +import pytest +from fastapi import HTTPException + +from skyrl.tinker import api +from skyrl.tinker.db_models import SamplingSessionDB + + +class _SamplingSession: + base_model = None + model_path = "tinker://model_a/sampler_weights/checkpoint_a" + + +class _Session: + def __init__(self, reads: list[str]): + self._reads = reads + + async def get(self, model, sampling_session_id): + assert model is SamplingSessionDB + self._reads.append(sampling_session_id) + await asyncio.sleep(0.01) + return _SamplingSession() + + +def _request(): + return SimpleNamespace( + sampling_session_id="sampling_a", + base_model=None, + model_path=None, + ) + + +def _server_request(external_inference_client=object()): + return SimpleNamespace( + app=SimpleNamespace( + state=SimpleNamespace( + external_inference_client=external_inference_client, + external_sampling_targets={}, + external_sampling_target_locks={}, + ) + ) + ) + + +@pytest.mark.asyncio +async def test_external_sampling_target_resolves_once_under_concurrency(monkeypatch): + reads = [] + get_model = AsyncMock() + validate_checkpoint = AsyncMock() + monkeypatch.setattr(api, "get_model", get_model) + monkeypatch.setattr(api, "validate_checkpoint", validate_checkpoint) + request = _request() + server_request = _server_request() + + targets = await asyncio.gather( + *(api.get_sampling_target(request, server_request, _Session(reads)) for _ in range(512)) + ) + + assert targets == [api.SamplingTarget(None, "model_a", "checkpoint_a")] * 512 + assert reads == ["sampling_a"] + get_model.assert_awaited_once() + validate_checkpoint.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_external_sampling_target_caches_only_successful_validation(monkeypatch): + reads = [] + monkeypatch.setattr(api, "get_model", AsyncMock()) + validate_checkpoint = AsyncMock( + side_effect=[ + HTTPException(status_code=425, detail="Checkpoint is still being created"), + None, + ] + ) + monkeypatch.setattr(api, "validate_checkpoint", validate_checkpoint) + request = _request() + server_request = _server_request() + + with pytest.raises(HTTPException, match="Checkpoint is still being created"): + await api.get_sampling_target(request, server_request, _Session(reads)) + target = await api.get_sampling_target(request, server_request, _Session(reads)) + cached_target = await api.get_sampling_target(request, server_request, _Session(reads)) + + assert target == cached_target == api.SamplingTarget(None, "model_a", "checkpoint_a") + assert reads == ["sampling_a", "sampling_a"] + assert validate_checkpoint.await_count == 2 + + +@pytest.mark.asyncio +async def test_internal_sampling_target_preserves_database_validation(monkeypatch): + reads = [] + get_model = AsyncMock() + validate_checkpoint = AsyncMock() + monkeypatch.setattr(api, "get_model", get_model) + monkeypatch.setattr(api, "validate_checkpoint", validate_checkpoint) + request = _request() + server_request = _server_request(external_inference_client=None) + + await api.get_sampling_target(request, server_request, _Session(reads)) + await api.get_sampling_target(request, server_request, _Session(reads)) + + assert reads == ["sampling_a", "sampling_a"] + assert get_model.await_count == 2 + assert validate_checkpoint.await_count == 2