diff --git a/.github/workflows/pr-test.yml b/.github/workflows/pr-test.yml index 34e15bcbeb..7f5481aa53 100644 --- a/.github/workflows/pr-test.yml +++ b/.github/workflows/pr-test.yml @@ -573,6 +573,30 @@ jobs: "num_gpus": 0, "test_file": "test_megatron_argument_validation.py" }, + { + "num_gpus": 0, + "test_file": "test_rs_exact_refill_utils.py" + }, + { + "num_gpus": 0, + "test_file": "test_rs_exact_refill_actor.py" + }, + { + "num_gpus": 0, + "test_file": "test_rs_exact_refill_rollout_manager.py" + }, + { + "num_gpus": 0, + "test_file": "test_rs_exact_refill_sglang_rollout.py" + }, + { + "num_gpus": 0, + "test_file": "test_rs_exact_refill_train_async.py" + }, + { + "num_gpus": 0, + "test_file": "test_data_source_epoch_wrap.py" + }, { "num_gpus": 0, "test_file": "test_deep_ep_tms_patch.py" diff --git a/.github/workflows/pr-test.yml.j2 b/.github/workflows/pr-test.yml.j2 index 6ed6147126..fab8a28524 100644 --- a/.github/workflows/pr-test.yml.j2 +++ b/.github/workflows/pr-test.yml.j2 @@ -69,6 +69,12 @@ 'extra_pip_deps': 'transformers wandb', 'tests': [ {'test_file': 'test_megatron_argument_validation.py', 'num_gpus': 0}, + {'test_file': 'test_rs_exact_refill_utils.py', 'num_gpus': 0}, + {'test_file': 'test_rs_exact_refill_actor.py', 'num_gpus': 0}, + {'test_file': 'test_rs_exact_refill_rollout_manager.py', 'num_gpus': 0}, + {'test_file': 'test_rs_exact_refill_sglang_rollout.py', 'num_gpus': 0}, + {'test_file': 'test_rs_exact_refill_train_async.py', 'num_gpus': 0}, + {'test_file': 'test_data_source_epoch_wrap.py', 'num_gpus': 0}, {'test_file': 'test_deep_ep_tms_patch.py', 'num_gpus': 0}, {'test_file': 'test_stateless_adam.py', 'num_gpus': 0}, {'test_file': 'utils/test_megatron_server_arguments.py', 'num_gpus': 0}, diff --git a/docs/en/get_started/usage.md b/docs/en/get_started/usage.md index 2ecebaba65..7aa0edfdd1 100644 --- a/docs/en/get_started/usage.md +++ b/docs/en/get_started/usage.md @@ -205,6 +205,7 @@ The recommended contract is to put the source identifier in `metadata["source_na Note: On-policy distillation (OPD) is now orthogonal to the advantage estimator. Use `--use-opd` and `--opd-kl-coef` to enable OPD on top of any estimator. - `--calculate-per-token-loss`: By default, slime calculates loss on a per-sample basis, i.e., `mean(sum(sample_i) / len(sample_i))`. Enable this flag to calculate loss on a per-token basis, i.e., `sum(sum(sample_i)) / sum(len(sample_i))`. - `--use-tis`: Enable this setting to use TIS (Truncated Importance Sampling) (https://fengyao.notion.site/off-policy-rl). +- `--rs-batch-refill`: With sequence/geometric RS on disaggregated `train_async.py`, preflight complete prompt groups and reactively generate only the missing groups before an optimizer step. `--rs-refill-max-rounds` bounds retries; exhaustion aborts the step instead of silently using an underfilled effective batch. This is a correctness option with additional actor preflight and rollout latency, not an unconditional speedup. See [Rollout Correction Methods](../../../examples/train_infer_mismatch_helper/README.md#effective-batch-refill-after-sequence-rs) for the current constraints and a complete command. #### GRPO Algorithm diff --git a/docs/zh/get_started/usage.md b/docs/zh/get_started/usage.md index 533729c40a..c5f304a373 100644 --- a/docs/zh/get_started/usage.md +++ b/docs/zh/get_started/usage.md @@ -209,6 +209,7 @@ sglang 的加载非常简单，只需要： 注意：在策略蒸馏 (OPD) 现在与 advantage estimator 正交，使用 `--use-opd` 和 `--opd-kl-coef` 可以在任意 estimator 之上启用 OPD。 - `--calculate-per-token-loss`：slime 中默认的方案是 per sample loss，即 `mean(sum(sample_i) / len(sample_i))`，如果需要计算 per token loss，即 `sum(sum(sample_i)) / sum(len(sample_i))`，可以开启 `--calculate-per-token-loss`； - `--use-tis`：如果需要开启 tis（https://fengyao.notion.site/off-policy-rl），可以开启这一设置； +- `--rs-batch-refill`：在训推分离的 `train_async.py` 中使用 sequence/geometric RS 时，先按完整 prompt group 做 actor preflight，再在 optimizer step 前只生成缺失的 group。`--rs-refill-max-rounds` 限制补齐轮数；超过上限会中止该 step，而不是静默使用缩水后的 effective batch。该选项增加 actor preflight 和补采样延迟，属于 correctness 机制，不保证无条件提速。完整命令与当前限制见 [Rollout Correction Methods](../../../examples/train_infer_mismatch_helper/README.md#effective-batch-refill-after-sequence-rs)。 #### GRPO 算法 diff --git a/examples/train_infer_mismatch_helper/README.md b/examples/train_infer_mismatch_helper/README.md index e943071148..e53b9e9acf 100644 --- a/examples/train_infer_mismatch_helper/README.md +++ b/examples/train_infer_mismatch_helper/README.md @@ -22,6 +22,61 @@ You may specify the **IS/RS configs** with a config file using `--custom-config- `--get-mismatch-metrics`: When you don't want to add TIS/MIS, but still want to monitor the mismatch-related metrics (e.g. rollout-training KL). It will **only return mismatch metrics** but not change the loss in any way. +### Effective-batch refill after sequence RS + +Sequence-level RS can leave an optimizer step with fewer independent prompt groups than configured. In disaggregated +`train_async.py`, `--rs-batch-refill` makes that loss of effective batch cardinality fail-safe and explicit: + +1. Generate exactly `rollout_batch_size` initial prompt groups; there is no speculative over-generation. +2. Before `optimizer.step`, recompute proximal log probabilities with the actor and apply the configured sequence or + geometric RS gate atomically to each complete prompt group. +3. Keep only the selected groups and their in-memory proximal log-probability cache. Generate the exact deficit, + rounded only to the smallest DP/VPP scheduling multiple, from the current rollout policy. +4. Repeat for at most `--rs-refill-max-rounds`. If the target batch is still incomplete, fail the job before an + optimizer step instead of silently training on an underfilled batch. + +After an initial or replacement candidate generation returns, every coordinator and actor wait in the refill loop is +bounded by `--rs-refill-rpc-timeout-seconds` (30 minutes by default). Candidate generation retains the rollout +backend's existing timeout and health-monitor behavior. `--rs-refill-max-candidate-cache-bytes` (1 GiB by default) +bounds the proximal-logprob tensor payload retained by any one process: an actor rank checks its current candidate +round before allocating pinned CPU memory, and the RolloutManager checks its accumulated accepted cache before pulling +the selected tensors from Ray. `peak_actor_candidate_logprob_cache_bytes` is the per-actor high-water mark to compare +with that limit; `aggregate_candidate_logprob_cache_bytes` and +`peak_aggregate_candidate_logprob_cache_bytes` report cumulative and per-round aggregate payload across reporting +actors. Selected-transfer and manager-retained metrics cover the coordinator side. Leave headroom for Python +containers, Ray object-store/transport buffers, and other process memory, which are not included in this tensor-payload +limit. The limit also does not reserve CUDA allocator headroom: training pipeline ranks that receive a final batch may +materialize their DP-local proximal-logprob shard on device, so long-response jobs must budget GPU memory for it. + +The refill path still applies TIS during training; completing the batch does not make stale initial trajectories +on-policy. Initial candidates are limited to one policy version of staleness and reactive replacements must match the +actor version used for preflight. + +```bash +python train_async.py \ + ... \ + --rs-batch-refill \ + --rs-refill-max-rounds 2 \ + --rs-refill-rpc-timeout-seconds 1800 \ + --rs-refill-max-candidate-cache-bytes 1073741824 \ + --update-weights-interval 1 \ + --use-tis \ + --custom-config-path examples/train_infer_mismatch_helper/mis_refill.yaml +``` + +This first implementation intentionally supports fresh/finetune runs from rollout 0 only. It requires the +disaggregated Megatron path, a global rollout dataset, one optimizer step per rollout batch, resident actor and rollout +engines, full NCCL weight updates, dynamic batching, zero actor dropout, the default model provider, and non-quantized +Megatron training. Checkpoint resume, fully async rollout, token-level RS, external rollout engines, and custom +model/loss/data hooks are rejected during argument validation. Dynamic sampling filters and custom rollout logging are +handled conservatively: dynamic sampling filters are rejected because they may change the exact candidate set, while a +custom rollout logger or debug saver is allowed only if a final train-data fingerprint proves that it left policy +inputs unchanged. + +Refill adds actor preflight work and replacement latency, so it is a correctness option rather than an unconditional +wall-clock speedup. Its cost depends on the observed rejection rate and the rollout/training time balance. + + ## Algorithms We give examples of the algorithms for solving the training-inference mismatch issue. diff --git a/examples/train_infer_mismatch_helper/mis_refill.yaml b/examples/train_infer_mismatch_helper/mis_refill.yaml new file mode 100644 index 0000000000..5b2905d6cc --- /dev/null +++ b/examples/train_infer_mismatch_helper/mis_refill.yaml @@ -0,0 +1,22 @@ +# TIS weights correct the policy-gradient loss but do not mask individual tokens. +# Keep batch normalization disabled: RS refill cannot yet normalize globally +# across DP-local dynamic microbatches. +use_tis: true +tis_level: "token" +tis_mode: "truncate" +tis_lower_bound: 0.1 +tis_upper_bound: 10.0 +tis_batch_normalize: false + +# Proximal replay also requires rollout_temperature > 0, 0 < rollout_top_p <= 1, +# rollout_top_k: -1, attention_dropout: 0, and hidden_dropout: 0 in the CLI or +# surrounding custom configuration. +attention_dropout: 0.0 +hidden_dropout: 0.0 + +# RS is a sequence-level gate so a rejected unit maps to an unambiguous batch deficit. +use_rs: true +rs_level: "geometric" +rs_lower_bound: 0.01 +rs_upper_bound: 100.0 +rs_veto_threshold: null diff --git a/slime/backends/megatron_utils/actor.py b/slime/backends/megatron_utils/actor.py index 1d20abc96f..5f55179a58 100644 --- a/slime/backends/megatron_utils/actor.py +++ b/slime/backends/megatron_utils/actor.py @@ -29,6 +29,7 @@ reload_process_groups, ) from slime.utils.routing_replay import RoutingReplay +from slime.utils.rs_refill import clone_rs_masks, compute_sequence_rs_masks, validate_final_rs_masks from slime.utils.types import RolloutBatch from ...utils.tensor_backper import TensorBackuper @@ -70,6 +71,7 @@ def init( monkey_patch_torch_dist() super().init(args, role, with_ref, with_opd_teacher) + self._rs_candidate_log_probs: dict[int, dict[int, torch.Tensor]] = {} # Destroying and recreating WORLD invalidates raw dist.group.WORLD references cached by external code. # Set SLIME_DESTROY_WORLD_PROCESS_GROUP=0 when such references may outlive a train sleep/wake cycle. if os.getenv("SLIME_DESTROY_WORLD_PROCESS_GROUP", "1").lower() not in {"0", "false", "no"}: @@ -281,7 +283,7 @@ def _get_rollout_data(self, rollout_data_ref: Box) -> RolloutBatch: for mm_dict in rollout_data["multimodal_train_inputs"] ] - for key in ["rollout_log_probs", "teacher_log_probs"]: + for key in ["rollout_log_probs", "rs_preflight_log_probs", "teacher_log_probs"]: if key not in rollout_data: continue rollout_data[key] = [ @@ -371,6 +373,152 @@ def compute_log_prob( use_rollout_top_p_replay=True, ) + def score_rs_candidates(self, rollout_id: int, rollout_data_ref: Box): + """Score candidates with the proximal policy and retain logprobs locally.""" + + rollout_data = self._get_rollout_data(rollout_data_ref) + data_iterator = get_data_iterator(rollout_data) + target_tag = "old_actor" if self.args.keep_old_actor else "actor" + previous_tag = self._active_model_tag + if previous_tag != target_tag: + self._switch_model(target_tag) + try: + rollout_data.update( + self.compute_log_prob( + data_iterator, + rollout_data["num_microbatches"], + store_prefix="", + ) + ) + finally: + if previous_tag is not None and self._active_model_tag != previous_tag: + self._switch_model(previous_tag) + + if not mpu.is_pipeline_last_stage(): + return [] + + train_log_probs = rollout_data.get("log_probs") + if not train_log_probs: + raise RuntimeError("RS refill preflight did not produce training log probabilities") + if len(train_log_probs) != len(rollout_data["sample_indices"]): + raise RuntimeError( + "RS refill preflight logprob/sample count mismatch: " + f"{len(train_log_probs)} != {len(rollout_data['sample_indices'])}" + ) + if len(train_log_probs) != len(rollout_data["group_indices"]): + raise RuntimeError( + "RS refill preflight logprob/group count mismatch: " + f"{len(train_log_probs)} != {len(rollout_data['group_indices'])}" + ) + + original_loss_masks = clone_rs_masks(rollout_data["loss_masks"]) + modified_masks = compute_sequence_rs_masks( + args=self.args, + train_log_probs=train_log_probs, + rollout_log_probs=rollout_data["rollout_log_probs"], + loss_masks=rollout_data["loss_masks"], + ) + + # Every TP rank executes the gate; only TP rank 0 owns the replicated + # report and actor-local cache transferred back through Ray. + if mpu.get_tensor_model_parallel_rank() != 0: + return [] + if rollout_id in self._rs_candidate_log_probs: + raise RuntimeError(f"RS candidate logprob cache already exists for rollout_id={rollout_id}") + if len(original_loss_masks) != len(modified_masks): + raise RuntimeError( + "RS preflight gate returned a different sample count: " + f"original={len(original_loss_masks)}, modified={len(modified_masks)}" + ) + if len(original_loss_masks) != len(rollout_data["loss_masks"]): + raise RuntimeError("RS preflight gate mutated the input loss-mask list length") + validate_final_rs_masks(original_loss_masks, rollout_data["loss_masks"]) + + cache_bytes_by_sample = [train_log_prob.numel() * 4 for train_log_prob in train_log_probs] + candidate_cache_bytes = sum(cache_bytes_by_sample) + if candidate_cache_bytes > self.args.rs_refill_max_candidate_cache_bytes: + raise RuntimeError( + "RS candidate proximal-logprob cache exceeds the per-actor limit before pinned-memory allocation: " + f"required={candidate_cache_bytes}, limit={self.args.rs_refill_max_candidate_cache_bytes}. " + "Increase --rs-refill-max-candidate-cache-bytes or reduce the candidate batch/response length." + ) + + stats = [] + for original_mask, modified_mask in zip(original_loss_masks, modified_masks, strict=True): + valid_tokens = original_mask.sum().to(torch.long) + if modified_mask.shape == original_mask.shape: + gate_passed = torch.all(modified_mask.float() == original_mask.float()).to(torch.long) + else: + gate_passed = original_mask.new_tensor(0, dtype=torch.long) + stats.append(torch.stack((valid_tokens, gate_passed))) + + stats = torch.stack(stats) + if stats.is_cuda: + stats_cpu = torch.empty_like(stats, device="cpu", pin_memory=True) + stats_cpu.copy_(stats, non_blocking=True) + proximal_log_probs = [] + for train_log_prob in train_log_probs: + cpu_log_prob = torch.empty( + train_log_prob.shape, + dtype=torch.float32, + device="cpu", + pin_memory=True, + ) + cpu_log_prob.copy_(train_log_prob.detach(), non_blocking=True) + proximal_log_probs.append(cpu_log_prob) + torch.cuda.current_stream(device=stats.device).synchronize() + else: + stats_cpu = stats.cpu() + proximal_log_probs = [ + train_log_prob.detach().float().cpu().contiguous() for train_log_prob in train_log_probs + ] + + stats_rows = stats_cpu.tolist() + sample_indices = [int(sample_index) for sample_index in rollout_data["sample_indices"]] + if len(sample_indices) != len(set(sample_indices)): + raise RuntimeError("RS preflight received duplicate sample indices on one actor rank") + self._rs_candidate_log_probs[rollout_id] = dict(zip(sample_indices, proximal_log_probs, strict=True)) + + reports = [] + for sample_index, group_index, stats_row, cache_bytes in zip( + sample_indices, + rollout_data["group_indices"], + stats_rows, + cache_bytes_by_sample, + strict=True, + ): + valid_tokens, gate_passed = stats_row + reports.append( + { + "sample_index": sample_index, + "group_index": int(group_index), + "valid_tokens": valid_tokens, + "gate_passed": bool(gate_passed), + "policy_version": str(self.weight_updater.weight_version), + "candidate_cache_bytes": cache_bytes, + } + ) + return reports + + def take_rs_candidate_log_probs(self, rollout_id: int, selected_sample_indices: list[int]): + """Return selected local caches once and discard all other candidates.""" + + if not mpu.is_pipeline_last_stage() or mpu.get_tensor_model_parallel_rank() != 0: + return {} + + selected = [int(sample_index) for sample_index in selected_sample_indices] + if len(selected) != len(set(selected)): + raise ValueError("Selected RS sample indices must be unique") + cache = self._rs_candidate_log_probs.pop(rollout_id, None) + if cache is None: + raise RuntimeError(f"No RS candidate logprob cache for rollout_id={rollout_id}") + return {sample_index: cache[sample_index] for sample_index in selected if sample_index in cache} + + def discard_rs_candidate_log_probs(self, rollout_id: int) -> None: + """Idempotently discard a candidate cache after coordination fails.""" + + self._rs_candidate_log_probs.pop(rollout_id, None) + def train(self, rollout_id: int, rollout_data_ref: Box, external_data=None): if self.args.debug_rollout_only: return None @@ -422,6 +570,42 @@ def train_critic(self, rollout_id: int, rollout_data: RolloutBatch): return {} def train_actor(self, rollout_id: int, rollout_data: RolloutBatch, external_data=None) -> None: + rs_preflight_log_probs = rollout_data.pop("rs_preflight_log_probs", None) + has_rs_preflight_log_probs = rs_preflight_log_probs is not None + rs_preflight_response_lengths = None + rs_preflight_loss_masks = None + if getattr(self.args, "rs_batch_refill", False): + if not has_rs_preflight_log_probs: + raise RuntimeError("RS-refilled training batch is missing its proximal logprob cache") + if "log_probs" in rollout_data: + raise RuntimeError("RS-refilled training batch contains conflicting proximal logprob fields") + sample_count = len(rollout_data["response_lengths"]) + if not (len(rs_preflight_log_probs) == len(rollout_data["loss_masks"]) == sample_count): + raise RuntimeError( + "RS-refilled proximal cache/sample count mismatch: " + f"log_probs={len(rs_preflight_log_probs)}, masks={len(rollout_data['loss_masks'])}, " + f"response_lengths={sample_count}" + ) + for sample_position, (log_probs, loss_mask, response_length) in enumerate( + zip( + rs_preflight_log_probs, + rollout_data["loss_masks"], + rollout_data["response_lengths"], + strict=True, + ) + ): + if log_probs.shape != loss_mask.shape or log_probs.numel() != int(response_length): + raise RuntimeError( + "RS-refilled proximal cache shape mismatch: " + f"sample={sample_position}, log_probs={tuple(log_probs.shape)}, " + f"mask={tuple(loss_mask.shape)}, response_length={response_length}" + ) + rollout_data["log_probs"] = rs_preflight_log_probs + rs_preflight_response_lengths = tuple(rollout_data["response_lengths"]) + rs_preflight_loss_masks = clone_rs_masks(rollout_data["loss_masks"]) + elif has_rs_preflight_log_probs: + raise RuntimeError("Received an RS preflight cache while --rs-batch-refill is disabled") + # Create data iterator for log_probs and train. data_iterator = get_data_iterator(rollout_data) num_microbatches = rollout_data["num_microbatches"] @@ -471,8 +655,10 @@ def train_actor(self, rollout_id: int, rollout_data: RolloutBatch, external_data and self.args.advantage_estimator != "gspo" ) if ( - not self.args.use_rollout_logprobs or self.args.get_mismatch_metrics - ) and not can_reuse_log_probs_in_loss: + (not self.args.use_rollout_logprobs or self.args.get_mismatch_metrics) + and not can_reuse_log_probs_in_loss + and not has_rs_preflight_log_probs + ): if self.args.use_routing_replay: if self.args.use_rollout_routing_replay: os.environ["ROUTING_REPLAY_STAGE"] = "replay_forward" @@ -511,6 +697,11 @@ def train_actor(self, rollout_id: int, rollout_data: RolloutBatch, external_data rollout_data, ) + if has_rs_preflight_log_probs: + if tuple(rollout_data["response_lengths"]) != rs_preflight_response_lengths: + raise RuntimeError("RS-refilled response lengths changed after proximal preflight") + validate_final_rs_masks(rs_preflight_loss_masks, rollout_data["loss_masks"]) + # Train if self.args.use_routing_replay: os.environ["ROUTING_REPLAY_STAGE"] = "replay_backward" diff --git a/slime/backends/megatron_utils/loss.py b/slime/backends/megatron_utils/loss.py index e7ac97c19e..574ea2d4b2 100644 --- a/slime/backends/megatron_utils/loss.py +++ b/slime/backends/megatron_utils/loss.py @@ -22,6 +22,7 @@ get_reinforce_plus_plus_baseline_advantages, get_reinforce_plus_plus_returns, ) +from slime.utils.rs_refill import apply_rs_refill_tis, clone_rs_masks, validate_final_rs_masks from slime.utils.types import RolloutBatch from .cp_utils import ( @@ -1071,11 +1072,16 @@ def policy_loss_function( "response_lengths": response_lengths, } - if args.custom_tis_function_path is not None: + original_rs_masks = clone_rs_masks(batch["loss_masks"]) if getattr(args, "rs_batch_refill", False) else None + if getattr(args, "rs_batch_refill", False): + tis_func = apply_rs_refill_tis + elif args.custom_tis_function_path is not None: tis_func = load_function(args.custom_tis_function_path) else: tis_func = vanilla_tis_function pg_loss, modified_response_masks, tis_metrics = tis_func(**tis_kwargs) + if getattr(args, "rs_batch_refill", False): + validate_final_rs_masks(original_rs_masks, batch["loss_masks"], modified_response_masks) # [decouple IS and rejection] Rebuild sum_of_sample_mean with # modified_response_masks for numerator correction (rejected tokens diff --git a/slime/observability/rollout_data_utils.py b/slime/observability/rollout_data_utils.py index de2c7db136..837dade984 100644 --- a/slime/observability/rollout_data_utils.py +++ b/slime/observability/rollout_data_utils.py @@ -13,6 +13,7 @@ "tokens": torch.long, "loss_masks": torch.int, "rollout_log_probs": torch.float32, + "rs_preflight_log_probs": torch.float32, "rollout_top_p_token_ids": torch.int32, "rollout_top_p_token_offsets": torch.int32, "teacher_log_probs": torch.float32, diff --git a/slime/ray/actor_group.py b/slime/ray/actor_group.py index 7662a0de69..8e20970a79 100644 --- a/slime/ray/actor_group.py +++ b/slime/ray/actor_group.py @@ -148,6 +148,24 @@ def async_train(self, rollout_id, rollout_data_ref, external_data=None): for actor in self._actor_handlers ] + def async_score_rs_candidates(self, rollout_id, rollout_data_ref): + """Return ObjectRefs for exact proximal-policy RS reports.""" + + return [actor.score_rs_candidates.remote(rollout_id, rollout_data_ref) for actor in self._actor_handlers] + + def async_take_rs_candidate_log_probs(self, rollout_id, selected_sample_indices): + """Return ObjectRefs containing only selected actor-local caches.""" + + return [ + actor.take_rs_candidate_log_probs.remote(rollout_id, selected_sample_indices) + for actor in self._actor_handlers + ] + + def async_discard_rs_candidate_log_probs(self, rollout_id): + """Discard one candidate cache on every actor rank.""" + + return [actor.discard_rs_candidate_log_probs.remote(rollout_id) for actor in self._actor_handlers] + def save_model(self, rollout_id, force_sync=False): """Save actor model""" ret = ray.get([actor.save_model.remote(rollout_id, force_sync=force_sync) for actor in self._actor_handlers]) diff --git a/slime/ray/rollout.py b/slime/ray/rollout.py index 246e828fb1..490c675774 100644 --- a/slime/ray/rollout.py +++ b/slime/ray/rollout.py @@ -1,3 +1,4 @@ +import copy import dataclasses import itertools import logging @@ -32,6 +33,21 @@ from slime.utils.health_monitor import RolloutHealthMonitor from slime.utils.http_utils import _wrap_ipv6, find_available_port, get_host_info, init_http_client from slime.utils.misc import Box, load_function +from slime.utils.rs_refill import ( + attach_proximal_log_probs, + fingerprint_rs_train_data, + merge_replacement_metrics, + merge_selected_log_prob_caches, + plan_topology_aligned_rs_refill, + select_accepted_groups, + snapshot_sample_masks, + validate_initial_policy_staleness, + validate_refill_rollout_ids, + validate_replacement_policy_version, + validate_rs_refill_target_batch_alignment, + validate_rs_train_data_fingerprint, + validate_sample_masks, +) from slime.utils.types import Sample from .rollout_validation import validate_server_group_gpu_indices @@ -406,6 +422,7 @@ def __init__(self, args, pg): runtime_env={"env_vars": add_default_ray_env_vars()}, ).remote() self.rollout_id = -1 + self._pending_rs_batches: dict[int, dict[str, Any]] = {} self._health_monitors = [] if not self.args.debug_train_only and self.args.use_fault_tolerance: @@ -490,6 +507,9 @@ def get_num_rollout_per_epoch(self): return len(self.data_source) // self.args.rollout_batch_size def generate(self, rollout_id): + if getattr(self.args, "rs_batch_refill", False): + return self._generate_rs_candidates(rollout_id) + start_time = time.time() self.rollout_id = rollout_id set_current_rollout_id(rollout_id) @@ -510,6 +530,441 @@ def generate(self, rollout_id): data = self._convert_samples_to_train_data(data) return self._split_train_data_by_dp(data) + def _call_rollout_for_group_count( + self, + rollout_id: int, + group_count: int, + *, + known_rollout_ids: set[Any] | None = None, + ): + # Evaluation shares this process-global hook context, so restore the + # pending batch ID before every reactive generation. + self.rollout_id = rollout_id + set_current_rollout_id(rollout_id) + # Custom rollout functions and sample hooks receive this namespace. Keep + # their nested mutations local to one candidate-generation call so they + # cannot change the topology or batching of later refill rounds. + call_args = copy.deepcopy(self.args) + call_args.rollout_batch_size = group_count + # A refill should request only the missing groups. The default rollout + # loop otherwise inherits the original full-batch sampling granularity. + call_args.over_sampling_batch_size = group_count + output = call_rollout_fn( + self.generate_rollout, + call_args, + rollout_id, + self.data_source, + evaluation=False, + ) + groups = output.samples + if len(groups) != group_count: + raise RuntimeError(f"RS refill requested {group_count} prompt groups but rollout returned {len(groups)}") + for group in groups: + if not isinstance(group, list) or len(group) != self.args.n_samples_per_prompt: + raise RuntimeError( + "RS batch refill requires the rollout function to return one list[Sample] per prompt " + f"with n_samples_per_prompt={self.args.n_samples_per_prompt}." + ) + if any(isinstance(sample, list) for sample in group): + raise RuntimeError("RS batch refill does not support fan-out/compact nested rollout samples yet.") + if any(sample.index is None or sample.group_index is None for sample in group): + raise RuntimeError("RS batch refill requires stable sample.index and sample.group_index values.") + candidate_rollout_ids = validate_refill_rollout_ids(groups, known_rollout_ids=known_rollout_ids) + return groups, output.metrics or {}, candidate_rollout_ids + + def _generate_rs_candidates(self, rollout_id: int): + if rollout_id in self._pending_rs_batches: + raise RuntimeError(f"RS candidate batch {rollout_id} already exists") + start_time = time.perf_counter() + self.rollout_id = rollout_id + set_current_rollout_id(rollout_id) + self.health_monitoring_resume() + if self.args.ci_test and self.args.use_fault_tolerance and rollout_id >= 2: + self._try_ci_fault_injection() + candidate_count = self.args.rollout_batch_size + groups, metrics, candidate_rollout_ids = self._call_rollout_for_group_count(rollout_id, candidate_count) + self._pending_rs_batches[rollout_id] = { + "accepted": [], + "unscored": groups, + "initial_candidate_count": candidate_count, + "round": 0, + "seen_sample_indices": set(), + "seen_group_indices": set(), + "seen_rollout_ids": candidate_rollout_ids, + "awaiting_log_prob_indices": None, + "awaiting_log_prob_bytes": None, + "proximal_log_probs_by_sample_index": {}, + "retained_logprob_cache_bytes": 0, + "accepted_mask_fingerprints": {}, + "metrics": dict(metrics), + "initial_generation_seconds": time.perf_counter() - start_time, + } + return rollout_id + + def prepare_rs_candidate_data(self, rollout_id: int): + pending = self._pending_rs_batches.get(rollout_id) + if pending is None: + raise RuntimeError(f"No pending RS candidate batch for rollout_id={rollout_id}") + if pending["awaiting_log_prob_indices"] is not None or pending["awaiting_log_prob_bytes"] is not None: + raise RuntimeError(f"RS candidate batch {rollout_id} has selected log probabilities awaiting collection") + groups = pending["unscored"] + if not groups: + raise RuntimeError(f"RS candidate batch {rollout_id} has no unscored groups") + + samples = list(itertools.chain.from_iterable(groups)) + data = self._convert_samples_to_train_data(samples, preflight=True) + # Preflight is forward-only, so score all candidates in one logical + # step regardless of the final optimizer global batch size. + return self._split_train_data_by_dp(data, global_batch_size=len(set(data["rollout_ids"]))) + + def apply_rs_candidate_reports(self, rollout_id: int, actor_report_refs, preflight_seconds: float): + pending = self._pending_rs_batches.get(rollout_id) + if pending is None: + raise RuntimeError(f"No pending RS candidate batch for rollout_id={rollout_id}") + if pending["awaiting_log_prob_indices"] is not None or pending["awaiting_log_prob_bytes"] is not None: + raise RuntimeError(f"RS candidate batch {rollout_id} already has an uncollected preflight result") + + report_wait_start = time.perf_counter() + actor_reports = ray.get(actor_report_refs, timeout=self.args.rs_refill_rpc_timeout_seconds) + preflight_seconds += time.perf_counter() - report_wait_start + reports = [report for worker_reports in actor_reports if worker_reports for report in worker_reports] + groups = pending["unscored"] + remaining_target = self.args.rollout_batch_size - len(pending["accepted"]) + selection = select_accepted_groups( + groups, + reports, + target_size=remaining_target, + known_sample_indices=pending["seen_sample_indices"], + known_group_indices=pending["seen_group_indices"], + ) + if pending["round"] == 0: + initial_staleness = validate_initial_policy_staleness(groups, reports) + else: + validate_replacement_policy_version(groups, reports) + + samples_by_index = {sample.index: sample for group in groups for sample in group} + cache_bytes_by_sample_index = {} + actor_cache_bytes = [] + float32_bytes = torch.finfo(torch.float32).bits // 8 + for worker_reports in actor_reports: + worker_cache_bytes = 0 + for report in worker_reports or []: + sample_index = report["sample_index"] + cache_bytes = report.get("candidate_cache_bytes") + if isinstance(cache_bytes, bool) or not isinstance(cache_bytes, int) or cache_bytes < 0: + raise RuntimeError( + "RS candidate cache byte reports must be non-negative integers: " + f"sample_index={sample_index}, value={cache_bytes!r}" + ) + expected_cache_bytes = samples_by_index[sample_index].response_length * float32_bytes + if cache_bytes != expected_cache_bytes: + raise RuntimeError( + "RS candidate cache byte report must exactly match the expected float32 payload: " + f"sample_index={sample_index}, reported={cache_bytes}, expected={expected_cache_bytes}" + ) + cache_bytes_by_sample_index[sample_index] = cache_bytes + worker_cache_bytes += cache_bytes + actor_cache_bytes.append(worker_cache_bytes) + + accepted_sample_indices = [sample.index for group in selection.accepted_groups for sample in group] + incoming_cache_bytes = sum(cache_bytes_by_sample_index[index] for index in accepted_sample_indices) + retained_cache_bytes = pending["retained_logprob_cache_bytes"] + if ( + isinstance(retained_cache_bytes, bool) + or not isinstance(retained_cache_bytes, int) + or retained_cache_bytes < 0 + ): + raise RuntimeError(f"Invalid retained RS proximal-logprob cache size: {retained_cache_bytes!r}") + required_cache_bytes = retained_cache_bytes + incoming_cache_bytes + cache_limit = self.args.rs_refill_max_candidate_cache_bytes + peak_actor_cache_bytes = max(actor_cache_bytes, default=0) + if peak_actor_cache_bytes > cache_limit: + raise RuntimeError( + "RS candidate proximal-logprob cache exceeds the per-actor limit according to actor reports: " + f"required={peak_actor_cache_bytes}, limit={cache_limit}. " + "Increase --rs-refill-max-candidate-cache-bytes or reduce the candidate batch/response length." + ) + if required_cache_bytes > cache_limit: + raise RuntimeError( + "RS selected proximal-logprob cache would exceed the RolloutManager retained-payload limit " + "before transfer: " + f"retained={retained_cache_bytes}, incoming={incoming_cache_bytes}, " + f"required={required_cache_bytes}, limit={cache_limit}. " + "Increase --rs-refill-max-candidate-cache-bytes or reduce the batch/response length." + ) + + pending["seen_sample_indices"].update(sample.index for group in groups for sample in group) + pending["seen_group_indices"].update(group[0].group_index for group in groups) + + for group in selection.accepted_groups: + pending["accepted_mask_fingerprints"].update(snapshot_sample_masks(group)) + pending["awaiting_log_prob_indices"] = accepted_sample_indices + pending["awaiting_log_prob_bytes"] = incoming_cache_bytes + pending["accepted"].extend(selection.accepted_groups) + if pending["round"] == 0: + pending["metrics"]["rollout/rs_refill/initial_policy_staleness"] = initial_staleness + + pending["metrics"].update( + { + "rollout/rs_refill/candidate_groups": pending["initial_candidate_count"], + "rollout/rs_refill/rejected_groups": pending["metrics"].get("rollout/rs_refill/rejected_groups", 0) + + len(selection.rejected_groups), + "rollout/rs_refill/surplus_groups": pending["metrics"].get("rollout/rs_refill/surplus_groups", 0) + + len(selection.surplus_groups), + } + ) + pending["metrics"]["rollout/rs_refill/scored_groups"] = pending["metrics"].get( + "rollout/rs_refill/scored_groups", 0 + ) + len(groups) + pending["metrics"]["rollout/rs_refill/scored_trainable_tokens"] = pending["metrics"].get( + "rollout/rs_refill/scored_trainable_tokens", 0 + ) + sum(int(report["valid_tokens"]) for report in reports) + aggregate_cache_bytes = sum(cache_bytes_by_sample_index.values()) + pending["metrics"]["rollout/rs_refill/aggregate_candidate_logprob_cache_bytes"] = ( + pending["metrics"].get("rollout/rs_refill/aggregate_candidate_logprob_cache_bytes", 0) + + aggregate_cache_bytes + ) + pending["metrics"]["rollout/rs_refill/peak_aggregate_candidate_logprob_cache_bytes"] = max( + pending["metrics"].get("rollout/rs_refill/peak_aggregate_candidate_logprob_cache_bytes", 0), + aggregate_cache_bytes, + ) + pending["metrics"]["rollout/rs_refill/peak_actor_candidate_logprob_cache_bytes"] = max( + pending["metrics"].get("rollout/rs_refill/peak_actor_candidate_logprob_cache_bytes", 0), + peak_actor_cache_bytes, + ) + pending["metrics"]["rollout/rs_refill/logprob_cache_limit_bytes"] = cache_limit + pending["metrics"]["rollout/rs_refill/preflight_seconds"] = ( + pending["metrics"].get("rollout/rs_refill/preflight_seconds", 0.0) + preflight_seconds + ) + + deficit = self.args.rollout_batch_size - len(pending["accepted"]) + pending["unscored"] = [] + return { + "complete": deficit == 0, + "exhausted": deficit > 0 and pending["round"] >= self.args.rs_refill_max_rounds, + "deficit": deficit, + "round": pending["round"], + "accepted_groups": len(pending["accepted"]), + "target_groups": self.args.rollout_batch_size, + "accepted_sample_indices": accepted_sample_indices, + } + + def generate_rs_replacement_candidates(self, rollout_id: int): + """Generate the next bounded replacement set after actor caches are released.""" + + pending = self._pending_rs_batches.get(rollout_id) + if pending is None: + raise RuntimeError(f"No pending RS candidate batch for rollout_id={rollout_id}") + if pending["awaiting_log_prob_indices"] is not None or pending["awaiting_log_prob_bytes"] is not None: + raise RuntimeError( + f"Cannot generate RS replacements for batch {rollout_id} before collecting selected caches" + ) + if pending["unscored"]: + raise RuntimeError(f"RS candidate batch {rollout_id} already has unscored replacement groups") + + deficit = self.args.rollout_batch_size - len(pending["accepted"]) + if deficit <= 0: + raise RuntimeError(f"RS candidate batch {rollout_id} is already complete") + if pending["round"] >= self.args.rs_refill_max_rounds: + raise RuntimeError( + "RS batch refill exhausted its retry budget before optimizer.step: " + f"accepted={len(pending['accepted'])}, target={self.args.rollout_batch_size}, " + f"rounds={pending['round']}, remaining={deficit}." + ) + + refill_start = time.perf_counter() + replacement_count = plan_topology_aligned_rs_refill( + self.args, + self.train_parallel_config, + deficit, + ) + refill_round = pending["round"] + 1 + pending["metrics"][f"rollout/rs_refill/round_{refill_round}/candidate_groups"] = replacement_count + replacement_groups, replacement_metrics, replacement_rollout_ids = self._call_rollout_for_group_count( + rollout_id, + replacement_count, + known_rollout_ids=pending["seen_rollout_ids"], + ) + pending["unscored"] = replacement_groups + pending["seen_rollout_ids"].update(replacement_rollout_ids) + pending["round"] += 1 + merge_replacement_metrics( + pending["metrics"], + replacement_metrics, + round_index=pending["round"], + ) + pending["metrics"]["rollout/rs_refill/generated_replacement_groups"] = ( + pending["metrics"].get("rollout/rs_refill/generated_replacement_groups", 0) + replacement_count + ) + pending["metrics"]["rollout/rs_refill/replacement_generation_seconds"] = pending["metrics"].get( + "rollout/rs_refill/replacement_generation_seconds", 0.0 + ) + (time.perf_counter() - refill_start) + return { + "round": pending["round"], + "candidate_groups": replacement_count, + } + + def store_rs_accepted_log_probs(self, rollout_id: int, actor_cache_refs): + """Store only caches selected by the group-atomic manager decision.""" + + pending = self._pending_rs_batches.get(rollout_id) + if pending is None: + raise RuntimeError(f"No pending RS candidate batch for rollout_id={rollout_id}") + expected = pending["awaiting_log_prob_indices"] + expected_bytes = pending["awaiting_log_prob_bytes"] + if expected is None or expected_bytes is None: + raise RuntimeError(f"RS candidate batch {rollout_id} has no selected log probabilities to collect") + + transfer_start = time.perf_counter() + worker_caches = ray.get(actor_cache_refs, timeout=self.args.rs_refill_rpc_timeout_seconds) + selected_cache = merge_selected_log_prob_caches(worker_caches, expected) + overlap = set(selected_cache) & set(pending["proximal_log_probs_by_sample_index"]) + if overlap: + raise RuntimeError(f"RS proximal logprob cache was already stored for sample indices {sorted(overlap)}") + + accepted_by_index = { + sample.index: sample for group in pending["accepted"] for sample in group if sample.index in selected_cache + } + transferred_bytes = 0 + for sample_index, proximal_log_probs in selected_cache.items(): + sample = accepted_by_index[sample_index] + if ( + not isinstance(proximal_log_probs, torch.Tensor) + or proximal_log_probs.device.type != "cpu" + or proximal_log_probs.dtype != torch.float32 + or proximal_log_probs.ndim != 1 + or proximal_log_probs.numel() != sample.response_length + or not proximal_log_probs.is_contiguous() + ): + raise RuntimeError( + "RS selected proximal-logprob cache must contain contiguous one-dimensional CPU float32 tensors " + "matching each response length: " + f"sample_index={sample_index}, type={type(proximal_log_probs).__name__}, " + f"device={getattr(proximal_log_probs, 'device', None)}, " + f"dtype={getattr(proximal_log_probs, 'dtype', None)}, " + f"shape={getattr(proximal_log_probs, 'shape', None)}, " + f"response_length={sample.response_length}" + ) + transferred_bytes += proximal_log_probs.numel() * proximal_log_probs.element_size() + if transferred_bytes != expected_bytes: + raise RuntimeError( + "RS selected proximal-logprob cache byte report does not match the transferred tensors: " + f"reported={expected_bytes}, actual={transferred_bytes}" + ) + + retained_cache_bytes = pending["retained_logprob_cache_bytes"] + required_cache_bytes = retained_cache_bytes + transferred_bytes + cache_limit = self.args.rs_refill_max_candidate_cache_bytes + if required_cache_bytes > cache_limit: + raise RuntimeError( + "RS selected proximal-logprob cache exceeds the RolloutManager retained-payload limit after transfer: " + f"retained={retained_cache_bytes}, incoming={transferred_bytes}, " + f"required={required_cache_bytes}, limit={cache_limit}." + ) + + pending["proximal_log_probs_by_sample_index"].update(selected_cache) + pending["retained_logprob_cache_bytes"] = required_cache_bytes + pending["metrics"]["rollout/rs_refill/selected_logprob_transfer_bytes"] = ( + pending["metrics"].get("rollout/rs_refill/selected_logprob_transfer_bytes", 0) + transferred_bytes + ) + pending["metrics"]["rollout/rs_refill/retained_logprob_cache_bytes"] = required_cache_bytes + pending["metrics"]["rollout/rs_refill/peak_retained_logprob_cache_bytes"] = max( + pending["metrics"].get("rollout/rs_refill/peak_retained_logprob_cache_bytes", 0), + required_cache_bytes, + ) + pending["metrics"]["rollout/rs_refill/selected_logprob_transfer_seconds"] = pending["metrics"].get( + "rollout/rs_refill/selected_logprob_transfer_seconds", 0.0 + ) + (time.perf_counter() - transfer_start) + pending["awaiting_log_prob_indices"] = None + pending["awaiting_log_prob_bytes"] = None + + def abort_rs_batch(self, rollout_id: int) -> bool: + """Discard transient manager state after a fatal coordination error.""" + + return self._pending_rs_batches.pop(rollout_id, None) is not None + + def finalize_rs_batch(self, rollout_id: int, coordinator_seconds: float): + pending = self._pending_rs_batches.get(rollout_id) + if pending is None: + raise RuntimeError(f"No pending RS candidate batch for rollout_id={rollout_id}") + if pending["awaiting_log_prob_indices"] is not None or pending["awaiting_log_prob_bytes"] is not None: + raise RuntimeError(f"Cannot finalize RS candidate batch {rollout_id} before collecting selected caches") + groups = pending["accepted"] + if len(groups) != self.args.rollout_batch_size: + raise RuntimeError( + f"Cannot finalize underfilled RS batch: accepted={len(groups)}, target={self.args.rollout_batch_size}" + ) + + samples = list(itertools.chain.from_iterable(groups)) + pending["metrics"]["rollout/rs_refill/rounds"] = pending["round"] + pending["metrics"]["rollout/rs_refill/accepted_groups"] = len(groups) + scored_groups = pending["metrics"]["rollout/rs_refill/scored_groups"] + rejected_groups = pending["metrics"].get("rollout/rs_refill/rejected_groups", 0) + pending["metrics"]["rollout/rs_refill/gate_acceptance_rate"] = ( + scored_groups - rejected_groups + ) / scored_groups + pending["metrics"]["rollout/rs_refill/selection_utilization"] = len(groups) / scored_groups + pending["metrics"]["rollout/rs_refill/coordinator_seconds"] = coordinator_seconds + effective_tokens = sum( + int(torch.as_tensor(sample.loss_mask).sum().item()) if sample.loss_mask is not None else 0 + for sample in samples + ) + pending["metrics"]["rollout/rs_refill/effective_trainable_tokens"] = effective_tokens + pending["metrics"]["rollout/rs_refill/initial_candidate_generation_seconds"] = pending[ + "initial_generation_seconds" + ] + refill_path_seconds = pending["initial_generation_seconds"] + coordinator_seconds + pending["metrics"]["rollout/rs_refill/refill_path_seconds"] = refill_path_seconds + pending["metrics"]["rollout/rs_refill/effective_tokens_per_refill_path_second"] = ( + effective_tokens / refill_path_seconds if refill_path_seconds > 0 else 0.0 + ) + data = self._convert_samples_to_train_data(samples) + observability_fingerprint = None + if ( + getattr(self.args, "custom_rollout_log_function_path", None) is not None + or self.args.save_debug_rollout_data is not None + ): + observability_fingerprint = fingerprint_rs_train_data( + data, + group_indices=[sample.group_index for sample in samples], + weight_versions=[sample.weight_versions for sample in samples], + ) + save_debug_rollout_data( + self.args.save_debug_rollout_data, + samples, + rollout_id=rollout_id, + evaluation=False, + ) + rollout_log_args = ( + copy.deepcopy(self.args) + if getattr(self.args, "custom_rollout_log_function_path", None) is not None + else self.args + ) + log_rollout_data( + rollout_id, + rollout_log_args, + samples, + pending["metrics"], + refill_path_seconds, + ) + if observability_fingerprint is not None: + data = self._convert_samples_to_train_data(samples) + validate_rs_train_data_fingerprint( + data, + observability_fingerprint, + group_indices=[sample.group_index for sample in samples], + weight_versions=[sample.weight_versions for sample in samples], + ) + validate_sample_masks(samples, pending["accepted_mask_fingerprints"]) + attach_proximal_log_probs( + data, + samples, + pending["proximal_log_probs_by_sample_index"], + ) + result = self._split_train_data_by_dp(data) + del self._pending_rs_batches[rollout_id] + return result + def eval(self, rollout_id): if self.args.debug_train_only: # if debug train only, we don't generate evaluation data @@ -632,14 +1087,22 @@ def _post_process_rewards(self, samples: list[Sample] | list[list[Sample]]): return raw_rewards, raw_rewards - def _convert_samples_to_train_data(self, samples: list[Sample] | list[list[Sample]]): + def _convert_samples_to_train_data( + self, + samples: list[Sample] | list[list[Sample]], + *, + preflight: bool = False, + ): """ Convert inference generated samples to training data. """ - if self.custom_convert_samples_to_train_data_func is not None: + if self.custom_convert_samples_to_train_data_func is not None and not preflight: return self.custom_convert_samples_to_train_data_func(self.args, samples) - raw_rewards, rewards = self._post_process_rewards(samples) + if preflight: + raw_rewards = rewards = [0.0] * len(samples) + else: + raw_rewards, rewards = self._post_process_rewards(samples) assert len(raw_rewards) == len(samples) assert len(rewards) == len(samples) @@ -665,6 +1128,8 @@ def _convert_samples_to_train_data(self, samples: list[Sample] | list[list[Sampl "sample_indices": [sample.index for sample in samples], "rollout_ids": rollout_ids, } + if preflight: + train_data["group_indices"] = [sample.group_index for sample in samples] # loss mask # TODO: compress the loss mask @@ -752,9 +1217,11 @@ def _convert_samples_to_train_data(self, samples: list[Sample] | list[list[Sampl return train_data def set_train_parallel_config(self, config: dict): + if getattr(self.args, "rs_batch_refill", False): + validate_rs_refill_target_batch_alignment(self.args, config) self.train_parallel_config = config - def _split_train_data_by_dp(self, data): + def _split_train_data_by_dp(self, data, global_batch_size: int | None = None): """Compute the DP/mbs schedule and package each rank's rollout_data into a Ray Box. The schedule itself is computed by :func:`build_dp_schedule` so it stays unit-testable without Ray/sglang. @@ -774,7 +1241,7 @@ def _split_train_data_by_dp(self, data): self.args, self.train_parallel_config, total_lengths, - global_batch_size=self.args.global_batch_size, + global_batch_size=global_batch_size or self.args.global_batch_size, rollout_indices=data["rollout_ids"], ) @@ -792,9 +1259,11 @@ def _split_train_data_by_dp(self, data): "loss_masks", "round_number", "sample_indices", + "group_indices", "rollout_ids", "rollout_mask_sums", "rollout_log_probs", + "rs_preflight_log_probs", "rollout_top_p_token_ids", "rollout_top_p_token_offsets", "rollout_routed_experts", diff --git a/slime/rollout/data_source.py b/slime/rollout/data_source.py index ca7171ecae..2b6f4b4c6a 100644 --- a/slime/rollout/data_source.py +++ b/slime/rollout/data_source.py @@ -90,17 +90,32 @@ def __init__(self, args): def get_samples(self, num_samples): # TODO further improve code if self.dataset is not None: - if self.sample_offset + num_samples <= len(self.dataset): - prompt_samples = self.dataset.samples[self.sample_offset : self.sample_offset + num_samples] - self.sample_offset += num_samples - else: - prompt_samples = self.dataset.samples[self.sample_offset :] - num_samples -= len(prompt_samples) + dataset_size = len(self.dataset) + if dataset_size == 0 and num_samples: + raise RuntimeError("Cannot draw rollout samples from an empty dataset") + if not 0 <= self.sample_offset <= dataset_size: + raise RuntimeError( + "Rollout data source has an invalid dataset offset: " + f"sample_offset={self.sample_offset}, dataset_size={dataset_size}" + ) + + prompt_samples = [] + remaining = num_samples + while remaining > 0: + available = dataset_size - self.sample_offset + take = min(remaining, available) + if take: + prompt_samples.extend(self.dataset.samples[self.sample_offset : self.sample_offset + take]) + self.sample_offset += take + remaining -= take + + if remaining == 0: + break + self.epoch_id += 1 if self.args.rollout_shuffle: self.dataset.shuffle(self.epoch_id) - prompt_samples += self.dataset.samples[:num_samples] - self.sample_offset = num_samples + self.sample_offset = 0 else: prompt_samples = [Sample() for _ in range(num_samples)] diff --git a/slime/rollout/sglang_rollout.py b/slime/rollout/sglang_rollout.py index 2cbcb89662..2dfedc4ab1 100644 --- a/slime/rollout/sglang_rollout.py +++ b/slime/rollout/sglang_rollout.py @@ -133,13 +133,19 @@ def reset(self) -> None: self.pendings = set() self.aborted = False - def submit_generate_tasks(self, samples: list[list[Sample]]) -> None: + def submit_generate_tasks( + self, + samples: list[list[Sample]], + *, + args: Namespace | None = None, + ) -> None: + call_args = self.args if args is None else args for group in samples: self.pendings.add( asyncio.create_task( # submit a group of samples as a single task. generate_and_rm_group( - self.args, + call_args, group, sampling_params=self.sampling_params.copy(), evaluation=False, @@ -408,7 +414,7 @@ async def generate_rollout_async( while state.remaining_batch_size < target_data_size: # get samples from the buffer and submit the generation requests. samples = data_source(args.over_sampling_batch_size) - state.submit_generate_tasks(samples) + state.submit_generate_tasks(samples, args=args) # wait for the generation to finish done, state.pendings = await asyncio.wait(state.pendings, return_when=asyncio.FIRST_COMPLETED) diff --git a/slime/utils/arguments.py b/slime/utils/arguments.py index 33381cf0a3..378565518f 100644 --- a/slime/utils/arguments.py +++ b/slime/utils/arguments.py @@ -2,6 +2,7 @@ import copy import json import logging +import math import os from typing import Any @@ -1081,6 +1082,36 @@ def add_algo_arguments(parser): default=None, help="Path to the custom TIS/RS function (e.g., examples/train_infer_mismatch_helper/mis.py:compute_mis_weights_with_cp).", ) + parser.add_argument( + "--rs-batch-refill", + action="store_true", + default=False, + help=( + "For disaggregated train_async.py, apply sequence rejection sampling before the optimizer " + "step and generate topology-aligned replacements until the effective batch is complete." + ), + ) + parser.add_argument( + "--rs-refill-max-rounds", + type=int, + default=2, + help="Maximum replacement rounds before failing the job without an optimizer step.", + ) + parser.add_argument( + "--rs-refill-rpc-timeout-seconds", + type=float, + default=1800.0, + help="Timeout for each post-generation exact-refill coordination or actor RPC before aborting the batch.", + ) + parser.add_argument( + "--rs-refill-max-candidate-cache-bytes", + type=int, + default=1 << 30, + help=( + "Maximum proximal-logprob tensor payload retained by one process: each actor rank per candidate " + "round and the RolloutManager across accepted refill rounds." + ), + ) parser.add_argument( "--custom-pg-loss-reducer-function-path", type=str, @@ -1903,6 +1934,8 @@ def slime_validate_args(args): args.rollout_external = args.rollout_external_engine_addrs is not None + if args.rs_batch_refill and args.rollout_external: + raise ValueError("--rs-batch-refill does not support external rollout engines yet.") if args.rollout_external and not args.debug_train_only: apply_external_engine_info_to_args(args, logger=logger) @@ -2010,14 +2043,225 @@ def slime_validate_args(args): if args.use_rollout_routing_replay: args.use_routing_replay = True + custom_config_keys = set() if args.custom_config_path: with open(args.custom_config_path) as f: data = yaml.safe_load(f) or {} + custom_config_keys = set(data) for k, v in data.items(): if hasattr(args, k): logger.info(f"Warning: Argument {k} is already set to {getattr(args, k)}, will override with {v}.") setattr(args, k, v) + if args.rs_batch_refill: + if args.train_backend != "megatron": + raise ValueError("--rs-batch-refill currently requires --train-backend megatron.") + if args.colocate: + raise ValueError("--rs-batch-refill is only supported by disaggregated train_async.py.") + if getattr(args, "offload", False): + raise ValueError("--rs-batch-refill does not support offload: true in --custom-config-path.") + external_config_keys = { + "rollout_external", + "rollout_external_engine_addrs", + "rollout_external_engine_infos", + } + if custom_config_keys & external_config_keys: + raise ValueError( + "--rs-batch-refill does not allow external rollout engine settings in --custom-config-path." + ) + if args.rollout_external or args.rollout_external_engine_addrs is not None: + raise ValueError("--rs-batch-refill does not support external rollout engines yet.") + if not args.rollout_global_dataset: + raise ValueError("--rs-batch-refill requires --rollout-global-dataset for stable sample and group IDs.") + if args.num_rollout is None or args.num_rollout <= 0: + raise ValueError("--rs-batch-refill requires a positive explicit --num-rollout.") + if args.update_weights_interval != 1: + raise ValueError("--rs-batch-refill requires --update-weights-interval 1.") + if not getattr(args, "finetune", False) or args.start_rollout_id != 0: + raise ValueError( + "--rs-batch-refill currently supports fresh or finetune runs starting at rollout 0; " + "checkpoint resume is not implemented yet." + ) + if getattr(args, "update_weight_start_version", 0) != 0: + raise ValueError("--rs-batch-refill requires --update-weight-start-version 0.") + if getattr(args, "ref_update_interval", None) is not None: + raise ValueError("--rs-batch-refill does not support --ref-update-interval yet.") + if args.num_steps_per_rollout not in {None, 1}: + raise ValueError("--rs-batch-refill requires --num-steps-per-rollout 1 when it is explicitly set.") + if args.rollout_batch_size <= 0 or args.n_samples_per_prompt <= 0: + raise ValueError("--rs-batch-refill requires positive rollout and per-prompt sample batch sizes.") + rollout_sample_count = args.rollout_batch_size * args.n_samples_per_prompt + if args.global_batch_size != rollout_sample_count: + raise ValueError( + "--rs-batch-refill requires exactly one optimizer step per rollout batch: " + "--global-batch-size must equal --rollout-batch-size * --n-samples-per-prompt " + f"({rollout_sample_count})." + ) + if args.megatron_config_path is not None: + raise ValueError("--rs-batch-refill does not support --megatron-config-path yet.") + if getattr(args, "custom_model_provider_path", None) is not None: + raise ValueError("--rs-batch-refill does not support --custom-model-provider-path yet.") + if getattr(args, "custom_megatron_init_path", None) is not None: + raise ValueError("--rs-batch-refill does not support --custom-megatron-init-path yet.") + if getattr(args, "custom_megatron_before_log_prob_hook_path", None) is not None: + raise ValueError( + "--rs-batch-refill does not support --custom-megatron-before-log-prob-hook-path because " + "preflight and training must observe the same model state." + ) + if getattr(args, "custom_megatron_before_train_step_hook_path", None) is not None: + raise ValueError( + "--rs-batch-refill does not support --custom-megatron-before-train-step-hook-path because " + "preflight and training must observe the same model state." + ) + if args.update_weight_mode != "full" or args.update_weight_transport != "nccl": + raise ValueError("--rs-batch-refill currently requires full weight updates over NCCL.") + if args.use_rollout_logprobs: + raise ValueError("--rs-batch-refill does not support --use-rollout-logprobs with proximal preflight.") + if getattr(args, "rollout_top_k", -1) != -1: + raise ValueError("--rs-batch-refill currently requires --rollout-top-k -1.") + rollout_temperature = getattr(args, "rollout_temperature", 1.0) + rollout_top_p = getattr(args, "rollout_top_p", 1.0) + if not math.isfinite(rollout_temperature) or rollout_temperature <= 0: + raise ValueError("--rs-batch-refill requires a finite positive --rollout-temperature.") + if not math.isfinite(rollout_top_p) or not 0 < rollout_top_p <= 1: + raise ValueError("--rs-batch-refill requires finite --rollout-top-p in (0, 1].") + if not args.use_tis: + raise ValueError("--rs-batch-refill requires --use-tis.") + if args.custom_tis_function_path is not None: + raise ValueError("--rs-batch-refill uses its built-in TIS/RS path; custom TIS functions are unsupported.") + if not getattr(args, "use_rs", False): + raise ValueError("--rs-batch-refill requires use_rs: true in --custom-config-path.") + if getattr(args, "rs_level", None) not in {"sequence", "geometric"}: + raise ValueError("--rs-batch-refill requires rs_level: sequence or geometric.") + if getattr(args, "tis_mode", None) not in {"truncate", "clip"}: + raise ValueError("--rs-batch-refill requires tis_mode: truncate or clip.") + if getattr(args, "tis_level", None) not in {"token", "sequence", "geometric"}: + raise ValueError("--rs-batch-refill requires tis_level: token, sequence, or geometric.") + if getattr(args, "tis_batch_normalize", False): + raise ValueError("--rs-batch-refill does not support tis_batch_normalize: true yet.") + if getattr(args, "attention_dropout", None) != 0.0 or getattr(args, "hidden_dropout", None) != 0.0: + raise ValueError( + "--rs-batch-refill requires --attention-dropout 0 and --hidden-dropout 0 so cached " + "preflight probabilities remain valid during training." + ) + fp8 = getattr(args, "fp8", None) + if fp8 is not None and fp8 is not False: + raise ValueError( + "--rs-batch-refill currently requires non-FP8 Megatron training because FP8 preflight " + "may update quantizer state or depend on candidate microbatch packing." + ) + if getattr(args, "fp8_param_gather", False): + raise ValueError("--rs-batch-refill does not support --fp8-param-gather yet.") + fp4 = getattr(args, "fp4", None) + if fp4 is not None and fp4 is not False: + raise ValueError( + "--rs-batch-refill currently requires non-FP4 Megatron training because FP4 preflight " + "may update quantizer state or depend on candidate microbatch packing." + ) + if getattr(args, "te_precision_config_file", None) is not None: + raise ValueError("--rs-batch-refill does not support a custom Transformer Engine precision config yet.") + if ( + getattr(args, "kitchen_config_file", None) is not None + or getattr(args, "kitchen_recipe_number", None) is not None + ): + raise ValueError("--rs-batch-refill does not support Megatron Kitchen quantization yet.") + deepgemm_forward_settings = ( + getattr(args, "megatron_deepgemm_forward_layers", None), + getattr(args, "megatron_deepgemm_forward_modules", None), + getattr(args, "megatron_deepgemm_moe_forward_layers", None), + getattr(args, "megatron_deepgemm_moe_forward_modules", None), + ) + if any(setting is not None for setting in deepgemm_forward_settings): + raise ValueError("--rs-batch-refill does not support DeepGEMM FP8 forward overrides yet.") + if getattr(args, "moe_input_jitter_eps", None) not in {None, 0}: + raise ValueError("--rs-batch-refill requires --moe-input-jitter-eps 0 or unset.") + if getattr(args, "moe_router_force_load_balancing", False): + raise ValueError("--rs-batch-refill does not support --moe-router-force-load-balancing.") + if getattr(args, "moe_expert_capacity_factor", None) is not None: + raise ValueError("--rs-batch-refill requires --moe-expert-capacity-factor to be unset.") + moe_load_balancing = getattr(args, "moe_router_load_balancing_type", "aux_loss") + if moe_load_balancing == "sinkhorn" or ( + isinstance(moe_load_balancing, (list, tuple)) and "sinkhorn" in moe_load_balancing + ): + raise ValueError("--rs-batch-refill does not support Sinkhorn MoE routing.") + if args.context_parallel_size != 1: + raise ValueError("--rs-batch-refill currently requires --context-parallel-size 1.") + if not args.use_dynamic_batch_size or args.max_tokens_per_gpu is None: + raise ValueError("--rs-batch-refill requires dynamic batching and --max-tokens-per-gpu.") + if args.log_probs_max_tokens_per_gpu is None: + args.log_probs_max_tokens_per_gpu = args.max_tokens_per_gpu + if getattr(args, "use_rollout_entropy", False): + raise ValueError("--rs-batch-refill does not support --use-rollout-entropy yet.") + if not args.compute_advantages_and_returns: + raise ValueError("--rs-batch-refill requires advantage and return computation.") + if getattr(args, "custom_advantage_function_path", None) is not None: + raise ValueError("--rs-batch-refill does not support --custom-advantage-function-path yet.") + if getattr(args, "loss_type", "policy_loss") != "policy_loss": + raise ValueError("--rs-batch-refill currently requires --loss-type policy_loss.") + if getattr(args, "custom_pg_loss_reducer_function_path", None) is not None: + raise ValueError("--rs-batch-refill does not support a custom policy-loss reducer yet.") + if getattr(args, "use_opsm", False) or args.use_opd: + raise ValueError("--rs-batch-refill does not support OPSM or on-policy distillation yet.") + if args.advantage_estimator == "ppo" or args.use_critic: + raise ValueError("--rs-batch-refill does not support critic training yet.") + if args.offload_train or args.release_train or args.offload_rollout: + raise ValueError("--rs-batch-refill requires resident actor and rollout engines.") + if args.keep_old_actor: + raise ValueError("--rs-batch-refill does not support --keep-old-actor.") + if args.partial_rollout: + raise ValueError("--rs-batch-refill does not support --partial-rollout yet.") + if getattr(args, "dynamic_sampling_filter_path", None) is not None: + raise ValueError("--rs-batch-refill does not support --dynamic-sampling-filter-path yet.") + if args.use_rollout_routing_replay or args.use_routing_replay: + raise ValueError("--rs-batch-refill does not support routing replay yet.") + if args.rollout_function_path == "slime.rollout.fully_async_rollout.generate_rollout_fully_async": + raise ValueError("--rs-batch-refill does not support the persistent fully-async rollout queue yet.") + if args.custom_convert_samples_to_train_data_path is not None: + raise ValueError("--rs-batch-refill does not support a custom train-data converter yet.") + if args.custom_reward_post_process_path is not None: + raise ValueError("--rs-batch-refill does not support a custom reward postprocessor yet.") + if args.rollout_data_postprocess_path is not None: + raise ValueError("--rs-batch-refill does not support a rollout-data postprocessor yet.") + if args.load_debug_rollout_data is not None or args.debug_train_only or args.debug_rollout_only: + raise ValueError("--rs-batch-refill requires both rollout and training backends.") + if ( + isinstance(args.rs_refill_max_rounds, bool) + or not isinstance(args.rs_refill_max_rounds, int) + or args.rs_refill_max_rounds < 0 + ): + raise ValueError("--rs-refill-max-rounds must be a non-negative integer.") + if not math.isfinite(args.rs_refill_rpc_timeout_seconds) or args.rs_refill_rpc_timeout_seconds <= 0: + raise ValueError("--rs-refill-rpc-timeout-seconds must be finite and positive.") + if ( + isinstance(args.rs_refill_max_candidate_cache_bytes, bool) + or not isinstance(args.rs_refill_max_candidate_cache_bytes, int) + or args.rs_refill_max_candidate_cache_bytes <= 0 + ): + raise ValueError("--rs-refill-max-candidate-cache-bytes must be a positive integer.") + + tis_upper_bound = getattr(args, "tis_upper_bound", None) + tis_lower_bound = getattr(args, "tis_lower_bound", None) + rs_lower_bound = getattr(args, "rs_lower_bound", None) + rs_upper_bound = getattr(args, "rs_upper_bound", None) + if tis_upper_bound is None or not math.isfinite(tis_upper_bound) or tis_upper_bound <= 0: + raise ValueError("--rs-batch-refill requires a finite positive tis_upper_bound.") + if tis_lower_bound is not None and (not math.isfinite(tis_lower_bound) or tis_lower_bound < 0): + raise ValueError("tis_lower_bound must be finite and non-negative for --rs-batch-refill.") + effective_tis_lower_bound = tis_lower_bound if tis_lower_bound is not None else 1.0 / tis_upper_bound + if args.tis_mode == "clip" and not effective_tis_lower_bound < tis_upper_bound: + raise ValueError("--rs-batch-refill requires TIS clip bounds with 0 <= lower < upper.") + effective_rs_lower_bound = rs_lower_bound if rs_lower_bound is not None else effective_tis_lower_bound + effective_rs_upper_bound = rs_upper_bound if rs_upper_bound is not None else tis_upper_bound + if ( + not math.isfinite(effective_rs_lower_bound) + or not math.isfinite(effective_rs_upper_bound) + or not 0 <= effective_rs_lower_bound < effective_rs_upper_bound + ): + raise ValueError("--rs-batch-refill requires finite RS bounds with 0 <= lower < upper.") + rs_veto_threshold = getattr(args, "rs_veto_threshold", None) + if rs_veto_threshold is not None and (not math.isfinite(rs_veto_threshold) or rs_veto_threshold < 0): + raise ValueError("rs_veto_threshold must be finite and non-negative for --rs-batch-refill.") + if args.eval_max_context_len is None: logger.info( f"args.eval_max_context_len is not set. Use args.rollout_max_context_len {args.rollout_max_context_len} as default value." diff --git a/slime/utils/rs_refill.py b/slime/utils/rs_refill.py new file mode 100644 index 0000000000..0e584d8071 --- /dev/null +++ b/slime/utils/rs_refill.py @@ -0,0 +1,879 @@ +"""Internal coordination and math helpers for bounded, group-atomic RS batch refill.""" + +from __future__ import annotations + +import hashlib +import logging +import math +import operator +import struct +from collections.abc import Callable, Mapping +from dataclasses import dataclass +from typing import Any + +import numpy as np +import torch + +logger = logging.getLogger(__name__) + + +@dataclass(frozen=True) +class RefillSelection: + """Result of applying per-sample RS reports to prompt groups.""" + + accepted_groups: list[list[Any]] + rejected_groups: list[list[Any]] + surplus_groups: list[list[Any]] + target_size: int + + @property + def deficit(self) -> int: + return max(0, self.target_size - len(self.accepted_groups)) + + +def _require_integer(value: Any, name: str, *, positive: bool = False) -> int: + if isinstance(value, bool): + raise ValueError(f"{name} must be an integer, not a boolean") + try: + result = operator.index(value) + except TypeError as error: + raise ValueError(f"{name} must be an integer, got {value!r}") from error + if positive and result <= 0: + raise ValueError(f"{name} must be positive, got {result}") + return result + + +def _resolve_rs_admission_bounds(args) -> tuple[Any, Any]: + tis_lower_bound = args.tis_lower_bound if args.tis_lower_bound is not None else 1.0 / args.tis_upper_bound + lower_bound = args.rs_lower_bound if args.rs_lower_bound is not None else tis_lower_bound + upper_bound = args.rs_upper_bound if args.rs_upper_bound is not None else args.tis_upper_bound + return lower_bound, upper_bound + + +def _resolve_rs_admission_log_bounds(args) -> tuple[float, float]: + lower_bound, upper_bound = _resolve_rs_admission_bounds(args) + lower_log_bound = -math.inf if lower_bound == 0 else math.log(lower_bound) + return lower_log_bound, math.log(upper_bound) + + +def get_rs_refill_candidate_group_multiple(args, train_parallel_config: Mapping[str, Any]) -> int: + """Return the prompt-group quantum that makes every preflight schedulable.""" + + try: + dp_size = _require_integer(train_parallel_config["dp_size"], "dp_size", positive=True) + vpp_size = _require_integer(train_parallel_config["vpp_size"], "vpp_size", positive=True) + microbatch_group_size = _require_integer( + train_parallel_config["microbatch_group_size_per_vp_stage"], + "microbatch_group_size_per_vp_stage", + positive=True, + ) + samples_per_group = _require_integer(args.n_samples_per_prompt, "n_samples_per_prompt", positive=True) + except KeyError as error: + raise ValueError("Invalid training topology for RS refill planning") from error + + aligned_samples = dp_size * (microbatch_group_size if vpp_size > 1 else 1) + return aligned_samples // math.gcd(aligned_samples, samples_per_group) + + +def validate_rs_refill_target_batch_alignment(args, train_parallel_config: Mapping[str, Any]) -> None: + """Fail before rollout when the final effective batch cannot always be scheduled.""" + + candidate_group_multiple = get_rs_refill_candidate_group_multiple(args, train_parallel_config) + if args.rollout_batch_size % candidate_group_multiple == 0: + return + + dp_size = int(train_parallel_config["dp_size"]) + vpp_size = int(train_parallel_config["vpp_size"]) + microbatch_group_size = int(train_parallel_config["microbatch_group_size_per_vp_stage"]) + aligned_samples = dp_size * (microbatch_group_size if vpp_size > 1 else 1) + target_samples = args.rollout_batch_size * args.n_samples_per_prompt + raise ValueError( + "--rs-batch-refill requires the final effective sample count to be divisible by the DP/VPP " + "microbatch alignment so every possible dynamic packing remains schedulable: " + f"rollout_batch_size * n_samples_per_prompt = {target_samples}, alignment = {aligned_samples}. " + f"Choose a rollout_batch_size divisible by {candidate_group_multiple}." + ) + + +def plan_topology_aligned_rs_refill( + args, + train_parallel_config: Mapping[str, Any], + required_successes: int, +) -> int: + """Round an exact deficit up to the smallest schedulable group count.""" + + if not isinstance(required_successes, int) or isinstance(required_successes, bool) or required_successes <= 0: + raise ValueError("required_successes must be a positive integer") + candidate_multiple = get_rs_refill_candidate_group_multiple(args, train_parallel_config) + return ((required_successes + candidate_multiple - 1) // candidate_multiple) * candidate_multiple + + +def compute_sequence_rs_masks( + args, + *, + train_log_probs: list[torch.Tensor], + rollout_log_probs: list[torch.Tensor], + loss_masks: list[torch.Tensor], +) -> list[torch.Tensor]: + """Compute the sequence-local RS admission mask used before and during training.""" + + if not (len(train_log_probs) == len(rollout_log_probs) == len(loss_masks)): + raise ValueError( + "RS admission inputs must have the same sample count: " + f"train={len(train_log_probs)}, rollout={len(rollout_log_probs)}, masks={len(loss_masks)}" + ) + if args.rs_level not in {"sequence", "geometric"}: + raise ValueError(f"RS batch refill does not support rs_level={args.rs_level!r}") + + lower_log_bound, upper_log_bound = _resolve_rs_admission_log_bounds(args) + modified_masks = [] + for sample_index, (train, rollout, loss_mask) in enumerate( + zip(train_log_probs, rollout_log_probs, loss_masks, strict=True) + ): + rollout = torch.as_tensor(rollout, device=train.device) + mask = torch.as_tensor(loss_mask, device=train.device).float() + if train.shape != rollout.shape or train.shape != mask.shape: + raise ValueError( + "RS admission sample shapes must match: " + f"sample={sample_index}, train={tuple(train.shape)}, rollout={tuple(rollout.shape)}, " + f"mask={tuple(mask.shape)}" + ) + + raw_log_ratio = train - rollout + finite_log_ratio = torch.isfinite(raw_log_ratio).all() + safe_log_ratio = torch.where(torch.isfinite(raw_log_ratio), raw_log_ratio, 0.0) + sequence_log_ratio = (safe_log_ratio * mask).sum() + if args.rs_level == "geometric": + sequence_log_ratio = sequence_log_ratio / torch.clamp_min(mask.sum(), 1) + accepted = finite_log_ratio & (sequence_log_ratio >= lower_log_bound) & (sequence_log_ratio <= upper_log_bound) + if args.rs_veto_threshold is not None: + veto_log_threshold = torch.log( + torch.tensor(args.rs_veto_threshold, device=raw_log_ratio.device, dtype=torch.float32) + ) + accepted = accepted & ~((raw_log_ratio < veto_log_threshold) & mask.bool()).any() + modified_masks.append((mask * accepted).detach()) + + return modified_masks + + +def apply_rs_refill_tis( + args, + *, + pg_loss: torch.Tensor, + train_log_probs: list[torch.Tensor], + rollout_log_probs: list[torch.Tensor], + loss_masks: list[torch.Tensor], + **_: Any, +) -> tuple[torch.Tensor, list[torch.Tensor], dict[str, torch.Tensor]]: + """Apply bounded TIS weights and exactly the same RS rule as preflight.""" + + if not (len(train_log_probs) == len(rollout_log_probs) == len(loss_masks)): + raise ValueError( + "RS refill TIS inputs must have the same sample count: " + f"train={len(train_log_probs)}, rollout={len(rollout_log_probs)}, masks={len(loss_masks)}" + ) + if not train_log_probs: + raise ValueError("RS refill TIS requires at least one sample") + if args.tis_batch_normalize: + raise ValueError("RS refill does not support DP-local TIS batch normalization") + + metrics: dict[str, list[torch.Tensor]] = {} + + def append_metric(key: str, value: torch.Tensor) -> None: + metrics.setdefault(key, []).append(value.clone().detach()) + + def masked_sum(value: torch.Tensor, mask: torch.Tensor, *, expand: bool = False) -> torch.Tensor: + result = (value * mask).sum() + return result.expand_as(value) if expand else result + + def masked_mean(value: torch.Tensor, mask: torch.Tensor, *, expand: bool = False) -> torch.Tensor: + result = masked_sum(value, mask) / torch.clamp_min(mask.sum(), 1) + return result.expand_as(value) if expand else result + + def aggregate_log_ratio(raw_log_ratio: torch.Tensor, mask: torch.Tensor, level: str) -> torch.Tensor: + if level == "token": + return raw_log_ratio + if level == "sequence": + return masked_sum(raw_log_ratio, mask, expand=True) + if level == "geometric": + return masked_mean(raw_log_ratio, mask, expand=True) + raise ValueError(f"RS refill TIS does not support tis_level={level!r}") + + tis_lower_bound = args.tis_lower_bound if args.tis_lower_bound is not None else 1.0 / args.tis_upper_bound + all_weights = [] + normalized_train_log_probs = [] + normalized_rollout_log_probs = [] + normalized_loss_masks = [] + for sample_index, (train, rollout, loss_mask) in enumerate( + zip(train_log_probs, rollout_log_probs, loss_masks, strict=True) + ): + rollout = torch.as_tensor(rollout, device=train.device) + mask = torch.as_tensor(loss_mask, device=train.device).float() + if train.shape != rollout.shape or train.shape != mask.shape: + raise ValueError( + "RS refill TIS sample shapes must match: " + f"sample={sample_index}, train={tuple(train.shape)}, rollout={tuple(rollout.shape)}, " + f"mask={tuple(mask.shape)}" + ) + + raw_log_ratio = train - rollout + mean_train_log_prob = masked_mean(train, mask, expand=True) + mean_rollout_log_prob = masked_mean(rollout, mask, expand=True) + training_log_ppl = -mean_train_log_prob + rollout_log_ppl = -mean_rollout_log_prob + log_ppl_diff = mean_rollout_log_prob - mean_train_log_prob + append_metric("mis_training_log_ppl", training_log_ppl) + append_metric("mis_training_ppl", torch.exp(training_log_ppl)) + append_metric("mis_rollout_log_ppl", rollout_log_ppl) + append_metric("mis_rollout_ppl", torch.exp(rollout_log_ppl)) + append_metric("mis_kl", rollout - train) + append_metric("mis_k3_kl", torch.exp(raw_log_ratio) - raw_log_ratio - 1) + append_metric("mis_log_ppl_diff", log_ppl_diff) + append_metric("mis_log_ppl_abs_diff", log_ppl_diff.abs()) + append_metric("mis_ppl_ratio", torch.exp(log_ppl_diff)) + safe_raw_log_ratio = torch.clamp(raw_log_ratio, min=-20.0, max=20.0) + chi2_token = masked_mean(torch.exp(safe_raw_log_ratio).square(), mask) - 1.0 + append_metric("mis_chi2_token", chi2_token.expand_as(train)) + sequence_log_ratio = torch.clamp(masked_sum(raw_log_ratio, mask, expand=True), min=-20.0, max=20.0) + append_metric("mis_chi2_seq", torch.exp(2.0 * sequence_log_ratio) - 1.0) + + log_ratio = aggregate_log_ratio(raw_log_ratio, mask, args.tis_level) + weights = torch.exp(torch.clamp(log_ratio, min=-20.0, max=20.0)) + append_metric("mis_tis_weight_before_bound", weights) + if args.tis_mode == "truncate": + append_metric("mis_tis_truncate_fraction", (weights > args.tis_upper_bound).int()) + weights = weights.clamp(0, args.tis_upper_bound) * mask + elif args.tis_mode == "clip": + append_metric("mis_tis_clip_fraction_low", (weights < tis_lower_bound).int()) + append_metric("mis_tis_clip_fraction_high", (weights > args.tis_upper_bound).int()) + weights = weights.clamp(tis_lower_bound, args.tis_upper_bound) * mask + else: + raise ValueError(f"RS refill TIS does not support tis_mode={args.tis_mode!r}") + append_metric("mis_tis_weight_after_bound", weights) + append_metric("mis_is_ratio_mean_after_tis_rs", weights) + + all_weights.append(weights.detach()) + normalized_train_log_probs.append(train) + normalized_rollout_log_probs.append(rollout) + normalized_loss_masks.append(mask) + + modified_masks = compute_sequence_rs_masks( + args, + train_log_probs=normalized_train_log_probs, + rollout_log_probs=normalized_rollout_log_probs, + loss_masks=normalized_loss_masks, + ) + + lower_log_bound, upper_log_bound = _resolve_rs_admission_log_bounds(args) + for train, rollout, mask in zip( + normalized_train_log_probs, normalized_rollout_log_probs, normalized_loss_masks, strict=True + ): + raw_log_ratio = train - rollout + rs_log_ratio = aggregate_log_ratio(raw_log_ratio, mask, args.rs_level) + append_metric("mis_rs_mask_fraction_low", (rs_log_ratio < lower_log_bound).int()) + append_metric("mis_rs_mask_fraction_high", (rs_log_ratio > upper_log_bound).int()) + if args.rs_veto_threshold is not None: + veto_log_threshold = torch.log( + torch.tensor(args.rs_veto_threshold, device=raw_log_ratio.device, dtype=torch.float32) + ) + catastrophic_tokens = (raw_log_ratio < veto_log_threshold) & mask.bool() + append_metric("mis_rs_catastrophic_token_fraction", catastrophic_tokens.int()) + append_metric("mis_rs_catastrophic_seq_fraction", catastrophic_tokens.any().int().expand_as(mask)) + + for weight, mask in zip(all_weights, normalized_loss_masks, strict=True): + valid = mask.bool() + mean = masked_mean(weight, mask, expand=True) + minimum = weight[valid].min() if valid.any() else weight.new_tensor(0.0) + maximum = weight[valid].max() if valid.any() else weight.new_tensor(0.0) + append_metric("mis_is_ratio_mean_final", mean) + append_metric("mis_is_ratio_min_final", minimum.expand_as(weight)) + append_metric("mis_is_ratio_max_final", maximum.expand_as(weight)) + + flat_weights = torch.cat(all_weights, dim=0) + if flat_weights.shape != pg_loss.shape: + raise ValueError( + "RS refill TIS weights must match the policy-gradient loss shape: " + f"weights={tuple(flat_weights.shape)}, pg_loss={tuple(pg_loss.shape)}" + ) + flat_metrics = {key: torch.cat(values, dim=0) for key, values in metrics.items()} + return pg_loss * flat_weights, modified_masks, flat_metrics + + +def merge_replacement_metrics( + destination: dict[str, Any], + source: dict[str, Any], + *, + round_index: int, +) -> None: + """Merge counters with known semantics and namespace every other metric.""" + + for key, value in source.items(): + if key.startswith("rollout/dynamic_filter/drop_"): + destination[key] = destination.get(key, 0) + value + else: + destination[f"rollout/rs_refill/replacement_round_{round_index}/{key}"] = value + + +def validate_refill_rollout_ids( + groups: list[list[Any]], + *, + known_rollout_ids: set[Any] | None = None, +) -> set[Any]: + """Require explicit effective rollout IDs to be unique across refill rounds.""" + + seen = set(known_rollout_ids or ()) + candidate_ids = set() + for group in groups: + for sample in group: + rollout_id = sample.rollout_id + if rollout_id is None: + continue + try: + duplicate = rollout_id in seen + except TypeError as error: + raise ValueError(f"RS refill rollout_id must be hashable, got {rollout_id!r}") from error + if duplicate: + raise ValueError( + "RS batch refill requires one unique effective rollout_id per training sample; " + f"duplicate={rollout_id!r}. Compact/fan-out rollouts are not supported yet." + ) + seen.add(rollout_id) + candidate_ids.add(rollout_id) + return candidate_ids + + +def merge_selected_log_prob_caches( + worker_caches: list[dict[int, Any] | None], + expected_sample_indices: list[int], +) -> dict[int, Any]: + """Merge actor-local caches and require exactly the selected samples.""" + + expected_indices = [_require_integer(index, "selected sample index") for index in expected_sample_indices] + expected = set(expected_indices) + if len(expected) != len(expected_indices): + raise ValueError("selected RS sample indices must be unique") + + merged: dict[int, Any] = {} + for worker_cache in worker_caches: + if not worker_cache: + continue + for sample_index, log_probs in worker_cache.items(): + sample_index = _require_integer(sample_index, "proximal logprob cache sample index") + if sample_index in merged: + raise ValueError(f"duplicate proximal logprob cache for sample_index={sample_index}") + merged[sample_index] = log_probs + + actual = set(merged) + if actual != expected: + raise ValueError( + "Selected RS proximal logprob cache is incomplete: " + f"missing={sorted(expected - actual)}, extra={sorted(actual - expected)}" + ) + return merged + + +def attach_proximal_log_probs( + train_data: dict[str, Any], + samples: list[Any], + log_probs_by_sample_index: dict[int, Any], +) -> None: + """Attach an in-memory preflight cache to the final batch exactly once.""" + + if "rs_preflight_log_probs" in train_data: + raise ValueError("RS preflight log probabilities are already attached to the final batch") + sample_indices = [_require_integer(sample.index, "RS-refilled sample index") for sample in samples] + if len(sample_indices) != len(set(sample_indices)): + raise ValueError("RS-refilled samples must have unique sample indices") + + expected = set(sample_indices) + cached = set(log_probs_by_sample_index) + if cached != expected: + raise ValueError( + "RS proximal logprob cache does not match the final batch: " + f"missing={sorted(expected - cached)}, extra={sorted(cached - expected)}" + ) + train_data["rs_preflight_log_probs"] = [log_probs_by_sample_index[index] for index in sample_indices] + + +def snapshot_sample_masks(samples: list[Any]) -> dict[int, tuple[int, bytes]]: + """Snapshot response lengths and fixed-size mask digests used by the preflight gate.""" + + fingerprints: dict[int, tuple[int, bytes]] = {} + for sample in samples: + sample_index = _require_integer(sample.index, "RS-refilled sample index") + if sample_index in fingerprints: + raise ValueError(f"duplicate sample index in RS mask snapshot: {sample_index}") + if sample.loss_mask is None: + raise ValueError(f"sample_index={sample_index} has no loss mask after RS preflight") + response_length = _require_integer(sample.response_length, "RS-refilled response length") + mask = torch.as_tensor(sample.loss_mask) + if mask.ndim != 1 or mask.numel() != response_length: + raise ValueError( + "RS-refilled loss mask must be one-dimensional and match response_length: " + f"sample_index={sample_index}, mask_shape={tuple(mask.shape)}, response_length={response_length}" + ) + if not torch.logical_or(mask == 0, mask == 1).all().item(): + raise ValueError(f"sample_index={sample_index} has a non-binary loss mask after RS preflight") + mask_bytes = mask.to(dtype=torch.uint8, device="cpu").contiguous().numpy().tobytes() + fingerprints[sample_index] = (response_length, hashlib.sha256(mask_bytes).digest()) + return fingerprints + + +def validate_sample_masks( + samples: list[Any], + expected: dict[int, tuple[int, bytes]], +) -> None: + """Reject post-gate mutations that would change the effective batch.""" + + actual = snapshot_sample_masks(samples) + if actual != expected: + expected_ids = set(expected) + actual_ids = set(actual) + changed = sorted(index for index in expected_ids & actual_ids if expected[index] != actual[index]) + raise RuntimeError( + "RS-refilled sample masks changed after proximal preflight: " + f"missing={sorted(expected_ids - actual_ids)}, extra={sorted(actual_ids - expected_ids)}, " + f"changed={changed}" + ) + + +def _rs_numpy_dtype_descriptor(dtype: np.dtype) -> dict[str, Any]: + subdtype = None + if dtype.subdtype is not None: + base, shape = dtype.subdtype + subdtype = { + "base": _rs_numpy_dtype_descriptor(base), + "shape": list(shape), + } + fields = None + if dtype.names is not None: + fields = [] + for name in dtype.names: + field = dtype.fields[name] + fields.append( + { + "name": name, + "dtype": _rs_numpy_dtype_descriptor(field[0]), + "offset": field[1], + "title": field[2] if len(field) > 2 else None, + } + ) + return { + "str": dtype.str, + "itemsize": dtype.itemsize, + "alignment": dtype.alignment, + "isalignedstruct": dtype.isalignedstruct, + "subdtype": subdtype, + "fields": fields, + "metadata": dict(dtype.metadata) if dtype.metadata is not None else None, + } + + +def _update_rs_value_fingerprint( + digest: Any, + value: Any, + active_containers: set[int], +) -> None: + """Stream a canonical train-data value into one digest.""" + + def add(tag: bytes, payload: Any = b"") -> None: + view = memoryview(payload) + if not view.contiguous: + raise TypeError("RS train-data fingerprint requires contiguous byte payloads") + byte_view = view.cast("B") + digest.update(len(tag).to_bytes(2, "big")) + digest.update(tag) + digest.update(byte_view.nbytes.to_bytes(8, "big")) + digest.update(byte_view) + + if value is None: + add(b"none") + elif isinstance(value, bool): + add(b"bool", bytes([value])) + elif isinstance(value, int): + add(b"int", str(value).encode("ascii")) + elif isinstance(value, float): + add(b"float", struct.pack(">d", value)) + elif isinstance(value, str): + add(b"str", value.encode("utf-8")) + elif isinstance(value, bytes): + add(b"bytes", value) + elif isinstance(value, torch.Tensor): + if value.device.type == "meta": + raise TypeError("RS train-data fingerprint does not support meta tensors") + add(b"tensor-layout", str(value.layout).encode("ascii")) + tensor = value.detach() + if tensor.layout != torch.strided: + tensor = tensor.to_dense() + tensor = tensor.to(device="cpu").resolve_conj().resolve_neg().contiguous() + add(b"tensor-dtype", str(tensor.dtype).encode("ascii")) + add(b"tensor-shape", ",".join(str(size) for size in tensor.shape).encode("ascii")) + add(b"tensor-data", tensor.reshape(-1).view(torch.uint8).numpy()) + elif isinstance(value, np.generic): + scalar = np.asarray(value) + if scalar.dtype.hasobject: + raise TypeError("RS train-data fingerprint does not support object-dtype scalars") + add(b"numpy-scalar-dtype") + _update_rs_value_fingerprint(digest, _rs_numpy_dtype_descriptor(scalar.dtype), active_containers) + add(b"numpy-scalar-data", scalar.reshape(1).view(np.uint8)) + elif isinstance(value, np.ndarray): + if value.dtype.hasobject: + raise TypeError("RS train-data fingerprint does not support object-dtype arrays") + array = np.ascontiguousarray(value) + add(b"ndarray-dtype") + _update_rs_value_fingerprint(digest, _rs_numpy_dtype_descriptor(array.dtype), active_containers) + add(b"ndarray-shape", ",".join(str(size) for size in array.shape).encode("ascii")) + add(b"ndarray-data", array.view(np.uint8).reshape(-1)) + elif isinstance(value, Mapping): + identity = id(value) + if identity in active_containers: + raise TypeError("RS train-data fingerprint does not support cyclic mappings") + active_containers.add(identity) + try: + items = sorted( + ((_rs_value_fingerprint(key, active_containers), key, item) for key, item in value.items()), + key=operator.itemgetter(0), + ) + if any(items[index - 1][0] == items[index][0] for index in range(1, len(items))): + raise TypeError("RS train-data fingerprint found ambiguous mapping keys") + add(b"mapping-size", str(len(items)).encode("ascii")) + for _, key, item in items: + add(b"mapping-key") + _update_rs_value_fingerprint(digest, key, active_containers) + add(b"mapping-value") + _update_rs_value_fingerprint(digest, item, active_containers) + finally: + active_containers.remove(identity) + elif isinstance(value, (list, tuple)): + identity = id(value) + if identity in active_containers: + raise TypeError("RS train-data fingerprint does not support cyclic sequences") + active_containers.add(identity) + try: + add(b"list" if isinstance(value, list) else b"tuple", str(len(value)).encode("ascii")) + primitive_type = type(value[0]) if value else None + if primitive_type in {bool, int, float} and all(type(item) is primitive_type for item in value): + add({bool: b"bool-sequence", int: b"int-sequence", float: b"float-sequence"}[primitive_type]) + for offset in range(0, len(value), 4096): + chunk = value[offset : offset + 4096] + add(b"primitive-chunk-size", len(chunk).to_bytes(4, "big")) + if primitive_type is bool: + payload = bytes(chunk) + elif primitive_type is int: + payload = ",".join(map(str, chunk)).encode("ascii") + else: + payload = struct.pack(f">{len(chunk)}d", *chunk) + add(b"primitive-chunk-data", payload) + else: + for item in value: + add(b"sequence-item") + _update_rs_value_fingerprint(digest, item, active_containers) + finally: + active_containers.remove(identity) + else: + value_type = type(value) + raise TypeError( + "RS train-data fingerprint does not support value type " + f"{value_type.__module__}.{value_type.__qualname__}" + ) + + +def _rs_value_fingerprint(value: Any, active_containers: set[int]) -> bytes: + """Hash supported train-data values without retaining their contents.""" + + digest = hashlib.sha256() + _update_rs_value_fingerprint(digest, value, active_containers) + return digest.digest() + + +def fingerprint_rs_train_data( + train_data: Mapping[str, Any], + *, + group_indices: list[Any], + weight_versions: list[Any], +) -> bytes: + """Fingerprint the exact ordered payload that will be split for training.""" + + if not isinstance(train_data, Mapping): + raise TypeError("RS train-data fingerprint requires a mapping") + return _rs_value_fingerprint( + { + "train_data": train_data, + "group_indices": group_indices, + "weight_versions": weight_versions, + }, + set(), + ) + + +def validate_rs_train_data_fingerprint( + train_data: Mapping[str, Any], + expected: bytes, + *, + group_indices: list[Any], + weight_versions: list[Any], +) -> None: + """Require rollout observability hooks to leave final training inputs unchanged.""" + + if not isinstance(expected, bytes) or len(expected) != hashlib.sha256().digest_size: + raise ValueError("invalid RS train-data fingerprint") + if ( + fingerprint_rs_train_data( + train_data, + group_indices=group_indices, + weight_versions=weight_versions, + ) + != expected + ): + raise RuntimeError( + "RS train data changed after proximal preflight; rollout observability hooks must be read-only." + ) + + +def clone_rs_masks(masks: list[Any]) -> list[torch.Tensor]: + """Clone masks before invoking an extension that may mutate its inputs.""" + + return [torch.as_tensor(mask).clone() for mask in masks] + + +def validate_final_rs_masks(original_masks: list[Any], *candidate_mask_sets: list[Any]) -> None: + """Require the training-time RS gate to preserve every preflight-accepted token.""" + + if not candidate_mask_sets: + raise ValueError("at least one candidate mask set is required") + + shape_changed: set[int] = set() + value_checks = [] + check_positions = [] + for candidate_masks in candidate_mask_sets: + if len(original_masks) != len(candidate_masks): + raise RuntimeError( + "RS training gate returned a different sample count from preflight: " + f"original={len(original_masks)}, modified={len(candidate_masks)}" + ) + for index, (original, modified) in enumerate(zip(original_masks, candidate_masks, strict=True)): + original = torch.as_tensor(original, dtype=torch.float32) + modified = torch.as_tensor(modified, dtype=torch.float32, device=original.device) + if original.shape != modified.shape: + shape_changed.add(index) + continue + value_checks.append(torch.all(original == modified)) + check_positions.append(index) + + changed = set(shape_changed) + if value_checks: + failed_checks = (~torch.stack(value_checks)).nonzero().flatten().tolist() + changed.update(check_positions[check_index] for check_index in failed_checks) + if changed: + raise RuntimeError( + "RS training gate rejected or changed samples accepted by preflight: " + f"microbatch_positions={sorted(changed)}" + ) + + +def validate_replacement_policy_version(groups: list[list[Any]], reports: list[dict[str, Any]]) -> str: + """Require every reactive replacement to come from the scored policy version.""" + + policy_versions = {str(report["policy_version"]) for report in reports} + if len(policy_versions) != 1: + raise ValueError(f"RS refill trainer ranks disagree on policy version: {policy_versions}") + expected_version = next(iter(policy_versions)) + + for group in groups: + for sample in group: + generated_versions = {str(version) for version in sample.weight_versions} + if generated_versions != {expected_version}: + raise ValueError( + "Reactive RS replacement was not generated entirely by the current rollout policy: " + f"sample_index={sample.index}, generated={sorted(generated_versions)}, " + f"expected={expected_version}." + ) + return expected_version + + +def validate_initial_policy_staleness( + groups: list[list[Any]], + reports: list[dict[str, Any]], + *, + max_staleness: int = 1, +) -> int: + """Require one initial rollout version no more than ``max_staleness`` behind the actor.""" + + if max_staleness < 0: + raise ValueError("max_staleness must be non-negative") + report_versions = {str(report["policy_version"]) for report in reports} + if len(report_versions) != 1: + raise ValueError(f"RS refill trainer ranks disagree on policy version: {report_versions}") + + generated_versions: set[str] = set() + for group in groups: + for sample in group: + sample_versions = {str(version) for version in sample.weight_versions} + if len(sample_versions) != 1: + raise ValueError( + "Each initial RS candidate must come from exactly one rollout policy version: " + f"sample_index={sample.index}, versions={sorted(sample_versions)}" + ) + generated_versions.update(sample_versions) + if len(generated_versions) != 1: + raise ValueError(f"Initial RS candidates span multiple rollout policy versions: {sorted(generated_versions)}") + + actor_version_text = next(iter(report_versions)) + rollout_version_text = next(iter(generated_versions)) + try: + actor_version = int(actor_version_text) + rollout_version = int(rollout_version_text) + except ValueError as error: + raise ValueError( + "RS batch refill requires integer actor and rollout policy versions, " + f"got actor={actor_version_text!r}, rollout={rollout_version_text!r}" + ) from error + + staleness = actor_version - rollout_version + if not 0 <= staleness <= max_staleness: + raise ValueError( + "Initial RS candidates exceed the hard policy-staleness bound: " + f"actor={actor_version}, rollout={rollout_version}, staleness={staleness}, " + f"allowed=[0, {max_staleness}]" + ) + return staleness + + +def select_accepted_groups( + groups: list[list[Any]], + reports: list[dict[str, Any]], + *, + target_size: int, + known_sample_indices: set[int] | None = None, + known_group_indices: set[int] | None = None, +) -> RefillSelection: + """Select complete prompt groups in input order without mismatch ranking.""" + + if target_size <= 0: + raise ValueError("target_size must be positive") + + report_by_index: dict[int, dict[str, Any]] = {} + for report in reports: + sample_index = report["sample_index"] + if sample_index in report_by_index: + raise ValueError(f"duplicate RS report for sample_index={sample_index}") + report_by_index[sample_index] = report + + accepted: list[list[Any]] = [] + rejected: list[list[Any]] = [] + surplus: list[list[Any]] = [] + seen_sample_indices = set(known_sample_indices or ()) + seen_group_indices = set(known_group_indices or ()) + candidate_sample_indices: set[int] = set() + + for group in groups: + if not group: + raise ValueError("RS refill received an empty prompt group") + group_indices = {sample.group_index for sample in group} + if len(group_indices) != 1: + raise ValueError(f"samples in one prompt group have different group_index values: {group_indices}") + group_index = next(iter(group_indices)) + if group_index in seen_group_indices: + raise ValueError(f"duplicate prompt group_index in RS candidates: {group_index}") + seen_group_indices.add(group_index) + + group_accepted = True + for sample in group: + if sample.index in seen_sample_indices: + raise ValueError(f"duplicate sample index in candidate groups: {sample.index}") + seen_sample_indices.add(sample.index) + candidate_sample_indices.add(sample.index) + if sample.index not in report_by_index: + raise ValueError(f"missing RS report for sample_index={sample.index}") + + report = report_by_index[sample.index] + if report.get("group_index") != sample.group_index: + raise ValueError( + f"RS report group mismatch for sample_index={sample.index}: " + f"sample={sample.group_index}, report={report.get('group_index')}" + ) + if int(report["valid_tokens"]) <= 0 or not bool(report["gate_passed"]): + group_accepted = False + + if group_accepted and len(accepted) < target_size: + accepted.append(group) + elif group_accepted: + surplus.append(group) + else: + rejected.append(group) + + unknown_reports = set(report_by_index) - candidate_sample_indices + if unknown_reports: + raise ValueError(f"RS reports contain unknown sample indices: {sorted(unknown_reports)}") + + return RefillSelection(accepted, rejected, surplus, target_size) + + +def run_rs_batch_refill( + actor_model, + rollout_manager, + rollout_id: int, + *, + resolve: Callable[[Any], Any], + clock: Callable[[], float], + rpc_timeout_seconds: float = 1800.0, +): + """Drive bounded refill rounds through Ray-like actor interfaces.""" + + if not math.isfinite(rpc_timeout_seconds) or rpc_timeout_seconds <= 0: + raise ValueError("rpc_timeout_seconds must be finite and positive") + + def resolve_rpc(value): + return resolve(value, timeout=rpc_timeout_seconds) + + coordinator_start = clock() + try: + while True: + preflight_start = clock() + candidate_refs = resolve_rpc(rollout_manager.prepare_rs_candidate_data.remote(rollout_id)) + report_refs = actor_model.async_score_rs_candidates(rollout_id, candidate_refs) + preflight_seconds = clock() - preflight_start + status = resolve_rpc( + rollout_manager.apply_rs_candidate_reports.remote(rollout_id, report_refs, preflight_seconds) + ) + if status["exhausted"]: + raise RuntimeError( + "RS batch refill exhausted its retry budget before optimizer.step: " + f"accepted={status['accepted_groups']}, target={status['target_groups']}, " + f"rounds={status['round']}, remaining={status['deficit']}." + ) + selected_cache_refs = actor_model.async_take_rs_candidate_log_probs( + rollout_id, + status["accepted_sample_indices"], + ) + resolve_rpc(rollout_manager.store_rs_accepted_log_probs.remote(rollout_id, selected_cache_refs)) + + if status["complete"]: + coordinator_seconds = clock() - coordinator_start + return resolve_rpc(rollout_manager.finalize_rs_batch.remote(rollout_id, coordinator_seconds)) + # Replacement generation retains the rollout backend's own timeout + # and health-monitor semantics. The coordination timeout applies + # only after generation has returned. + resolve(rollout_manager.generate_rs_replacement_candidates.remote(rollout_id)) + except Exception: + actor_cleanup = None + manager_cleanup = None + try: + actor_cleanup = actor_model.async_discard_rs_candidate_log_probs(rollout_id) + except Exception: + logger.exception("Failed to submit actor-local RS candidate cache cleanup after a coordination error") + try: + manager_cleanup = rollout_manager.abort_rs_batch.remote(rollout_id) + except Exception: + logger.exception("Failed to submit manager-local RS pending-state cleanup after a coordination error") + if actor_cleanup is not None: + try: + resolve_rpc(actor_cleanup) + except Exception: + logger.exception("Failed to discard actor-local RS candidate caches after a coordination error") + if manager_cleanup is not None: + try: + resolve_rpc(manager_cleanup) + except Exception: + logger.exception("Failed to discard manager-local RS pending state after a coordination error") + raise diff --git a/tests/test_data_source_epoch_wrap.py b/tests/test_data_source_epoch_wrap.py new file mode 100644 index 0000000000..966e466cff --- /dev/null +++ b/tests/test_data_source_epoch_wrap.py @@ -0,0 +1,87 @@ +from types import SimpleNamespace + +import pytest + +from slime.rollout.data_source import RolloutDataSource +from slime.utils.types import Sample + +NUM_GPUS = 0 + + +class _Dataset: + def __init__(self, prompts): + self.samples = [Sample(prompt=prompt) for prompt in prompts] + self.shuffle_epochs = [] + + def __len__(self): + return len(self.samples) + + def shuffle(self, epoch_id): + self.shuffle_epochs.append(epoch_id) + self.samples.reverse() + + +def _source(prompts, *, sample_offset=0, rollout_shuffle=False): + source = object.__new__(RolloutDataSource) + source.args = SimpleNamespace(n_samples_per_prompt=1, rollout_shuffle=rollout_shuffle) + source.dataset = _Dataset(prompts) + source.sample_offset = sample_offset + source.epoch_id = 0 + source.sample_group_index = 0 + source.sample_index = 0 + return source + + +def _prompts(groups): + return [group[0].prompt for group in groups] + + +def test_get_samples_wraps_across_every_required_epoch(): + source = _source(["a", "b"]) + + groups = source.get_samples(7) + + assert _prompts(groups) == ["a", "b", "a", "b", "a", "b", "a"] + assert source.epoch_id == 3 + assert source.sample_offset == 1 + assert [group[0].index for group in groups] == list(range(7)) + assert [group[0].group_index for group in groups] == list(range(7)) + + +def test_get_samples_wraps_from_nonzero_offset_and_preserves_boundary_position(): + source = _source(["a", "b", "c"], sample_offset=2) + + assert _prompts(source.get_samples(7)) == ["c", "a", "b", "c", "a", "b", "c"] + assert source.epoch_id == 2 + assert source.sample_offset == 3 + + assert _prompts(source.get_samples(1)) == ["a"] + assert source.epoch_id == 3 + assert source.sample_offset == 1 + + +def test_get_samples_shuffles_at_each_crossed_epoch(): + source = _source(["a", "b"], rollout_shuffle=True) + + assert _prompts(source.get_samples(6)) == ["a", "b", "b", "a", "a", "b"] + assert source.dataset.shuffle_epochs == [1, 2] + assert source.epoch_id == 2 + assert source.sample_offset == 2 + + +def test_get_samples_rejects_nonempty_request_from_empty_dataset(): + source = _source([]) + + with pytest.raises(RuntimeError, match="empty dataset"): + source.get_samples(1) + + +def test_get_samples_rejects_an_invalid_dataset_offset(): + source = _source(["a", "b"], sample_offset=3) + + with pytest.raises(RuntimeError, match="invalid dataset offset"): + source.get_samples(1) + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__])) diff --git a/tests/test_megatron_argument_validation.py b/tests/test_megatron_argument_validation.py index 20c133ed95..06d67029da 100644 --- a/tests/test_megatron_argument_validation.py +++ b/tests/test_megatron_argument_validation.py @@ -1,3 +1,4 @@ +import argparse import importlib.util import sys import types @@ -85,6 +86,28 @@ def make_qwen3_6_args(**overrides): return types.SimpleNamespace(**values) +@pytest.mark.unit +def test_rs_refill_cli_parser_contract(monkeypatch): + module = load_slime_arguments_module(monkeypatch) + parser = argparse.ArgumentParser() + module.get_slime_extra_args_provider()(parser) + + defaults = parser.parse_args(["--rollout-batch-size", "1"]) + enabled = parser.parse_args(["--rollout-batch-size", "1", "--rs-batch-refill", "--rs-refill-max-rounds", "7"]) + + assert defaults.rs_batch_refill is False + assert defaults.rs_refill_max_rounds == 2 + assert defaults.rs_refill_rpc_timeout_seconds == 1800.0 + assert defaults.rs_refill_max_candidate_cache_bytes == 1 << 30 + assert enabled.rs_batch_refill is True + assert enabled.rs_refill_max_rounds == 7 + help_text = parser.format_help() + assert "--rs-batch-refill" in help_text + assert "--rs-refill-max-rounds" in help_text + assert "--rs-refill-rpc-timeout-seconds" in help_text + assert "--rs-refill-max-candidate-cache-bytes" in help_text + + def make_qwen3_6_hf_config(): text_config = types.SimpleNamespace( hidden_size=2048, @@ -204,6 +227,21 @@ def make_slime_validate_args(**overrides): use_tis=False, get_mismatch_metrics=False, custom_tis_function_path=None, + custom_pg_loss_reducer_function_path=None, + rs_batch_refill=False, + rs_refill_max_rounds=2, + rs_refill_rpc_timeout_seconds=1800.0, + rs_refill_max_candidate_cache_bytes=1 << 30, + use_rs=False, + tis_mode="truncate", + tis_level="token", + tis_lower_bound=0.1, + tis_upper_bound=10.0, + tis_batch_normalize=False, + rs_level="token", + rs_lower_bound=0.8, + rs_upper_bound=1.2, + rs_veto_threshold=None, use_dynamic_batch_size=False, max_tokens_per_gpu=None, log_probs_max_tokens_per_gpu=None, @@ -229,6 +267,8 @@ def make_slime_validate_args(**overrides): rollout_num_gpus=8, eval_function_path=None, rollout_function_path="custom.rollout", + rollout_top_p=1.0, + rollout_top_k=-1, num_steps_per_rollout=None, rollout_batch_size=1, n_samples_per_prompt=1, @@ -243,6 +283,16 @@ def make_slime_validate_args(**overrides): use_rollout_routing_replay=False, use_routing_replay=False, custom_config_path=None, + megatron_config_path=None, + custom_model_provider_path=None, + custom_megatron_init_path=None, + custom_megatron_before_log_prob_hook_path=None, + custom_megatron_before_train_step_hook_path=None, + custom_convert_samples_to_train_data_path=None, + custom_reward_post_process_path=None, + rollout_data_postprocess_path=None, + custom_rollout_log_function_path=None, + dynamic_sampling_filter_path=None, eval_max_context_len=None, rollout_max_context_len=None, rollout_max_prompt_len=None, @@ -255,12 +305,151 @@ def make_slime_validate_args(**overrides): update_weight_disk_dir=None, update_weight_local_checkpoint_dir=None, update_weight_mode="full", + update_weights_interval=1, + update_weight_start_version=0, + ref_update_interval=None, rollout_temperature=1.0, + context_parallel_size=1, + attention_dropout=0.0, + hidden_dropout=0.0, + fp8=None, + fp8_param_gather=False, + fp4=None, + te_precision_config_file=None, + kitchen_config_file=None, + kitchen_recipe_number=None, + megatron_deepgemm_forward_layers=None, + megatron_deepgemm_forward_modules=None, + megatron_deepgemm_moe_forward_layers=None, + megatron_deepgemm_moe_forward_modules=None, + moe_input_jitter_eps=None, + moe_router_force_load_balancing=False, + moe_expert_capacity_factor=None, + moe_router_load_balancing_type="aux_loss", + compute_advantages_and_returns=True, + custom_advantage_function_path=None, + loss_type="policy_loss", + use_opsm=False, + use_rollout_entropy=False, + partial_rollout=False, ) values.update(overrides) return types.SimpleNamespace(**values) +def make_rs_refill_args(**overrides): + values = dict( + rs_batch_refill=True, + use_tis=True, + use_rs=True, + rs_level="geometric", + rollout_global_dataset=True, + rollout_batch_size=2, + n_samples_per_prompt=2, + global_batch_size=4, + use_dynamic_batch_size=True, + max_tokens_per_gpu=1024, + ) + values.update(overrides) + return make_slime_validate_args(**values) + + +@pytest.mark.unit +def test_rs_refill_accepts_the_supported_exact_configuration(monkeypatch): + module = load_slime_arguments_module(monkeypatch) + + args = make_rs_refill_args() + module.slime_validate_args(args) + + assert args.start_rollout_id == 0 + assert args.finetune is True + + +@pytest.mark.unit +def test_rs_refill_accepts_explicitly_disabled_fp_quantization(monkeypatch): + module = load_slime_arguments_module(monkeypatch) + + module.slime_validate_args(make_rs_refill_args(fp8=False, fp4=False)) + + +@pytest.mark.unit +def test_rs_refill_accepts_a_read_only_custom_rollout_logger(monkeypatch): + module = load_slime_arguments_module(monkeypatch) + + module.slime_validate_args(make_rs_refill_args(custom_rollout_log_function_path="custom.log")) + + +@pytest.mark.unit +@pytest.mark.parametrize( + ("overrides", "message"), + [ + ({"train_backend": "fsdp"}, "train-backend megatron"), + ({"colocate": True}, "disaggregated"), + ({"rollout_global_dataset": False}, "rollout-global-dataset"), + ({"update_weights_interval": 2}, "update-weights-interval 1"), + ({"start_rollout_id": 3}, "starting at rollout 0"), + ({"global_batch_size": 8}, "one optimizer step"), + ({"custom_model_provider_path": "custom.model"}, "custom-model-provider-path"), + ({"custom_megatron_init_path": "custom.init"}, "custom-megatron-init-path"), + ({"update_weight_transport": "disk"}, "full weight updates over NCCL"), + ({"rollout_top_k": 8}, "rollout-top-k -1"), + ({"use_tis": False}, "requires --use-tis"), + ({"custom_tis_function_path": "custom.tis"}, "built-in TIS/RS"), + ({"use_rs": False}, "requires use_rs"), + ({"rs_level": "token"}, "rs_level: sequence or geometric"), + ({"tis_mode": "mask"}, "tis_mode: truncate or clip"), + ({"tis_batch_normalize": True}, "tis_batch_normalize"), + ({"attention_dropout": 0.1}, "attention-dropout 0"), + ({"fp8": "hybrid"}, "non-FP8 Megatron training"), + ({"fp8_param_gather": True}, "fp8-param-gather"), + ({"fp4": "nvfp4"}, "non-FP4 Megatron training"), + ({"te_precision_config_file": "/tmp/te.yaml"}, "Transformer Engine precision config"), + ({"kitchen_config_file": "/tmp/kitchen.yaml"}, "Kitchen quantization"), + ({"kitchen_recipe_number": 1}, "Kitchen quantization"), + ({"megatron_deepgemm_forward_layers": [0]}, "DeepGEMM FP8 forward"), + ({"megatron_deepgemm_forward_modules": ["linear_qkv"]}, "DeepGEMM FP8 forward"), + ({"megatron_deepgemm_moe_forward_layers": [0]}, "DeepGEMM FP8 forward"), + ({"megatron_deepgemm_moe_forward_modules": ["mlp.experts"]}, "DeepGEMM FP8 forward"), + ({"moe_input_jitter_eps": 0.1}, "moe-input-jitter-eps"), + ({"moe_router_force_load_balancing": True}, "moe-router-force-load-balancing"), + ({"moe_expert_capacity_factor": 1.25}, "moe-expert-capacity-factor"), + ({"moe_router_load_balancing_type": "sinkhorn"}, "Sinkhorn MoE routing"), + ({"moe_router_load_balancing_type": ["aux_loss", "sinkhorn"]}, "Sinkhorn MoE routing"), + ({"context_parallel_size": 2}, "context-parallel-size 1"), + ({"partial_rollout": True}, "partial-rollout"), + ({"dynamic_sampling_filter_path": "custom.filter"}, "dynamic-sampling-filter-path"), + ( + {"rollout_function_path": "slime.rollout.fully_async_rollout.generate_rollout_fully_async"}, + "fully-async rollout queue", + ), + ({"rs_refill_max_rounds": -1}, "max-rounds"), + ({"rs_refill_max_rounds": 1.5}, "max-rounds"), + ({"rs_refill_max_rounds": False}, "max-rounds"), + ({"rs_refill_rpc_timeout_seconds": 0}, "rpc-timeout"), + ({"rs_refill_rpc_timeout_seconds": float("inf")}, "rpc-timeout"), + ({"rs_refill_max_candidate_cache_bytes": 0}, "candidate-cache"), + ({"rs_refill_max_candidate_cache_bytes": False}, "candidate-cache"), + ({"rs_lower_bound": 2.0, "rs_upper_bound": 1.0}, "finite RS bounds"), + ], +) +def test_rs_refill_rejects_unsupported_configurations(monkeypatch, overrides, message): + module = load_slime_arguments_module(monkeypatch) + + with pytest.raises(ValueError, match=message): + module.slime_validate_args(make_rs_refill_args(**overrides)) + + +@pytest.mark.unit +def test_rs_refill_rejects_checkpoint_resume(monkeypatch, tmp_path): + module = load_slime_arguments_module(monkeypatch) + checkpoint = tmp_path / "actor" + checkpoint.mkdir() + (checkpoint / "latest_checkpointed_iteration.txt").write_text("3", encoding="utf-8") + + with pytest.raises(ValueError, match="checkpoint resume is not implemented"): + module.slime_validate_args(make_rs_refill_args(load=str(checkpoint), start_rollout_id=4)) + + @pytest.mark.unit @pytest.mark.parametrize("megatron_to_hf_mode", ["raw", "bridge"]) def test_slime_validate_args_preserves_explicit_start_rollout_id(monkeypatch, megatron_to_hf_mode): diff --git a/tests/test_rollout_data_utils.py b/tests/test_rollout_data_utils.py index 82a72de105..ee08a0c51b 100644 --- a/tests/test_rollout_data_utils.py +++ b/tests/test_rollout_data_utils.py @@ -50,6 +50,7 @@ def test_tensorize_rollout_data_for_training_normalizes_cpu_tensors(): rollout_data = { "tokens": [readonly_tokens], "loss_masks": [[1, 0]], + "rs_preflight_log_probs": [torch.tensor([-0.2, -0.3], dtype=torch.float64, requires_grad=True)], "multimodal_train_inputs": [ { "pixel_values": torch.tensor([1.0], requires_grad=True), @@ -63,6 +64,10 @@ def test_tensorize_rollout_data_for_training_normalizes_cpu_tensors(): assert rollout_data["tokens"][0].dtype == torch.long assert rollout_data["loss_masks"][0].dtype == torch.int + assert rollout_data["rs_preflight_log_probs"][0].dtype == torch.float32 + assert rollout_data["rs_preflight_log_probs"][0].device.type == "cpu" + assert rollout_data["rs_preflight_log_probs"][0].is_contiguous() + assert not rollout_data["rs_preflight_log_probs"][0].requires_grad assert rollout_data["multimodal_train_inputs"][0]["metadata"] == "unchanged" assert not rollout_data["multimodal_train_inputs"][0]["pixel_values"].requires_grad assert rollout_data["rollout_mask_sums"].dtype == torch.float32 diff --git a/tests/test_rs_exact_refill_actor.py b/tests/test_rs_exact_refill_actor.py new file mode 100644 index 0000000000..b78ae49af7 --- /dev/null +++ b/tests/test_rs_exact_refill_actor.py @@ -0,0 +1,487 @@ +"""CPU contracts for actor-local exact RS refill scoring and cache reuse.""" + +import importlib.util +import sys +import types +from contextlib import nullcontext +from pathlib import Path +from types import SimpleNamespace + +import pytest +import torch + + +def _stub_module(name, **attrs): + module = types.ModuleType(name) + for key, value in attrs.items(): + setattr(module, key, value) + sys.modules[name] = module + return module + + +def _noop(*_args, **_kwargs): + return None + + +def _module_available(name): + try: + return importlib.util.find_spec(name) is not None + except (ImportError, ValueError): + return False + + +def _snapshot_modules(prefixes): + return { + name: module + for name, module in sys.modules.items() + if any(name == prefix or name.startswith(prefix + ".") for prefix in prefixes) + } + + +def _restore_modules(prefixes, snapshot): + for name in list(sys.modules): + if any(name == prefix or name.startswith(prefix + ".") for prefix in prefixes): + sys.modules.pop(name, None) + sys.modules.update(snapshot) + + +# The routine CPU job intentionally omits Slime's patched Megatron/TMS +# runtime. Keep actor.py and loss.py real and stub only unexercised backend +# edges needed while importing them. +_STUB_PREFIXES = ("megatron", "torch_memory_saver", "slime.backends.megatron_utils") +_STUB_SNAPSHOT = None +if "megatron.core.mpu" not in sys.modules and not _module_available("megatron"): + _STUB_SNAPSHOT = _snapshot_modules(_STUB_PREFIXES) + _megatron_utils = _stub_module("slime.backends.megatron_utils") + _megatron_utils.__path__ = [str(Path(__file__).resolve().parents[1] / "slime" / "backends" / "megatron_utils")] + _mpu = _stub_module( + "megatron.core.mpu", + is_pipeline_last_stage=_noop, + get_tensor_model_parallel_rank=_noop, + ) + _megatron_core = _stub_module("megatron.core", mpu=_mpu) + _stub_module("megatron", core=_megatron_core) + _stub_module( + "torch_memory_saver", + torch_memory_saver=types.SimpleNamespace(pause=_noop, resume=_noop), + ) + _stub_module( + "slime.backends.megatron_utils.cp_utils", + all_gather_with_cp=_noop, + get_logits_and_tokens_offset_with_cp=_noop, + get_sum_of_sample_mean=_noop, + prepare_routed_experts_for_routing_replay=_noop, + slice_log_prob_with_cp=_noop, + ) + _stub_module("slime.backends.megatron_utils.checkpoint", load_checkpoint=_noop) + _stub_module("slime.backends.megatron_utils.data", DataIterator=object, get_data_iterator=_noop) + _stub_module("slime.backends.megatron_utils.hf_checkpoint_saver", save_hf_model_to_path=_noop) + _stub_module("slime.backends.megatron_utils.initialize", init=_noop, is_megatron_main_rank=lambda: False) + _stub_module( + "slime.backends.megatron_utils.model", + forward_only=_noop, + initialize_model_and_optimizer=_noop, + save=_noop, + train=_noop, + ) + _stub_module("slime.backends.megatron_utils.update_weight.common", named_params_and_buffers=_noop) + for _module_name, _class_name in ( + ("update_weight_from_disk", "UpdateWeightFromDisk"), + ("update_weight_from_distributed", "UpdateWeightFromDistributed"), + ("update_weight_from_tensor", "UpdateWeightFromTensor"), + ): + _stub_module( + f"slime.backends.megatron_utils.update_weight.{_module_name}", + **{_class_name: type(_class_name, (), {})}, + ) + +try: + from slime.backends.megatron_utils import actor as actor_module + from slime.backends.megatron_utils import loss as loss_module + from slime.backends.megatron_utils.actor import MegatronTrainRayActor + from slime.ray.actor_group import RayTrainGroup +finally: + if _STUB_SNAPSHOT is not None: + _restore_modules(_STUB_PREFIXES, _STUB_SNAPSHOT) + _backends_package = sys.modules.get("slime.backends") + if _backends_package is not None: + if "slime.backends.megatron_utils" in _STUB_SNAPSHOT: + _backends_package.megatron_utils = _STUB_SNAPSHOT["slime.backends.megatron_utils"] + else: + _backends_package.__dict__.pop("megatron_utils", None) + +NUM_GPUS = 0 + + +def _make_args(**overrides): + values = { + "advantage_estimator": "grpo", + "calculate_per_token_loss": False, + "compute_advantages_and_returns": True, + "custom_pg_loss_reducer_function_path": None, + "custom_tis_function_path": None, + "entropy_coef": 0.0, + "eps_clip": 0.2, + "eps_clip_c": None, + "eps_clip_high": 0.2, + "get_mismatch_metrics": True, + "keep_old_actor": False, + "kl_coef": 0.0, + "loss_type": "policy_loss", + "ref_update_interval": None, + "rollout_top_p": 1.0, + "rs_batch_refill": True, + "rs_refill_max_candidate_cache_bytes": 1 << 20, + "rs_level": "sequence", + "rs_lower_bound": 0.5, + "rs_upper_bound": 2.0, + "rs_veto_threshold": None, + "save_debug_train_data": None, + "tis_lower_bound": None, + "tis_upper_bound": 2.0, + "use_critic": False, + "use_kl_loss": False, + "use_opd": False, + "use_opsm": False, + "use_rollout_logprobs": False, + "use_rollout_routing_replay": False, + "use_routing_replay": False, + "use_tis": True, + } + values.update(overrides) + return SimpleNamespace(**values) + + +class _WeightsBackuper: + def __init__(self): + self.backup_tags = {"actor"} + self.backups = [] + self.restores = [] + + def backup(self, tag): + self.backups.append(tag) + + def restore(self, tag): + self.restores.append(tag) + + +def _make_actor(args=None): + actor = object.__new__(MegatronTrainRayActor) + actor.args = args or _make_args() + actor._active_model_tag = "actor" + actor._rs_candidate_log_probs = {} + actor.weights_backuper = _WeightsBackuper() + actor.weight_updater = SimpleNamespace(weight_version=9, pop_metrics=lambda: {}) + actor.prof = SimpleNamespace(step=lambda **_kwargs: None) + actor.model = object() + actor.optimizer = object() + actor.opt_param_scheduler = object() + actor.rollout_data_postprocess = None + return actor + + +def _patch_last_pipeline_rank(monkeypatch): + monkeypatch.setattr(actor_module.mpu, "is_pipeline_last_stage", lambda: True) + monkeypatch.setattr(actor_module.mpu, "get_tensor_model_parallel_rank", lambda: 0) + + +def test_preflight_scores_and_transfers_only_selected_cache_once(monkeypatch): + actor = _make_actor() + train_log_probs = [ + torch.tensor([0.0, 0.0], dtype=torch.float64, requires_grad=True), + torch.tensor([1.0, 1.0], dtype=torch.float64, requires_grad=True), + ] + rollout_data = { + "group_indices": [3, 4], + "loss_masks": [torch.ones(2), torch.ones(2)], + "num_microbatches": [2], + "rollout_log_probs": [torch.zeros(2), torch.zeros(2)], + "sample_indices": [10, 11], + } + actor._get_rollout_data = lambda _ref: rollout_data + actor.compute_log_prob = lambda *_args, **_kwargs: {"log_probs": train_log_probs} + monkeypatch.setattr(actor_module, "get_data_iterator", lambda data: data) + _patch_last_pipeline_rank(monkeypatch) + + reports = actor.score_rs_candidates(27, object()) + + assert reports == [ + { + "sample_index": 10, + "group_index": 3, + "valid_tokens": 2, + "gate_passed": True, + "policy_version": "9", + "candidate_cache_bytes": 8, + }, + { + "sample_index": 11, + "group_index": 4, + "valid_tokens": 2, + "gate_passed": False, + "policy_version": "9", + "candidate_cache_bytes": 8, + }, + ] + cached = actor._rs_candidate_log_probs[27] + assert set(cached) == {10, 11} + assert all(value.device.type == "cpu" and value.dtype == torch.float32 for value in cached.values()) + assert all(not value.requires_grad for value in cached.values()) + + selected = actor.take_rs_candidate_log_probs(27, [10, 99]) + + assert set(selected) == {10} + torch.testing.assert_close(selected[10], torch.zeros(2)) + assert 27 not in actor._rs_candidate_log_probs + with pytest.raises(RuntimeError, match="No RS candidate logprob cache"): + actor.take_rs_candidate_log_probs(27, [10]) + + +def test_preflight_metadata_mismatch_does_not_publish_cache(monkeypatch): + actor = _make_actor() + rollout_data = { + "group_indices": [], + "loss_masks": [torch.ones(1)], + "num_microbatches": [1], + "rollout_log_probs": [torch.zeros(1)], + "sample_indices": [5], + } + actor._get_rollout_data = lambda _ref: rollout_data + actor.compute_log_prob = lambda *_args, **_kwargs: {"log_probs": [torch.zeros(1)]} + monkeypatch.setattr(actor_module, "get_data_iterator", lambda data: data) + _patch_last_pipeline_rank(monkeypatch) + + with pytest.raises(RuntimeError, match="logprob/group count mismatch"): + actor.score_rs_candidates(12, object()) + + assert 12 not in actor._rs_candidate_log_probs + + +def test_preflight_rejects_candidate_cache_before_pinned_allocation(monkeypatch): + actor = _make_actor(_make_args(rs_refill_max_candidate_cache_bytes=7)) + rollout_data = { + "group_indices": [3], + "loss_masks": [torch.ones(2)], + "num_microbatches": [1], + "rollout_log_probs": [torch.zeros(2)], + "sample_indices": [10], + } + actor._get_rollout_data = lambda _ref: rollout_data + actor.compute_log_prob = lambda *_args, **_kwargs: {"log_probs": [torch.zeros(2)]} + monkeypatch.setattr(actor_module, "get_data_iterator", lambda data: data) + _patch_last_pipeline_rank(monkeypatch) + + with pytest.raises(RuntimeError, match=r"required=8, limit=7"): + actor.score_rs_candidates(27, object()) + + assert 27 not in actor._rs_candidate_log_probs + + +def test_duplicate_take_request_preserves_cache_for_cleanup(monkeypatch): + actor = _make_actor() + actor._rs_candidate_log_probs[4] = {2: torch.zeros(1)} + _patch_last_pipeline_rank(monkeypatch) + + with pytest.raises(ValueError, match="must be unique"): + actor.take_rs_candidate_log_probs(4, [2, 2]) + + assert 4 in actor._rs_candidate_log_probs + actor.discard_rs_candidate_log_probs(4) + actor.discard_rs_candidate_log_probs(4) + assert actor._rs_candidate_log_probs == {} + + +def _patch_train_actor_dependencies(monkeypatch, train_calls): + monkeypatch.setattr(actor_module, "get_data_iterator", lambda data: data) + monkeypatch.setattr(actor_module, "compute_advantages_and_returns", lambda _args, _data: None) + monkeypatch.setattr(actor_module.train_metric_utils, "log_rollout_data", lambda *_args, **_kwargs: None) + monkeypatch.setattr(actor_module.train_metric_utils, "log_perf_data", lambda *_args, **_kwargs: None) + monkeypatch.setattr(actor_module.train_data_utils, "save_debug_train_data", lambda *_args, **_kwargs: None) + monkeypatch.setattr(actor_module, "inverse_timer", lambda _name: nullcontext()) + monkeypatch.setattr(actor_module, "timer", lambda _name: nullcontext()) + monkeypatch.setattr(actor_module, "train", lambda *args, **kwargs: train_calls.append((args, kwargs))) + + +def _training_batch(): + return { + "global_batch_sizes": [1], + "loss_masks": [torch.tensor([1, 1], dtype=torch.int)], + "num_microbatches": [1], + "response_lengths": [2], + "rs_preflight_log_probs": [torch.tensor([-0.2, -0.3])], + } + + +def test_train_actor_reuses_preflight_log_probs_without_recompute(monkeypatch): + actor = _make_actor() + train_calls = [] + _patch_train_actor_dependencies(monkeypatch, train_calls) + actor.compute_log_prob = lambda *_args, **_kwargs: pytest.fail("proximal logprobs were recomputed") + rollout_data = _training_batch() + cached = rollout_data["rs_preflight_log_probs"] + + actor.train_actor(8, rollout_data) + + assert rollout_data["log_probs"] is cached + assert "rs_preflight_log_probs" not in rollout_data + assert len(train_calls) == 1 + assert train_calls[0][0][4] is rollout_data + + +@pytest.mark.parametrize( + ("mutation", "error"), + [ + ("response_lengths", "response lengths changed after proximal preflight"), + ("loss_masks", "changed samples accepted by preflight"), + ], +) +def test_train_actor_rejects_post_preflight_batch_mutation(monkeypatch, mutation, error): + actor = _make_actor() + train_calls = [] + _patch_train_actor_dependencies(monkeypatch, train_calls) + actor.compute_log_prob = lambda *_args, **_kwargs: pytest.fail("proximal logprobs were recomputed") + + def mutate(_args, _rollout_id, data): + if mutation == "response_lengths": + data["response_lengths"][0] = 1 + else: + data["loss_masks"][0][0] = 0 + + actor.rollout_data_postprocess = mutate + + with pytest.raises(RuntimeError, match=error): + actor.train_actor(8, _training_batch()) + + assert train_calls == [] + + +@pytest.mark.parametrize( + ("enabled", "rollout_data", "error"), + [ + (True, {}, "missing its proximal logprob cache"), + (True, {"rs_preflight_log_probs": [], "log_probs": []}, "conflicting proximal logprob fields"), + ( + True, + { + "loss_masks": [torch.ones(1)], + "response_lengths": [1], + "rs_preflight_log_probs": [], + }, + "proximal cache/sample count mismatch", + ), + ( + True, + { + "loss_masks": [torch.ones(2)], + "response_lengths": [2], + "rs_preflight_log_probs": [torch.zeros(1)], + }, + "proximal cache shape mismatch", + ), + (False, {"rs_preflight_log_probs": []}, "while --rs-batch-refill is disabled"), + ], +) +def test_train_actor_rejects_invalid_preflight_cache_contract(enabled, rollout_data, error): + actor = _make_actor(_make_args(rs_batch_refill=enabled)) + + with pytest.raises(RuntimeError, match=error): + actor.train_actor(1, rollout_data) + + +def test_policy_loss_routes_rs_refill_through_shared_gate(monkeypatch): + args = _make_args(custom_tis_function_path="must-not-load") + cached_log_probs = [torch.zeros(2)] + loss_masks = [torch.ones(2)] + batch = { + "advantages": [torch.ones(2)], + "log_probs": cached_log_probs, + "loss_masks": loss_masks, + "response_lengths": [2], + "rollout_log_probs": [torch.zeros(2)], + "rollout_mask_sums": torch.tensor([2.0]), + "total_lengths": [3], + "unconcat_tokens": [torch.zeros(3, dtype=torch.long)], + } + calls = {} + + monkeypatch.setattr( + loss_module, + "get_log_probs_and_entropy", + lambda *_args, **_kwargs: ( + None, + {"entropy": [torch.zeros(2)], "log_probs": [torch.full((2,), 0.1)]}, + ), + ) + monkeypatch.setattr(loss_module, "get_sum_of_sample_mean", lambda *_args, **_kwargs: lambda value: value.mean()) + monkeypatch.setattr(loss_module, "load_function", lambda _path: pytest.fail("custom TIS was loaded")) + + def apply_gate(**kwargs): + calls["gate"] = kwargs + return kwargs["pg_loss"], [mask.clone() for mask in kwargs["loss_masks"]], {} + + def validate_masks(original, *candidates): + calls["validation"] = (original, candidates) + + monkeypatch.setattr(loss_module, "apply_rs_refill_tis", apply_gate) + monkeypatch.setattr(loss_module, "validate_final_rs_masks", validate_masks) + + loss_module.policy_loss_function( + args, + batch, + torch.zeros((1, 2, 2), requires_grad=True), + lambda value: value.mean(), + ) + + assert calls["gate"]["train_log_probs"] is cached_log_probs + assert calls["gate"]["loss_masks"] is loss_masks + original, candidates = calls["validation"] + assert len(candidates) == 2 + assert candidates[0] is loss_masks + assert original[0] is not loss_masks[0] + torch.testing.assert_close(original[0], loss_masks[0]) + + +class _RemoteMethod: + def __init__(self, actor_index, method, calls): + self.actor_index = actor_index + self.method = method + self.calls = calls + + def remote(self, *args): + call = (self.actor_index, self.method, args) + self.calls.append(call) + return call + + +class _ActorHandle: + def __init__(self, actor_index, calls): + self.score_rs_candidates = _RemoteMethod(actor_index, "score", calls) + self.take_rs_candidate_log_probs = _RemoteMethod(actor_index, "take", calls) + self.discard_rs_candidate_log_probs = _RemoteMethod(actor_index, "discard", calls) + + +def test_actor_group_fans_out_exact_refill_rpcs(): + calls = [] + group = object.__new__(RayTrainGroup) + group._actor_handlers = [_ActorHandle(0, calls), _ActorHandle(1, calls)] + rollout_ref = object() + + score_refs = group.async_score_rs_candidates(17, rollout_ref) + take_refs = group.async_take_rs_candidate_log_probs(17, [2, 5]) + discard_refs = group.async_discard_rs_candidate_log_probs(17) + + assert score_refs == calls[:2] + assert take_refs == calls[2:4] + assert discard_refs == calls[4:] + assert calls == [ + (0, "score", (17, rollout_ref)), + (1, "score", (17, rollout_ref)), + (0, "take", (17, [2, 5])), + (1, "take", (17, [2, 5])), + (0, "discard", (17,)), + (1, "discard", (17,)), + ] + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__])) diff --git a/tests/test_rs_exact_refill_rollout_manager.py b/tests/test_rs_exact_refill_rollout_manager.py new file mode 100644 index 0000000000..4f3fefaed4 --- /dev/null +++ b/tests/test_rs_exact_refill_rollout_manager.py @@ -0,0 +1,898 @@ +import importlib.util +import sys +import types +from types import SimpleNamespace + +import pytest +import torch + + +def _stub_module(name, **attrs): + module = types.ModuleType(name) + for key, value in attrs.items(): + setattr(module, key, value) + sys.modules[name] = module + return module + + +def _module_available(name): + try: + return importlib.util.find_spec(name) is not None + except (ImportError, ValueError): + return False + + +def _snapshot_modules(prefixes): + return { + name: module + for name, module in sys.modules.items() + if any(name == prefix or name.startswith(prefix + ".") for prefix in prefixes) + } + + +def _restore_modules(prefixes, snapshot): + for name in list(sys.modules): + if any(name == prefix or name.startswith(prefix + ".") for prefix in prefixes): + sys.modules.pop(name, None) + sys.modules.update(snapshot) + + +# The routine CPU job intentionally omits the SGLang CUDA runtime. These are +# import-only serving edges; every RolloutManager method under test stays real. +_STUB_PREFIXES = ("sglang", "slime.backends.sglang_utils", "slime.ray.rollout") +_STUB_SNAPSHOT = None +if "sglang.srt.constants" not in sys.modules and not _module_available("sglang"): + _STUB_SNAPSHOT = _snapshot_modules(_STUB_PREFIXES) + _sglang = _stub_module("sglang") + _sglang_srt = _stub_module("sglang.srt") + _sglang_constants = _stub_module( + "sglang.srt.constants", + GPU_MEMORY_TYPE_CUDA_GRAPH="cuda_graph", + GPU_MEMORY_TYPE_KV_CACHE="kv_cache", + GPU_MEMORY_TYPE_WEIGHTS="weights", + ) + _sglang.srt = _sglang_srt + _sglang_srt.constants = _sglang_constants + _stub_module( + "slime.backends.sglang_utils.sglang_engine", + SGLangEngine=type("SGLangEngine", (), {}), + ) + +try: + import slime.ray.rollout as rollout_module + from slime.ray.rollout import RolloutManager + from slime.utils.rs_refill import run_rs_batch_refill + from slime.utils.types import Sample +finally: + if _STUB_SNAPSHOT is not None: + _restore_modules(_STUB_PREFIXES, _STUB_SNAPSHOT) + for _parent_name, _attribute, _module_name in ( + ("slime.backends", "sglang_utils", "slime.backends.sglang_utils"), + ("slime.ray", "rollout", "slime.ray.rollout"), + ): + _parent = sys.modules.get(_parent_name) + if _parent is None: + continue + if _module_name in _STUB_SNAPSHOT: + setattr(_parent, _attribute, _STUB_SNAPSHOT[_module_name]) + else: + _parent.__dict__.pop(_attribute, None) + +NUM_GPUS = 0 + + +class _RemoteMethod: + def __init__(self, name, function, events): + self.name = name + self.function = function + self.events = events + + def remote(self, *args): + self.events.append((self.name, args)) + return self.function(*args) + + +class _LifecycleActor: + def __init__(self, reports, cache_refs): + self.reports = reports + self.cache_refs = cache_refs + self.events = [] + self.cache_live = True + + def async_score_rs_candidates(self, rollout_id, candidates): + self.events.append(("score", rollout_id, candidates)) + if self.reports and isinstance(self.reports[0], list): + return self.reports + return [self.reports] + + def async_take_rs_candidate_log_probs(self, rollout_id, selected): + self.events.append(("take", rollout_id, list(selected))) + if isinstance(self.cache_refs, list): + self.cache_live = False + return self.cache_refs + + def async_discard_rs_candidate_log_probs(self, rollout_id): + self.events.append(("discard", rollout_id)) + self.cache_live = False + return True + + +class _RoundLifecycleActor: + def __init__(self, report_rounds, cache_rounds): + self.report_rounds = list(report_rounds) + self.cache_rounds = list(cache_rounds) + self.events = [] + self.current_cache = None + + def async_score_rs_candidates(self, rollout_id, candidates): + self.events.append(("score", rollout_id, candidates)) + self.current_cache = self.cache_rounds.pop(0) + reports = self.report_rounds.pop(0) + if reports and isinstance(reports[0], list): + return reports + return [reports] + + def async_take_rs_candidate_log_probs(self, rollout_id, selected): + self.events.append(("take", rollout_id, list(selected))) + cache = self.current_cache + self.current_cache = None + return [{index: cache[index] for index in selected}] + + def async_discard_rs_candidate_log_probs(self, rollout_id): + self.events.append(("discard", rollout_id)) + self.current_cache = None + return True + + +def _sample(index, group_index, *, weight_version="7", loss_mask=None): + if loss_mask is None: + loss_mask = [1, 1] + elif not isinstance(loss_mask, torch.Tensor): + loss_mask = list(loss_mask) + return Sample( + index=index, + group_index=group_index, + rollout_id=index, + tokens=[100 + index, 200 + index, 300 + index], + response_length=len(loss_mask), + loss_mask=loss_mask, + weight_versions=[weight_version], + rollout_log_probs=[-0.1] * len(loss_mask), + ) + + +def _reports(groups, *, policy_version="8"): + return [ + { + "sample_index": sample.index, + "group_index": sample.group_index, + "valid_tokens": sum(sample.loss_mask), + "gate_passed": True, + "policy_version": policy_version, + "candidate_cache_bytes": sample.response_length * 4, + } + for group in groups + for sample in group + ] + + +def _real_lifecycle_manager(groups, *, cache_limit=1 << 20): + manager = object.__new__(RolloutManager.__ray_actor_class__) + manager.args = SimpleNamespace( + rollout_batch_size=len(groups), + n_samples_per_prompt=len(groups[0]), + rs_refill_max_rounds=2, + rs_refill_rpc_timeout_seconds=123.0, + rs_refill_max_candidate_cache_bytes=cache_limit, + save_debug_rollout_data=None, + ) + manager.train_parallel_config = { + "dp_size": 1, + "vpp_size": 1, + "microbatch_group_size_per_vp_stage": 1, + } + manager._pending_rs_batches = { + 7: { + "accepted": [], + "unscored": groups, + "initial_candidate_count": len(groups), + "round": 0, + "seen_sample_indices": set(), + "seen_group_indices": set(), + "seen_rollout_ids": {sample.rollout_id for group in groups for sample in group}, + "awaiting_log_prob_indices": None, + "awaiting_log_prob_bytes": None, + "proximal_log_probs_by_sample_index": {}, + "retained_logprob_cache_bytes": 0, + "accepted_mask_fingerprints": {}, + "metrics": {"rollout/source_groups": len(groups)}, + "initial_generation_seconds": 0.5, + } + } + conversions = [] + splits = [] + debug_saves = [] + + def convert(samples, preflight=False): + conversions.append((preflight, [sample.index for sample in samples])) + return { + "rollout_ids": [sample.rollout_id for sample in samples], + "sample_indices": [sample.index for sample in samples], + } + + def split(data, global_batch_size=None): + splits.append(global_batch_size) + return {**data, "split_global_batch_size": global_batch_size} + + manager._convert_samples_to_train_data = convert + manager._split_train_data_by_dp = split + return manager, conversions, splits, debug_saves + + +def _manager_rpc(manager, events, snapshots): + def finalize(rollout_id, coordinator_seconds): + pending = manager._pending_rs_batches[rollout_id] + snapshots.append( + ( + "finalize", + pending["awaiting_log_prob_indices"], + sorted(pending["proximal_log_probs_by_sample_index"]), + ) + ) + return manager.finalize_rs_batch(rollout_id, coordinator_seconds) + + def abort(rollout_id): + pending = manager._pending_rs_batches[rollout_id] + awaiting = pending["awaiting_log_prob_indices"] + snapshots.append( + ( + "abort", + None if awaiting is None else list(awaiting), + sorted(pending["proximal_log_probs_by_sample_index"]), + ) + ) + return manager.abort_rs_batch(rollout_id) + + return SimpleNamespace( + prepare_rs_candidate_data=_RemoteMethod("prepare", manager.prepare_rs_candidate_data, events), + apply_rs_candidate_reports=_RemoteMethod("apply", manager.apply_rs_candidate_reports, events), + store_rs_accepted_log_probs=_RemoteMethod("store", manager.store_rs_accepted_log_probs, events), + generate_rs_replacement_candidates=_RemoteMethod( + "generate", manager.generate_rs_replacement_candidates, events + ), + finalize_rs_batch=_RemoteMethod("finalize", finalize, events), + abort_rs_batch=_RemoteMethod("abort", abort, events), + ) + + +def test_rollout_manager_generates_exact_initial_and_aligned_replacement_counts(): + manager = object.__new__(RolloutManager.__ray_actor_class__) + manager.args = SimpleNamespace( + ci_test=False, + use_fault_tolerance=False, + rollout_batch_size=8, + n_samples_per_prompt=2, + rs_refill_max_rounds=2, + ) + manager.train_parallel_config = { + "dp_size": 2, + "vpp_size": 2, + "microbatch_group_size_per_vp_stage": 4, + } + manager._pending_rs_batches = {} + manager.health_monitoring_resume = lambda: None + requested_counts = [] + + def generate(_rollout_id, group_count, *, known_rollout_ids=None): + requested_counts.append((group_count, known_rollout_ids)) + return [[SimpleNamespace(), SimpleNamespace()] for _ in range(group_count)], {}, set() + + manager._call_rollout_for_group_count = generate + + manager._generate_rs_candidates(7) + assert requested_counts == [(8, None)] + + manager._pending_rs_batches[7] = { + "accepted": [[SimpleNamespace()] for _ in range(4)], + "unscored": [], + "round": 0, + "awaiting_log_prob_indices": None, + "awaiting_log_prob_bytes": None, + "seen_rollout_ids": {"initial"}, + "metrics": {}, + } + manager.generate_rs_replacement_candidates(7) + + assert requested_counts == [(8, None), (4, {"initial"})] + assert manager._pending_rs_batches[7]["round"] == 1 + + +def test_rollout_manager_isolates_custom_rollout_nested_argument_mutation(): + manager = object.__new__(RolloutManager.__ray_actor_class__) + manager.args = SimpleNamespace( + rollout_batch_size=2, + over_sampling_batch_size=2, + n_samples_per_prompt=2, + apply_chat_template_kwargs={"nested": [1]}, + rollout_sample_hook_path=["custom.sample_hook"], + ) + manager.data_source = object() + observed_args = [] + + def custom_rollout(args, rollout_id, data_source, *, evaluation): + assert rollout_id == 7 + assert data_source is manager.data_source + assert evaluation is False + assert args.rollout_batch_size == 1 + assert args.over_sampling_batch_size == 1 + observed_args.append(args) + + # A custom rollout and the sample hooks it calls share this namespace. + args.rollout_batch_size = 99 + args.apply_chat_template_kwargs["nested"].append(2) + args.rollout_sample_hook_path.append("custom.second_hook") + return [[_sample(0, 0), _sample(1, 0)]] + + manager.generate_rollout = custom_rollout + + groups, metrics, rollout_ids = manager._call_rollout_for_group_count(7, 1) + + assert [[sample.index for sample in group] for group in groups] == [[0, 1]] + assert metrics == {} + assert rollout_ids == {0, 1} + assert observed_args[0] is not manager.args + assert manager.args.rollout_batch_size == 2 + assert manager.args.over_sampling_batch_size == 2 + assert manager.args.apply_chat_template_kwargs == {"nested": [1]} + assert manager.args.rollout_sample_hook_path == ["custom.sample_hook"] + + +def test_rollout_manager_rejects_an_unaligned_final_effective_batch(): + manager = object.__new__(RolloutManager.__ray_actor_class__) + manager.args = SimpleNamespace(rs_batch_refill=True, rollout_batch_size=5, n_samples_per_prompt=2) + topology = {"dp_size": 2, "vpp_size": 2, "microbatch_group_size_per_vp_stage": 4} + + with pytest.raises(ValueError, match=r"effective sample count.*10, alignment = 8"): + manager.set_train_parallel_config(topology) + + assert not hasattr(manager, "train_parallel_config") + + +def test_rollout_manager_real_lifecycle_applies_stores_and_finalizes(monkeypatch): + groups = [ + [_sample(0, 0), _sample(1, 0)], + [_sample(10, 1, loss_mask=[1, 0]), _sample(11, 1, loss_mask=[1, 0])], + ] + manager, conversions, splits, debug_saves = _real_lifecycle_manager(groups, cache_limit=32) + manager.args.custom_rollout_log_function_path = "custom.read_only_log" + manager.args.global_batch_size = 4 + manager.args.apply_chat_template_kwargs = {"nested": [1]} + cache = { + sample.index: torch.tensor([sample.index + 0.25, sample.index + 0.5]) for group in groups for sample in group + } + actor = _LifecycleActor( + _reports(groups), + [{0: cache[0], 10: cache[10]}, {1: cache[1], 11: cache[11]}], + ) + rpc_events = [] + state_snapshots = [] + rollout_logs = [] + rollout_log_args = [] + manager_rpc = _manager_rpc(manager, rpc_events, state_snapshots) + + manager_timeouts = [] + + def manager_get(value, *, timeout=None): + manager_timeouts.append(timeout) + return value + + monkeypatch.setattr(rollout_module.ray, "get", manager_get) + monkeypatch.setattr(rollout_module.time, "perf_counter", lambda: 100.0) + monkeypatch.setattr( + rollout_module, + "save_debug_rollout_data", + lambda path_template, samples, *, rollout_id, evaluation: debug_saves.append( + (path_template, rollout_id, evaluation, [sample.index for sample in samples]) + ), + ) + + def log_rollout(rollout_id, args, samples, metrics, elapsed): + rollout_log_args.append(args) + args.rollout_batch_size = 99 + args.global_batch_size = 99 + args.apply_chat_template_kwargs["nested"].append(2) + rollout_logs.append((rollout_id, [sample.index for sample in samples], dict(metrics), elapsed)) + + monkeypatch.setattr(rollout_module, "log_rollout_data", log_rollout) + + result = run_rs_batch_refill( + actor, + manager_rpc, + 7, + resolve=lambda value, **_kwargs: value, + clock=iter([0.0, 1.0, 2.0, 3.0]).__next__, + ) + + assert [name for name, _ in rpc_events] == ["prepare", "apply", "store", "finalize"] + assert actor.events[0] == ( + "score", + 7, + { + "rollout_ids": [0, 1, 10, 11], + "sample_indices": [0, 1, 10, 11], + "split_global_batch_size": 4, + }, + ) + assert actor.events[1] == ("take", 7, [0, 1, 10, 11]) + assert state_snapshots == [("finalize", None, [0, 1, 10, 11])] + assert conversions == [ + (True, [0, 1, 10, 11]), + (False, [0, 1, 10, 11]), + (False, [0, 1, 10, 11]), + ] + assert splits == [4, None] + assert debug_saves == [(None, 7, False, [0, 1, 10, 11])] + assert result["sample_indices"] == [0, 1, 10, 11] + assert result["split_global_batch_size"] is None + assert len(rollout_log_args) == 1 + assert rollout_log_args[0] is not manager.args + assert manager.args.rollout_batch_size == 2 + assert manager.args.global_batch_size == 4 + assert manager.args.apply_chat_template_kwargs == {"nested": [1]} + for actual, index in zip(result["rs_preflight_log_probs"], [0, 1, 10, 11], strict=True): + torch.testing.assert_close(actual, cache[index]) + + assert 7 not in manager._pending_rs_batches + assert actor.cache_live is False + assert [event[0] for event in actor.events] == ["score", "take"] + assert len(rollout_logs) == 1 + rollout_id, sample_indices, metrics, elapsed = rollout_logs[0] + assert rollout_id == 7 + assert sample_indices == [0, 1, 10, 11] + assert elapsed == 3.5 + assert metrics["rollout/rs_refill/initial_policy_staleness"] == 1 + assert metrics["rollout/rs_refill/candidate_groups"] == 2 + assert metrics["rollout/rs_refill/scored_groups"] == 2 + assert metrics["rollout/rs_refill/rejected_groups"] == 0 + assert metrics["rollout/rs_refill/surplus_groups"] == 0 + assert metrics["rollout/rs_refill/rounds"] == 0 + assert metrics["rollout/rs_refill/accepted_groups"] == 2 + assert metrics["rollout/rs_refill/scored_trainable_tokens"] == 6 + assert metrics["rollout/rs_refill/aggregate_candidate_logprob_cache_bytes"] == 32 + assert metrics["rollout/rs_refill/peak_aggregate_candidate_logprob_cache_bytes"] == 32 + assert metrics["rollout/rs_refill/peak_actor_candidate_logprob_cache_bytes"] == 32 + assert metrics["rollout/rs_refill/logprob_cache_limit_bytes"] == 32 + assert metrics["rollout/rs_refill/selected_logprob_transfer_bytes"] == 32 + assert metrics["rollout/rs_refill/retained_logprob_cache_bytes"] == 32 + assert metrics["rollout/rs_refill/peak_retained_logprob_cache_bytes"] == 32 + assert metrics["rollout/rs_refill/effective_trainable_tokens"] == 6 + assert metrics["rollout/rs_refill/gate_acceptance_rate"] == 1 + assert metrics["rollout/rs_refill/selection_utilization"] == 1 + assert manager_timeouts == [123.0, 123.0] + + +def test_rollout_manager_rejects_custom_logger_training_input_mutation(monkeypatch): + groups = [[_sample(0, 0), _sample(1, 0)]] + manager, _, _, _ = _real_lifecycle_manager(groups) + manager.args.custom_rollout_log_function_path = "custom.log" + original_convert = manager._convert_samples_to_train_data + + def convert(samples, preflight=False): + data = original_convert(samples, preflight=preflight) + data["tokens"] = [list(sample.tokens) for sample in samples] + return data + + manager._convert_samples_to_train_data = convert + actor = _LifecycleActor( + _reports(groups), + [{sample.index: torch.zeros(sample.response_length) for sample in groups[0]}], + ) + rpc_events = [] + state_snapshots = [] + manager_rpc = _manager_rpc(manager, rpc_events, state_snapshots) + monkeypatch.setattr(rollout_module.ray, "get", lambda value, *, timeout=None: value) + monkeypatch.setattr(rollout_module.time, "perf_counter", lambda: 100.0) + monkeypatch.setattr(rollout_module, "save_debug_rollout_data", lambda *_args, **_kwargs: None) + + def mutate_tokens(_rollout_id, _args, samples, _metrics, _elapsed): + samples[0].tokens[-1] += 1 + + monkeypatch.setattr(rollout_module, "log_rollout_data", mutate_tokens) + + with pytest.raises(RuntimeError, match="rollout observability hooks must be read-only"): + run_rs_batch_refill( + actor, + manager_rpc, + 7, + resolve=lambda value, **_kwargs: value, + clock=iter([0.0, 1.0, 2.0, 3.0]).__next__, + ) + + assert [name for name, _ in rpc_events] == ["prepare", "apply", "store", "finalize", "abort"] + assert actor.events[-1] == ("discard", 7) + assert state_snapshots[-1] == ("abort", None, [0, 1]) + assert 7 not in manager._pending_rs_batches + + +def test_rollout_manager_rejects_debug_saver_training_input_mutation(monkeypatch): + groups = [[_sample(0, 0), _sample(1, 0)]] + manager, _, _, _ = _real_lifecycle_manager(groups) + manager.args.save_debug_rollout_data = "/tmp/rollout-{rollout_id}.pt" + original_convert = manager._convert_samples_to_train_data + + def convert(samples, preflight=False): + data = original_convert(samples, preflight=preflight) + data["tokens"] = [list(sample.tokens) for sample in samples] + return data + + manager._convert_samples_to_train_data = convert + actor = _LifecycleActor( + _reports(groups), + [{sample.index: torch.zeros(sample.response_length) for sample in groups[0]}], + ) + rpc_events = [] + state_snapshots = [] + manager_rpc = _manager_rpc(manager, rpc_events, state_snapshots) + monkeypatch.setattr(rollout_module.ray, "get", lambda value, *, timeout=None: value) + monkeypatch.setattr(rollout_module.time, "perf_counter", lambda: 100.0) + + def mutate_tokens(_path, samples, *, rollout_id, evaluation): + samples[0].tokens[-1] += 1 + + monkeypatch.setattr(rollout_module, "save_debug_rollout_data", mutate_tokens) + monkeypatch.setattr(rollout_module, "log_rollout_data", lambda *_args, **_kwargs: None) + + with pytest.raises(RuntimeError, match="rollout observability hooks must be read-only"): + run_rs_batch_refill( + actor, + manager_rpc, + 7, + resolve=lambda value, **_kwargs: value, + clock=iter([0.0, 1.0, 2.0, 3.0]).__next__, + ) + + assert [name for name, _ in rpc_events] == ["prepare", "apply", "store", "finalize", "abort"] + assert actor.events[-1] == ("discard", 7) + assert state_snapshots[-1] == ("abort", None, [0, 1]) + assert 7 not in manager._pending_rs_batches + + +def test_rollout_manager_real_lifecycle_refills_a_rejected_group_and_accepts_tensor_masks(monkeypatch): + initial = [ + [_sample(0, 0), _sample(1, 0)], + [_sample(10, 1), _sample(11, 1)], + ] + replacement = [ + [ + _sample(20, 2, weight_version="8", loss_mask=torch.tensor([1, 0])), + _sample(21, 2, weight_version="8", loss_mask=[1, 0]), + ] + ] + initial_reports = _reports(initial) + for report in initial_reports: + if report["group_index"] == 1: + report["gate_passed"] = False + cache_rounds = [ + { + sample.index: torch.full((sample.response_length,), float(sample.index)) + for group in initial + for sample in group + }, + { + sample.index: torch.full((sample.response_length,), float(sample.index)) + for group in replacement + for sample in group + }, + ] + actor = _RoundLifecycleActor([initial_reports, _reports(replacement)], cache_rounds) + manager, conversions, splits, debug_saves = _real_lifecycle_manager(initial) + generated = [] + + def generate(rollout_id, group_count, *, known_rollout_ids=None): + generated.append((rollout_id, group_count, set(known_rollout_ids))) + return replacement, {"rollout/replacement_marker": 1}, {20, 21} + + manager._call_rollout_for_group_count = generate + rpc_events = [] + state_snapshots = [] + rollout_logs = [] + manager_rpc = _manager_rpc(manager, rpc_events, state_snapshots) + manager_timeouts = [] + + def manager_get(value, *, timeout=None): + manager_timeouts.append(timeout) + return value + + monkeypatch.setattr(rollout_module.ray, "get", manager_get) + monkeypatch.setattr(rollout_module.time, "perf_counter", lambda: 100.0) + monkeypatch.setattr( + rollout_module, + "save_debug_rollout_data", + lambda path_template, samples, *, rollout_id, evaluation: debug_saves.append( + (path_template, rollout_id, evaluation, [sample.index for sample in samples]) + ), + ) + monkeypatch.setattr( + rollout_module, + "log_rollout_data", + lambda rollout_id, args, samples, metrics, elapsed: rollout_logs.append( + (rollout_id, [sample.index for sample in samples], dict(metrics), elapsed) + ), + ) + + result = run_rs_batch_refill( + actor, + manager_rpc, + 7, + resolve=lambda value, **_kwargs: value, + clock=iter([0.0, 1.0, 2.0, 3.0, 4.0, 5.0]).__next__, + ) + + assert [name for name, _ in rpc_events] == [ + "prepare", + "apply", + "store", + "generate", + "prepare", + "apply", + "store", + "finalize", + ] + assert [event[0] for event in actor.events] == ["score", "take", "score", "take"] + assert generated == [(7, 1, {0, 1, 10, 11})] + assert result["sample_indices"] == [0, 1, 20, 21] + assert [tensor.tolist() for tensor in result["rs_preflight_log_probs"]] == [ + [0.0, 0.0], + [1.0, 1.0], + [20.0, 20.0], + [21.0, 21.0], + ] + assert conversions == [ + (True, [0, 1, 10, 11]), + (True, [20, 21]), + (False, [0, 1, 20, 21]), + ] + assert splits == [4, 2, None] + assert debug_saves == [(None, 7, False, [0, 1, 20, 21])] + assert state_snapshots == [("finalize", None, [0, 1, 20, 21])] + assert manager_timeouts == [123.0, 123.0, 123.0, 123.0] + metrics = rollout_logs[0][2] + assert metrics["rollout/rs_refill/rounds"] == 1 + assert metrics["rollout/rs_refill/scored_groups"] == 3 + assert metrics["rollout/rs_refill/rejected_groups"] == 1 + assert metrics["rollout/rs_refill/generated_replacement_groups"] == 1 + assert metrics["rollout/rs_refill/effective_trainable_tokens"] == 6 + assert metrics["rollout/rs_refill/aggregate_candidate_logprob_cache_bytes"] == 48 + assert metrics["rollout/rs_refill/peak_aggregate_candidate_logprob_cache_bytes"] == 32 + assert metrics["rollout/rs_refill/peak_actor_candidate_logprob_cache_bytes"] == 32 + assert metrics["rollout/rs_refill/selected_logprob_transfer_bytes"] == 32 + assert metrics["rollout/rs_refill/retained_logprob_cache_bytes"] == 32 + assert metrics["rollout/rs_refill/peak_retained_logprob_cache_bytes"] == 32 + assert metrics["rollout/rs_refill/replacement_round_1/rollout/replacement_marker"] == 1 + assert 7 not in manager._pending_rs_batches + + +def test_rollout_manager_real_lifecycle_aborts_after_cache_rpc_failure(monkeypatch): + groups = [ + [_sample(0, 0), _sample(1, 0)], + [_sample(10, 1), _sample(11, 1)], + ] + manager, conversions, splits, debug_saves = _real_lifecycle_manager(groups) + failed_cache_ref = object() + actor = _LifecycleActor(_reports(groups), failed_cache_ref) + rpc_events = [] + state_snapshots = [] + manager_rpc = _manager_rpc(manager, rpc_events, state_snapshots) + + def get(value, *, timeout=None): + assert timeout == 123.0 + if value is failed_cache_ref: + raise RuntimeError("cache RPC failed") + return value + + monkeypatch.setattr(rollout_module.ray, "get", get) + monkeypatch.setattr(rollout_module.time, "perf_counter", lambda: 100.0) + + with pytest.raises(RuntimeError, match="cache RPC failed"): + run_rs_batch_refill( + actor, + manager_rpc, + 7, + resolve=lambda value, **_kwargs: value, + clock=iter([0.0, 1.0, 2.0]).__next__, + ) + + assert [name for name, _ in rpc_events] == ["prepare", "apply", "store", "abort"] + assert actor.events == [ + ( + "score", + 7, + { + "rollout_ids": [0, 1, 10, 11], + "sample_indices": [0, 1, 10, 11], + "split_global_batch_size": 4, + }, + ), + ("take", 7, [0, 1, 10, 11]), + ("discard", 7), + ] + assert state_snapshots == [("abort", [0, 1, 10, 11], [])] + assert 7 not in manager._pending_rs_batches + assert manager.abort_rs_batch(7) is False + assert actor.cache_live is False + assert conversions == [(True, [0, 1, 10, 11])] + assert splits == [4] + assert debug_saves == [] + + +def test_rollout_manager_rejects_aggregate_selected_cache_before_actor_transfer(monkeypatch): + groups = [ + [_sample(0, 0), _sample(1, 0)], + [_sample(10, 1), _sample(11, 1)], + ] + manager, _, _, _ = _real_lifecycle_manager(groups, cache_limit=16) + reports = _reports(groups) + actor = _LifecycleActor([reports[:2], reports[2:]], [{}, {}]) + rpc_events = [] + state_snapshots = [] + manager_rpc = _manager_rpc(manager, rpc_events, state_snapshots) + monkeypatch.setattr(rollout_module.ray, "get", lambda value, *, timeout=None: value) + + with pytest.raises(RuntimeError, match=r"before transfer.*required=32, limit=16"): + run_rs_batch_refill( + actor, + manager_rpc, + 7, + resolve=lambda value, **_kwargs: value, + clock=iter([0.0, 1.0, 2.0]).__next__, + ) + + assert [name for name, _ in rpc_events] == ["prepare", "apply", "abort"] + assert [event[0] for event in actor.events] == ["score", "discard"] + assert state_snapshots == [("abort", None, [])] + assert 7 not in manager._pending_rs_batches + + +def test_rollout_manager_rejects_cross_round_retained_cache_before_second_transfer(monkeypatch): + initial = [ + [_sample(0, 0), _sample(1, 0)], + [_sample(10, 1), _sample(11, 1)], + ] + replacement = [[_sample(20, 2, weight_version="8"), _sample(21, 2, weight_version="8")]] + initial_reports = _reports(initial) + for report in initial_reports: + if report["group_index"] == 1: + report["gate_passed"] = False + cache_rounds = [ + {sample.index: torch.zeros(sample.response_length) for group in initial for sample in group}, + {sample.index: torch.zeros(sample.response_length) for group in replacement for sample in group}, + ] + replacement_reports = _reports(replacement) + actor = _RoundLifecycleActor( + [[initial_reports[:2], initial_reports[2:]], [replacement_reports]], + cache_rounds, + ) + manager, _, _, _ = _real_lifecycle_manager(initial, cache_limit=24) + manager._call_rollout_for_group_count = lambda *_args, **_kwargs: (replacement, {}, {20, 21}) + rpc_events = [] + state_snapshots = [] + manager_rpc = _manager_rpc(manager, rpc_events, state_snapshots) + monkeypatch.setattr(rollout_module.ray, "get", lambda value, *, timeout=None: value) + + with pytest.raises(RuntimeError, match=r"before transfer.*retained=16, incoming=16.*limit=24"): + run_rs_batch_refill( + actor, + manager_rpc, + 7, + resolve=lambda value, **_kwargs: value, + clock=iter([0.0, 1.0, 2.0, 3.0, 4.0, 5.0]).__next__, + ) + + assert [name for name, _ in rpc_events] == ["prepare", "apply", "store", "generate", "prepare", "apply", "abort"] + assert [event[0] for event in actor.events] == ["score", "take", "score", "discard"] + assert state_snapshots == [("abort", None, [0, 1])] + assert 7 not in manager._pending_rs_batches + + +@pytest.mark.parametrize( + "invalid_tensor", + [ + torch.zeros(2, dtype=torch.float64), + torch.zeros(2, dtype=torch.int32), + torch.zeros((2, 1), dtype=torch.float32), + torch.zeros(4, dtype=torch.float32)[::2], + torch.zeros(1, dtype=torch.float32), + torch.zeros(2, dtype=torch.float32, device="meta"), + [0.0, 0.0], + ], + ids=["float64", "int32", "two-dimensional", "non-contiguous", "wrong-length", "non-cpu", "non-tensor"], +) +def test_rollout_manager_rejects_invalid_transferred_cache_contract(monkeypatch, invalid_tensor): + groups = [[_sample(0, 0), _sample(1, 0)]] + manager, _, _, _ = _real_lifecycle_manager(groups) + cache = {sample.index: invalid_tensor for sample in groups[0]} + actor = _LifecycleActor(_reports(groups), [cache]) + rpc_events = [] + state_snapshots = [] + manager_rpc = _manager_rpc(manager, rpc_events, state_snapshots) + monkeypatch.setattr(rollout_module.ray, "get", lambda value, *, timeout=None: value) + + with pytest.raises(RuntimeError, match=r"contiguous one-dimensional CPU float32 tensors"): + run_rs_batch_refill( + actor, + manager_rpc, + 7, + resolve=lambda value, **_kwargs: value, + clock=iter([0.0, 1.0, 2.0, 3.0]).__next__, + ) + + assert [name for name, _ in rpc_events] == ["prepare", "apply", "store", "abort"] + assert [event[0] for event in actor.events] == ["score", "take", "discard"] + assert state_snapshots == [("abort", [0, 1], [])] + assert 7 not in manager._pending_rs_batches + + +def test_rollout_manager_rejects_underreported_cache_before_actor_transfer(monkeypatch): + groups = [[_sample(0, 0), _sample(1, 0)]] + reports = _reports(groups) + reports[0]["candidate_cache_bytes"] = 0 + manager, _, _, _ = _real_lifecycle_manager(groups) + actor = _LifecycleActor(reports, [{sample.index: torch.zeros(2) for sample in groups[0]}]) + rpc_events = [] + state_snapshots = [] + manager_rpc = _manager_rpc(manager, rpc_events, state_snapshots) + monkeypatch.setattr(rollout_module.ray, "get", lambda value, *, timeout=None: value) + + with pytest.raises(RuntimeError, match=r"reported=0, expected=8"): + run_rs_batch_refill( + actor, + manager_rpc, + 7, + resolve=lambda value, **_kwargs: value, + clock=iter([0.0, 1.0, 2.0]).__next__, + ) + + assert [name for name, _ in rpc_events] == ["prepare", "apply", "abort"] + assert [event[0] for event in actor.events] == ["score", "discard"] + assert state_snapshots == [("abort", None, [])] + assert 7 not in manager._pending_rs_batches + + +def test_rollout_manager_reports_distinct_aggregate_and_per_actor_cache_peaks(monkeypatch): + groups = [ + [_sample(0, 0), _sample(1, 0)], + [_sample(10, 1), _sample(11, 1)], + ] + reports = _reports(groups) + manager, _, _, _ = _real_lifecycle_manager(groups, cache_limit=32) + monkeypatch.setattr(rollout_module.ray, "get", lambda value, *, timeout=None: value) + + result = manager.apply_rs_candidate_reports(7, [reports[:2], reports[2:]], 0.5) + + assert result["complete"] is True + metrics = manager._pending_rs_batches[7]["metrics"] + assert metrics["rollout/rs_refill/aggregate_candidate_logprob_cache_bytes"] == 32 + assert metrics["rollout/rs_refill/peak_aggregate_candidate_logprob_cache_bytes"] == 32 + assert metrics["rollout/rs_refill/peak_actor_candidate_logprob_cache_bytes"] == 16 + + +@pytest.mark.parametrize("invalid_bytes", [True, -1, 1.5, None]) +def test_rollout_manager_rejects_invalid_candidate_cache_byte_reports(monkeypatch, invalid_bytes): + groups = [[_sample(0, 0), _sample(1, 0)]] + reports = _reports(groups) + reports[0]["candidate_cache_bytes"] = invalid_bytes + manager, _, _, _ = _real_lifecycle_manager(groups) + monkeypatch.setattr(rollout_module.ray, "get", lambda value, *, timeout=None: value) + + with pytest.raises(RuntimeError, match="non-negative integers"): + manager.apply_rs_candidate_reports(7, [reports], 0.5) + + pending = manager._pending_rs_batches[7] + assert pending["accepted"] == [] + assert pending["awaiting_log_prob_indices"] is None + assert pending["awaiting_log_prob_bytes"] is None + assert pending["proximal_log_probs_by_sample_index"] == {} + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__])) diff --git a/tests/test_rs_exact_refill_sglang_rollout.py b/tests/test_rs_exact_refill_sglang_rollout.py new file mode 100644 index 0000000000..b32a8fb555 --- /dev/null +++ b/tests/test_rs_exact_refill_sglang_rollout.py @@ -0,0 +1,78 @@ +import asyncio +from types import SimpleNamespace + +import pytest + +from slime.rollout import sglang_rollout + +NUM_GPUS = 0 + + +def _make_state(monkeypatch, default_batch_size): + args = SimpleNamespace( + hf_checkpoint="checkpoint", + n_samples_per_prompt=1, + rollout_max_response_len=128, + rollout_skip_special_tokens=False, + rollout_stop=None, + rollout_stop_token_ids=None, + rollout_temperature=1.0, + rollout_top_k=-1, + rollout_top_p=1.0, + sglang_dp_size=1, + sglang_enable_deterministic_inference=False, + sglang_server_concurrency=1, + ) + monkeypatch.setattr(sglang_rollout, "load_tokenizer", lambda *_args, **_kwargs: object()) + monkeypatch.setattr(sglang_rollout, "load_processor", lambda *_args, **_kwargs: None) + monkeypatch.setattr(sglang_rollout, "get_rollout_num_engines", lambda _args: 1) + state = object.__new__(sglang_rollout.GenerateState) + state.__init__(args) + state.args.rollout_batch_size = default_batch_size + return state + + +def test_generate_state_preserves_legacy_submit_call(monkeypatch): + state = _make_state(monkeypatch, default_batch_size=8) + seen = [] + + async def generate(args, group, **_kwargs): + seen.append((args.rollout_batch_size, group)) + return group + + monkeypatch.setattr(sglang_rollout, "generate_and_rm_group", generate) + + async def run_call(): + state.submit_generate_tasks([["initial"]]) + await asyncio.gather(*state.pendings) + + asyncio.run(run_call()) + + assert state.args.rollout_batch_size == 8 + assert seen == [(8, ["initial"])] + + +def test_generate_state_uses_each_calls_refill_argument_override(monkeypatch): + state = _make_state(monkeypatch, default_batch_size=8) + seen = [] + + async def generate(args, group, **_kwargs): + seen.append((args.rollout_batch_size, group)) + return group + + monkeypatch.setattr(sglang_rollout, "generate_and_rm_group", generate) + + async def run_calls(): + state.submit_generate_tasks([["initial"]]) + await asyncio.gather(*state.pendings) + state.pendings.clear() + state.submit_generate_tasks([["replacement"]], args=SimpleNamespace(rollout_batch_size=2)) + await asyncio.gather(*state.pendings) + + asyncio.run(run_calls()) + + assert seen == [(8, ["initial"]), (2, ["replacement"])] + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__])) diff --git a/tests/test_rs_exact_refill_train_async.py b/tests/test_rs_exact_refill_train_async.py new file mode 100644 index 0000000000..445711b420 --- /dev/null +++ b/tests/test_rs_exact_refill_train_async.py @@ -0,0 +1,107 @@ +import importlib.util +import sys +import types +from pathlib import Path +from types import SimpleNamespace +from unittest import mock + +import pytest + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__])) + +_TRAIN_ASYNC_PATH = Path(__file__).resolve().parents[1] / "train_async.py" +_TRAIN_ASYNC_SPEC = importlib.util.spec_from_file_location("test_train_async_module", _TRAIN_ASYNC_PATH) +assert _TRAIN_ASYNC_SPEC is not None and _TRAIN_ASYNC_SPEC.loader is not None +train_async = importlib.util.module_from_spec(_TRAIN_ASYNC_SPEC) +_ARGUMENTS_STUB = types.ModuleType("slime.utils.arguments") +_ARGUMENTS_STUB.parse_args = lambda: None +# This CPU scheduling contract does not need the SGLang-backed CLI parser. +with mock.patch.dict(sys.modules, {"slime.utils.arguments": _ARGUMENTS_STUB}): + _TRAIN_ASYNC_SPEC.loader.exec_module(train_async) + +NUM_GPUS = 0 + + +class _RemoteMethod: + def __init__(self, fn): + self._fn = fn + + def remote(self, *args, **kwargs): + return self._fn(*args, **kwargs) + + +class _RolloutManager: + def __init__(self, events): + self.generate = _RemoteMethod(lambda rollout_id: self._record(events, "generate", rollout_id)) + self.save = _RemoteMethod(lambda rollout_id: self._record(events, "manager_save", rollout_id)) + self.eval = _RemoteMethod(lambda rollout_id: self._record(events, "eval", rollout_id)) + self.dispose = _RemoteMethod(lambda: self._record(events, "dispose", None)) + + @staticmethod + def _record(events, name, rollout_id): + events.append((name, rollout_id)) + return (name, rollout_id) + + +class _ActorModel: + def __init__(self, events): + self.events = events + + def update_weights(self): + self.events.append(("update_weights", None)) + + def async_train(self, rollout_id, _data, external_data=None): + assert external_data is None + self.events.append(("train", rollout_id)) + return ("train_ref", rollout_id) + + def save_model(self, rollout_id, *, force_sync): + self.events.append(("actor_save", rollout_id, force_sync)) + + +def test_refill_checkpoint_finishes_before_the_next_rollout(monkeypatch): + events = [] + manager = _RolloutManager(events) + actor = _ActorModel(events) + + def resolve(value): + if isinstance(value, tuple) and value[0] == "generate": + return value[1] + return None + + monkeypatch.setattr(train_async.ray, "get", resolve) + monkeypatch.setattr(train_async, "create_placement_groups", lambda _args: {"rollout": object()}) + monkeypatch.setattr(train_async, "create_rollout_manager", lambda _args, _pg: (manager, None)) + monkeypatch.setattr(train_async, "create_training_models", lambda _args, _pgs, _manager: (actor, None)) + monkeypatch.setattr(train_async, "configure_logger", lambda: None) + monkeypatch.setattr(train_async, "init_tracking", lambda _args: None) + monkeypatch.setattr(train_async, "finish_tracking", lambda _args: None) + monkeypatch.setattr( + train_async, + "run_rs_batch_refill", + lambda _actor, _manager, rollout_id, **_kwargs: ("train_data", rollout_id), + ) + + args = SimpleNamespace( + colocate=False, + release_train=False, + check_weight_update_equal=False, + start_rollout_id=0, + num_rollout=2, + rs_batch_refill=True, + rs_refill_rpc_timeout_seconds=123.0, + use_critic=False, + num_critic_only_steps=0, + save_interval=1, + rollout_global_dataset=True, + update_weights_interval=1, + eval_interval=None, + ) + + train_async.train(args) + + manager_save = events.index(("manager_save", 0)) + next_generate = events.index(("generate", 1)) + updates_before_next = [i for i, event in enumerate(events[:next_generate]) if event[0] == "update_weights"] + assert manager_save < updates_before_next[-1] < next_generate diff --git a/tests/test_rs_exact_refill_utils.py b/tests/test_rs_exact_refill_utils.py new file mode 100644 index 0000000000..117ec03401 --- /dev/null +++ b/tests/test_rs_exact_refill_utils.py @@ -0,0 +1,756 @@ +import copy +import importlib.util +from pathlib import Path +from types import SimpleNamespace + +import numpy as np +import pytest +import torch + +from slime.utils.rs_refill import ( + apply_rs_refill_tis, + attach_proximal_log_probs, + clone_rs_masks, + compute_sequence_rs_masks, + fingerprint_rs_train_data, + get_rs_refill_candidate_group_multiple, + merge_replacement_metrics, + merge_selected_log_prob_caches, + plan_topology_aligned_rs_refill, + run_rs_batch_refill, + select_accepted_groups, + snapshot_sample_masks, + validate_final_rs_masks, + validate_initial_policy_staleness, + validate_refill_rollout_ids, + validate_replacement_policy_version, + validate_rs_refill_target_batch_alignment, + validate_rs_train_data_fingerprint, + validate_sample_masks, +) + +NUM_GPUS = 0 + +_MIS_PATH = Path(__file__).resolve().parents[1] / "examples" / "train_infer_mismatch_helper" / "mis.py" +_MIS_SPEC = importlib.util.spec_from_file_location("test_rs_exact_refill_existing_mis", _MIS_PATH) +assert _MIS_SPEC is not None and _MIS_SPEC.loader is not None +_MIS_MODULE = importlib.util.module_from_spec(_MIS_SPEC) +_MIS_SPEC.loader.exec_module(_MIS_MODULE) + + +def _args(**overrides): + values = { + "n_samples_per_prompt": 2, + "rollout_batch_size": 8, + "tis_level": "token", + "tis_mode": "clip", + "tis_lower_bound": 0.5, + "tis_upper_bound": 2.0, + "tis_batch_normalize": False, + "rs_level": "geometric", + "rs_lower_bound": 0.6, + "rs_upper_bound": 1.5, + "rs_veto_threshold": None, + } + values.update(overrides) + return SimpleNamespace(**values) + + +def _topology(**overrides): + values = { + "dp_size": 2, + "vpp_size": 2, + "microbatch_group_size_per_vp_stage": 4, + } + values.update(overrides) + return values + + +def test_train_data_fingerprint_covers_ordered_policy_inputs(): + train_data = { + "tokens": [[1, 2, 3], [4, 5]], + "sample_indices": [10, 11], + "rewards": [1.0, 0.0], + "raw_reward": [0.75, -0.25], + "loss_masks": [[1, 1], [1]], + "rollout_ids": [20, 21], + "rollout_log_probs": [torch.tensor([-0.1, -0.2]), torch.tensor([-0.3])], + "rollout_top_p_token_ids": [[7], [8]], + "rollout_top_p_token_offsets": [[0], [1]], + "rollout_routed_experts": [torch.tensor([0, 1]), torch.tensor([1])], + "multimodal_train_inputs": [{"pixel_values": torch.arange(4).reshape(2, 2)}, None], + "metadata": [{"loss": "policy"}, {"loss": "policy"}], + } + fingerprint_kwargs = {"group_indices": [0, 1], "weight_versions": [["7"], ["7"]]} + expected = fingerprint_rs_train_data(train_data, **fingerprint_kwargs) + + validate_rs_train_data_fingerprint(dict(reversed(train_data.items())), expected, **fingerprint_kwargs) + mutators = [ + lambda value: value["tokens"][0].__setitem__(2, 99), + lambda value: value["sample_indices"].reverse(), + lambda value: value["rewards"].__setitem__(0, 0.5), + lambda value: value["raw_reward"].__setitem__(1, 0.25), + lambda value: value["loss_masks"][0].__setitem__(1, 0), + lambda value: value["rollout_ids"].__setitem__(0, 99), + lambda value: value["rollout_log_probs"][0].add_(1), + lambda value: value["rollout_top_p_token_ids"][0].__setitem__(0, 9), + lambda value: value["rollout_top_p_token_offsets"][1].__setitem__(0, 0), + lambda value: value["rollout_routed_experts"][0].add_(1), + lambda value: value["multimodal_train_inputs"][0]["pixel_values"].add_(1), + lambda value: value["metadata"][0].__setitem__("loss", "other"), + ] + for mutate in mutators: + candidate = copy.deepcopy(train_data) + mutate(candidate) + with pytest.raises(RuntimeError, match="rollout observability hooks must be read-only"): + validate_rs_train_data_fingerprint(candidate, expected, **fingerprint_kwargs) + + with pytest.raises(RuntimeError, match="rollout observability hooks must be read-only"): + validate_rs_train_data_fingerprint( + train_data, + expected, + group_indices=[1, 0], + weight_versions=[["7"], ["7"]], + ) + with pytest.raises(RuntimeError, match="rollout observability hooks must be read-only"): + validate_rs_train_data_fingerprint( + train_data, + expected, + group_indices=[0, 1], + weight_versions=[["8"], ["7"]], + ) + + +def test_train_data_fingerprint_fails_closed_for_unsupported_values(): + with pytest.raises(TypeError, match=r"does not support value type builtins\.object"): + fingerprint_rs_train_data({"metadata": [object()]}, group_indices=[0], weight_versions=[["7"]]) + + +def test_train_data_fingerprint_preserves_numpy_scalar_dtype_and_tensor_layout(): + kwargs = {"group_indices": [0], "weight_versions": [["7"]]} + float32_digest = fingerprint_rs_train_data({"value": np.float32(1)}, **kwargs) + float64_digest = fingerprint_rs_train_data({"value": np.float64(1)}, **kwargs) + assert float32_digest != float64_digest + + indices = torch.tensor([[0, 1]]) + values = torch.tensor([1.0, 0.0]) + sparse = torch.sparse_coo_tensor(indices, values, (2,)) + dense = sparse.to_dense() + assert fingerprint_rs_train_data({"value": sparse}, **kwargs) != fingerprint_rs_train_data( + {"value": dense}, **kwargs + ) + + first_schema = np.zeros(1, dtype=[("a", "