diff --git a/tests/_unit_stubs.py b/tests/_unit_stubs.py index afc27ec5..23cf3136 100644 --- a/tests/_unit_stubs.py +++ b/tests/_unit_stubs.py @@ -43,19 +43,110 @@ def real_module_available(name: str) -> bool: return False +class _RayRemoteWrapper: + def __init__(self, target: Any): + self._target = target + if isinstance(target, type): + self.__ray_actor_class__ = target + + def options(self, *args, **kwargs): # noqa: ARG002 + return self + + def remote(self, *args, **kwargs): + return self._target(*args, **kwargs) + + +def _ray_remote(*args, **kwargs): # noqa: ARG001 + if len(args) == 1 and callable(args[0]) and not kwargs: + return _RayRemoteWrapper(args[0]) + + def decorator(target: Any): + return _RayRemoteWrapper(target) + + return decorator + + +class _RaySchedulingStrategy: + def __init__(self, **kwargs): + self.kwargs = kwargs + + +def _install_ray_module_stub() -> None: + ray_mod = types.ModuleType("ray") + ray_mod.__path__ = [] + ray_mod.get = lambda refs: refs + ray_mod.put = lambda value, **kwargs: value # noqa: ARG005 + ray_mod.wait = lambda refs, timeout=None: (refs, []) # noqa: ARG005 + ray_mod.kill = lambda actor, no_restart=False: None # noqa: ARG005 + ray_mod.nodes = lambda: [] + ray_mod.cluster_resources = lambda: {} + ray_mod.available_resources = lambda: {} + ray_mod.get_gpu_ids = lambda: [0] + ray_mod.remote = _ray_remote + ray_mod.ObjectRef = object + + actor_mod = types.ModuleType("ray.actor") + actor_mod.ActorHandle = object + + private_mod = types.ModuleType("ray._private") + private_mod.__path__ = [] + services_mod = types.ModuleType("ray._private.services") + services_mod.get_node_ip_address = lambda: "127.0.0.1" + private_mod.services = services_mod + + util_mod = types.ModuleType("ray.util") + util_mod.__path__ = [] + util_mod.get_node_ip_address = services_mod.get_node_ip_address + + scheduling_mod = types.ModuleType("ray.util.scheduling_strategies") + scheduling_mod.PlacementGroupSchedulingStrategy = _RaySchedulingStrategy + scheduling_mod.NodeAffinitySchedulingStrategy = _RaySchedulingStrategy + + placement_group_mod = types.ModuleType("ray.util.placement_group") + placement_group_mod.PlacementGroup = object + placement_group_mod.placement_group = lambda *args, **kwargs: types.SimpleNamespace(ready=lambda: None) # noqa: ARG005 + + ray_mod.actor = actor_mod + ray_mod._private = private_mod + ray_mod.util = util_mod + util_mod.scheduling_strategies = scheduling_mod + util_mod.placement_group = placement_group_mod + + sys.modules["ray"] = ray_mod + sys.modules["ray.actor"] = actor_mod + sys.modules["ray._private"] = private_mod + sys.modules["ray._private.services"] = services_mod + sys.modules["ray.util"] = util_mod + sys.modules["ray.util.scheduling_strategies"] = scheduling_mod + sys.modules["ray.util.placement_group"] = placement_group_mod + + def ensure_ray_stub() -> None: if real_module_available("ray"): return - ray = MagicMock() - sys.modules["ray"] = ray - sys.modules["ray._private"] = MagicMock() - sys.modules["ray._private.services"] = MagicMock() - sys.modules["ray.actor"] = MagicMock() + _install_ray_module_stub() + + +def install_pyarrow_parquet_stub() -> None: + """Avoid importing pyarrow's native parquet extension in CPU-only unit tests.""" + pyarrow_mod = types.ModuleType("pyarrow") + pyarrow_mod.__path__ = [] + parquet_mod = types.ModuleType("pyarrow.parquet") + + class ParquetFile: + def __init__(self, *args, **kwargs): # noqa: ARG002 + raise ImportError("pyarrow is required for parquet support") + + parquet_mod.ParquetFile = ParquetFile + pyarrow_mod.parquet = parquet_mod + sys.modules["pyarrow"] = pyarrow_mod + sys.modules["pyarrow.parquet"] = parquet_mod def install_rollout_optional_stubs() -> None: """Stub rollout-side optional imports when not installed.""" ensure_ray_stub() + install_pyarrow_parquet_stub() install_vllm_router_stub() @@ -216,14 +307,7 @@ def install_megatron_mpu_stub() -> MagicMock: def install_ray_stub() -> None: - ray_mod = types.ModuleType("ray") - ray_mod.get = lambda refs: refs - ray_mod.ObjectRef = object - ray_mod.actor = types.ModuleType("ray.actor") - ray_mod.actor.ActorHandle = object - ray_mod._private = types.SimpleNamespace(services=types.SimpleNamespace(get_node_ip_address=lambda: "127.0.0.1")) - sys.modules.setdefault("ray", ray_mod) - sys.modules.setdefault("ray.actor", ray_mod.actor) + _install_ray_module_stub() def install_vllm_cli_stubs() -> None: diff --git a/tests/test_sample.py b/tests/test_sample.py index bc9d3ff6..d8bf082c 100644 --- a/tests/test_sample.py +++ b/tests/test_sample.py @@ -50,6 +50,11 @@ def _make_sample(**overrides) -> Sample: loss_mask=[1, 1, 0, 1, 1], weight_versions=["v1"], rollout_log_probs=[-0.1, -0.2], + consistency_metadata={ + "schema_version": 1, + "sample": {"rollout_id": 7, "index": 42}, + "fingerprint": "sha256:unit-test", + }, rollout_top_p_token_ids=[10, 11, 12, 20], rollout_top_p_token_offsets=[0, 3, 4], rollout_routed_experts=[[0, 1], [2, 3]], @@ -133,6 +138,7 @@ def test_round_trip_preserves_every_field(): "loss_mask", "weight_versions", "rollout_log_probs", + "consistency_metadata", "rollout_top_p_token_ids", "rollout_top_p_token_offsets", "rollout_routed_experts", diff --git a/tests/utils/test_consistency_metadata.py b/tests/utils/test_consistency_metadata.py new file mode 100644 index 00000000..2bb5a9dc --- /dev/null +++ b/tests/utils/test_consistency_metadata.py @@ -0,0 +1,312 @@ +from __future__ import annotations + +import argparse +import sys +from pathlib import Path + +import pytest + +_tests_root = Path(__file__).resolve().parents[1] +if str(_tests_root) not in sys.path: + sys.path.insert(0, str(_tests_root)) + +import _unit_stubs + +_unit_stubs.install_rollout_optional_stubs() +_unit_stubs.install_vllm_cli_stubs() + +from vime.rollout.data_source import RolloutDataSourceWithBuffer # noqa: E402 +from vime.utils.consistency_metadata import ( + build_batch_layout_fingerprints, + build_requested_actual_provenance, + build_rollout_consistency_metadata, + ensure_sample_consistency_metadata, + get_consistency_mode, + raise_for_consistency_metadata_failures, + stable_fingerprint, + validate_samples_consistency_metadata, +) # noqa: E402 +from vime.utils.types import Sample # noqa: E402 + +NUM_GPUS = 0 + + +def _drop_rl_engine_modules() -> None: + for name in list(sys.modules): + if name == "rl_engine" or name.startswith("rl_engine."): + sys.modules.pop(name) + + +def _args(**overrides): + values = dict( + hf_checkpoint="unit/model", + padding_side="right", + num_gpus=2, + num_gpus_per_node=2, + rollout_num_gpus=2, + rollout_num_gpus_per_engine=1, + router_policy="consistent_hash", + vllm_router_ip="127.0.0.1", + vllm_router_port=8000, + rlk_consistency="audit", + buffer_filter_path=None, + rollout_global_dataset=False, + n_samples_per_prompt=2, + ) + values.update(overrides) + return argparse.Namespace(**values) + + +def _sample(**overrides) -> Sample: + values = dict( + index=5, + group_index=2, + rollout_id=11, + session_id="session-1", + tokens=[101, 201, 202, 203], + response_length=3, + loss_mask=[1, 0, 1], + rollout_log_probs=[-0.1, -0.2, -0.3], + weight_versions=["weights-v1"], + status=Sample.Status.COMPLETED, + metadata={}, + ) + values.update(overrides) + return Sample(**values) + + +def _complete_sample(**overrides) -> Sample: + sample = _sample(**overrides) + sample.consistency_metadata = build_rollout_consistency_metadata( + sample, + args=_args(), + sampling_params={"temperature": 0.7, "top_p": 0.95, "top_k": 50, "max_new_tokens": 128}, + logprob_contract_id="rlk.logp.fp32.v1", + requested_provenance={"backend": "native", "fallback": False}, + actual_provenance={"backend": "native", "fallback": False}, + ) + return sample + + +@pytest.mark.unit +def test_consistency_metadata_helpers_do_not_import_rl_engine(): + _drop_rl_engine_modules() + sample = _complete_sample() + + validation = validate_samples_consistency_metadata([sample], mode="strict") + + assert validation.ok + assert "alignment_profile" not in sample.consistency_metadata + assert "ScoreArtifact" not in sample.consistency_metadata + assert not any(name == "rl_engine" or name.startswith("rl_engine.") for name in sys.modules) + + +@pytest.mark.unit +def test_stable_fingerprint_is_order_insensitive_for_dicts(): + assert stable_fingerprint({"b": 2, "a": [1, 2]}) == stable_fingerprint({"a": [1, 2], "b": 2}) + assert stable_fingerprint({"a": [1, 2]}) != stable_fingerprint({"a": [2, 1]}) + + +@pytest.mark.unit +def test_rollout_metadata_records_compact_token_and_active_mask_fingerprints(): + sample = _complete_sample() + metadata = sample.consistency_metadata + + assert metadata["tokens"]["total_token_count"] == 4 + assert metadata["tokens"]["response_length"] == 3 + assert metadata["tokens"]["response_token_ids_fingerprint"] == stable_fingerprint([201, 202, 203]) + assert metadata["active_mask"]["active_token_count"] == 2 + assert metadata["active_mask"]["mask_fingerprint"] == stable_fingerprint([1, 0, 1]) + assert metadata["old_logp"]["source"] == "rollout_engine" + assert metadata["old_logp"]["contract_id"] == "rlk.logp.fp32.v1" + + +@pytest.mark.unit +def test_response_token_fingerprint_excludes_prompt_tokens(): + first = _complete_sample(tokens=[1, 10, 11], response_length=2, loss_mask=[1, 1]) + second = _complete_sample(tokens=[999, 10, 11], response_length=2, loss_mask=[1, 1]) + + assert ( + first.consistency_metadata["tokens"]["response_token_ids_fingerprint"] + == second.consistency_metadata["tokens"]["response_token_ids_fingerprint"] + ) + assert ( + first.consistency_metadata["tokens"]["token_ids_fingerprint"] + != second.consistency_metadata["tokens"]["token_ids_fingerprint"] + ) + + +@pytest.mark.unit +def test_metadata_uses_sample_index_as_default_rollout_identifier(): + sample = _complete_sample(rollout_id=None, index=42) + + assert sample.consistency_metadata["sample"]["rollout_id"] == 42 + validation = validate_samples_consistency_metadata([sample], mode="strict") + assert validation.ok + + +@pytest.mark.unit +def test_loss_mask_length_mismatch_is_rejected_before_audit_claims(): + sample = _sample(response_length=3, loss_mask=[1, 0]) + with pytest.raises(ValueError, match="loss_mask length"): + build_rollout_consistency_metadata( + sample, + args=_args(), + sampling_params={"temperature": 1.0}, + logprob_contract_id="contract", + ) + + +@pytest.mark.unit +def test_audit_mode_reports_missing_custom_metadata_as_structured_warning(): + validation = validate_samples_consistency_metadata([_sample(consistency_metadata=None)], mode="audit") + + assert validation.ok + assert [issue.code for issue in validation.warnings] == ["consistency_metadata_missing"] + assert validation.warnings[0].sample_index == 5 + assert validation.warnings[0].rollout_id == 11 + assert validation.failures == [] + + +@pytest.mark.unit +def test_strict_mode_fails_closed_when_required_metadata_is_missing(): + validation = validate_samples_consistency_metadata([_sample(consistency_metadata=None)], mode="strict") + + assert not validation.ok + assert [issue.code for issue in validation.failures] == ["consistency_metadata_missing"] + with pytest.raises(ValueError, match="Strict consistency metadata validation failed"): + raise_for_consistency_metadata_failures(validation) + + +@pytest.mark.unit +def test_complete_metadata_passes_strict_validation_and_counts_active_tokens(): + validation = validate_samples_consistency_metadata([_complete_sample()], mode="strict") + + assert validation.ok + assert validation.active_token_count == 2 + assert validation.zero_active_token_samples == [] + assert validation.warnings == [] + assert validation.failures == [] + + +@pytest.mark.unit +def test_zero_active_token_sample_is_identified_without_becoming_strict_failure(): + validation = validate_samples_consistency_metadata( + [_complete_sample(tokens=[1, 2], response_length=2, loss_mask=[0, 0])], + mode="strict", + ) + + assert validation.ok + assert validation.active_token_count == 0 + assert validation.zero_active_token_samples == [{"sample_index": 5, "rollout_id": 11}] + assert [issue.code for issue in validation.warnings] == ["zero_active_tokens"] + + +@pytest.mark.unit +def test_requested_actual_provenance_mismatch_is_audit_warning_and_strict_failure(): + sample = _complete_sample() + sample.consistency_metadata["provenance"] = build_requested_actual_provenance( + requested={"backend": "native", "fallback": False}, + actual={"backend": "fallback-native", "fallback": True}, + ) + + audit = validate_samples_consistency_metadata([sample], mode="audit") + assert {issue.code for issue in audit.warnings} == { + "requested_actual_provenance_mismatch", + "undeclared_runtime_fallback", + } + + strict = validate_samples_consistency_metadata([sample], mode="strict") + assert {issue.code for issue in strict.failures} == { + "requested_actual_provenance_mismatch", + "undeclared_runtime_fallback", + } + + +@pytest.mark.unit +def test_ensure_sample_consistency_metadata_does_not_overwrite_existing_record_by_default(): + sample = _sample(consistency_metadata={"schema_version": 1, "fingerprint": "existing"}) + + record = ensure_sample_consistency_metadata( + sample, + args=_args(), + sampling_params={"temperature": 1.0}, + logprob_contract_id="contract", + ) + + assert record == {"schema_version": 1, "fingerprint": "existing"} + assert sample.consistency_metadata == {"schema_version": 1, "fingerprint": "existing"} + + +@pytest.mark.unit +def test_batch_layout_fingerprints_map_samples_to_rank_and_microbatch(): + train_data = { + "tokens": [[1, 2], [3, 4, 5], [6, 7], [8, 9]], + "total_lengths": [2, 3, 2, 2], + "response_lengths": [2, 2, 2, 2], + "loss_masks": [[1, 1], [1, 0], [0, 1], [0, 0]], + "rollout_ids": [10, 11, 12, 13], + "sample_indices": [0, 1, 2, 3], + } + + layouts = build_batch_layout_fingerprints( + train_data, + partitions=[[0, 2], [1, 3]], + micro_batch_indices=[[[0], [1]], [[0, 1]]], + num_microbatches=[2], + global_batch_sizes=[2], + ) + + assert layouts[2]["dp_rank"] == 0 + assert layouts[2]["microbatch_id"] == 1 + assert layouts[2]["active_token_count"] == 1 + assert layouts[3]["dp_rank"] == 1 + assert layouts[3]["microbatch_size"] == 2 + assert layouts[3]["active_mask_density"] == 0.0 + assert layouts[0]["global_batch_shape"]["packed_order_fingerprint"] == stable_fingerprint([0, 2, 1, 3]) + + +@pytest.mark.unit +def test_batch_layout_fingerprints_mark_samples_dropped_by_existing_schedule_trim(): + train_data = { + "tokens": [[1, 2], [3, 4]], + "total_lengths": [2, 2], + "response_lengths": [2, 2], + "loss_masks": [[1, 1], [1, 0]], + "rollout_ids": [10, 11], + "sample_indices": [0, 1], + } + + layouts = build_batch_layout_fingerprints( + train_data, + partitions=[[0]], + micro_batch_indices=[[[0]]], + num_microbatches=[1], + global_batch_sizes=[1], + ) + + assert "dropped_by_schedule" not in layouts[0] + assert layouts[1]["dropped_by_schedule"] is True + assert layouts[1]["active_token_count"] == 1 + + +@pytest.mark.unit +def test_rollout_data_source_buffer_preserves_consistency_metadata(): + source = RolloutDataSourceWithBuffer(_args()) + sample = _complete_sample() + + source.add_samples([[sample, _complete_sample(index=6, rollout_id=12)]]) + returned = source.get_samples(1) + + assert returned[0][0].consistency_metadata == sample.consistency_metadata + + +@pytest.mark.unit +def test_get_consistency_mode_accepts_args_and_env_alias(monkeypatch): + assert get_consistency_mode(_args(rlk_consistency="strict")) == "strict" + + monkeypatch.setenv("VIME_RLK_CONSISTENCY", "audit") + assert get_consistency_mode(argparse.Namespace()) == "audit" + + with pytest.raises(ValueError, match="Unsupported RL-Kernel consistency mode"): + get_consistency_mode(_args(rlk_consistency="maybe")) diff --git a/tests/utils/test_consistency_metadata_rollout_integration.py b/tests/utils/test_consistency_metadata_rollout_integration.py new file mode 100644 index 00000000..146611d8 --- /dev/null +++ b/tests/utils/test_consistency_metadata_rollout_integration.py @@ -0,0 +1,106 @@ +from __future__ import annotations + +import sys +from pathlib import Path + +import pytest + +_tests_root = Path(__file__).resolve().parents[1] +if str(_tests_root) not in sys.path: + sys.path.insert(0, str(_tests_root)) + +import _unit_stubs + +_unit_stubs.install_rollout_optional_stubs() +_unit_stubs.install_vllm_cli_stubs() + +from vime.ray.rollout import RolloutManager # noqa: E402 +from vime.utils.consistency_metadata import build_rollout_consistency_metadata # noqa: E402 +from vime.utils.types import Sample # noqa: E402 + + +NUM_GPUS = 0 + + +class Args: + reward_key = None + advantage_estimator = "grpo" + rewards_normalization = False + grpo_std_normalization = False + rollout_top_p = 1.0 + rlk_consistency = "audit" + hf_checkpoint = "unit/model" + padding_side = "right" + num_gpus = 1 + rollout_num_gpus = 1 + rollout_num_gpus_per_engine = 1 + router_policy = "round_robin" + + +def _manager(mode: str): + cls = RolloutManager.__ray_actor_class__ + manager = cls.__new__(cls) + manager.args = Args() + manager.args.rlk_consistency = mode + manager.custom_reward_post_process_func = None + manager.custom_convert_samples_to_train_data_func = None + return manager + + +def _sample(index: int = 0) -> Sample: + return Sample( + index=index, + rollout_id=index, + tokens=[101, 201, 202], + response_length=2, + loss_mask=[1, 1], + rollout_log_probs=[-0.1, -0.2], + weight_versions=["w1"], + reward=1.0, + status=Sample.Status.COMPLETED, + metadata={}, + ) + + +@pytest.mark.unit +def test_convert_samples_to_train_data_leaves_consistency_fields_out_when_off(): + train_data = _manager("off")._convert_samples_to_train_data([_sample()]) + + assert "consistency_metadata" not in train_data + assert "consistency_metadata_validation" not in train_data + + +@pytest.mark.unit +def test_convert_samples_to_train_data_reports_audit_warning_for_missing_metadata(): + train_data = _manager("audit")._convert_samples_to_train_data([_sample()]) + + assert train_data["consistency_metadata"] == [None] + assert train_data["consistency_metadata_validation"]["ok"] is True + assert [issue["code"] for issue in train_data["consistency_metadata_validation"]["warnings"]] == [ + "consistency_metadata_missing" + ] + + +@pytest.mark.unit +def test_convert_samples_to_train_data_fails_closed_in_strict_mode(): + with pytest.raises(ValueError, match="Strict consistency metadata validation failed"): + _manager("strict")._convert_samples_to_train_data([_sample()]) + + +@pytest.mark.unit +def test_convert_samples_to_train_data_carries_sample_metadata_when_present(): + sample = _sample() + sample.consistency_metadata = build_rollout_consistency_metadata( + sample, + args=Args(), + sampling_params={"temperature": 1.0, "top_p": 1.0}, + logprob_contract_id="contract-v1", + requested_provenance={"backend": "native", "fallback": False}, + actual_provenance={"backend": "native", "fallback": False}, + ) + + train_data = _manager("strict")._convert_samples_to_train_data([sample]) + + assert train_data["consistency_metadata"] == [sample.consistency_metadata] + assert train_data["consistency_metadata_validation"]["ok"] is True + assert train_data["consistency_metadata_validation"]["active_token_count"] == 2 diff --git a/vime/ray/rollout.py b/vime/ray/rollout.py index 0607258d..0d1bd1e3 100644 --- a/vime/ray/rollout.py +++ b/vime/ray/rollout.py @@ -22,6 +22,13 @@ GPU_MEMORY_TYPE_CUDA_GRAPH = "cuda_graph" from vime.rollout.base_types import call_rollout_fn from vime.utils import logging_utils +from vime.utils.consistency_metadata import ( + build_batch_layout_fingerprints, + get_consistency_mode, + raise_for_consistency_metadata_failures, + sample_consistency_metadata, + validate_samples_consistency_metadata, +) from vime.utils.data import get_source from vime.utils.dp_schedule import build_dp_schedule from vime.utils.health_monitor import RolloutHealthMonitor @@ -752,6 +759,13 @@ def _convert_samples_to_train_data(self, samples: list[Sample] | list[list[Sampl loss_masks.append(sample.loss_mask) train_data["loss_masks"] = loss_masks + consistency_mode = get_consistency_mode(self.args) + if consistency_mode != "off": + validation = validate_samples_consistency_metadata(samples, mode=consistency_mode) + train_data["consistency_metadata_validation"] = validation.to_dict() + raise_for_consistency_metadata_failures(validation) + train_data["consistency_metadata"] = [sample_consistency_metadata(sample) for sample in samples] + # Per-rollout aggregate, precomputed at the step level (where we can # see every sample of every rollout) and broadcast per-sample so the # per-mb loss reducer uses the correct whole-rollout denominator even @@ -845,6 +859,15 @@ def _split_train_data_by_dp(self, data): rollout_indices=data["rollout_ids"], ) + if get_consistency_mode(self.args) != "off" or "consistency_metadata" in data: + data["consistency_batch_layout_fingerprints"] = build_batch_layout_fingerprints( + data, + partitions=partitions, + micro_batch_indices=micro_batch_indices, + num_microbatches=num_microbatches, + global_batch_sizes=global_batch_sizes, + ) + # Package per-rank rollout_data rollout_data_refs = [] for r in range(dp_size): @@ -868,6 +891,8 @@ def _split_train_data_by_dp(self, data): "source_names", "prompt", "teacher_log_probs", + "consistency_metadata", + "consistency_batch_layout_fingerprints", ]: if key not in data: continue @@ -877,6 +902,8 @@ def _split_train_data_by_dp(self, data): if key not in data: continue rollout_data[key] = data[key] + if "consistency_metadata_validation" in data: + rollout_data["consistency_metadata_validation"] = data["consistency_metadata_validation"] rollout_data["global_batch_sizes"] = global_batch_sizes rollout_data["num_microbatches"] = num_microbatches rollout_data["micro_batch_indices"] = micro_batch_indices[r] diff --git a/vime/rollout/vllm_rollout.py b/vime/rollout/vllm_rollout.py index 7e657e71..ed5cc593 100644 --- a/vime/rollout/vllm_rollout.py +++ b/vime/rollout/vllm_rollout.py @@ -19,6 +19,11 @@ from vime.rollout.base_types import RolloutFnEvalOutput, RolloutFnTrainOutput from vime.rollout.filter_hub.base_types import MetricGatherer, call_dynamic_filter from vime.utils.async_utils import run +from vime.utils.consistency_metadata import ( + ensure_sample_consistency_metadata, + get_consistency_mode, + stable_fingerprint, +) from vime.utils.data import Dataset from vime.utils.eval_config import EvalDatasetConfig from vime.utils.http_utils import get, get_rollout_num_engines, post @@ -44,6 +49,52 @@ _ABORT_RESWEEP_INTERVAL_S = 3.0 +def _iter_samples(node: Sample | list[Any]): + if isinstance(node, Sample): + yield node + return + if isinstance(node, list): + for item in node: + yield from _iter_samples(item) + + +def _refresh_consistency_fingerprint(record: dict[str, Any]) -> None: + payload = dict(record) + payload.pop("fingerprint", None) + record["fingerprint"] = stable_fingerprint(payload) + + +def _annotate_consistency_metadata( + args: Namespace, + node: Sample | list[Any], + *, + sampling_params: dict[str, Any], + source: str, +) -> None: + if get_consistency_mode(args) == "off": + return + for sample in _iter_samples(node): + old_logp_source = source if sample.rollout_log_probs is not None else None + ensure_sample_consistency_metadata( + sample, + args=args, + sampling_params=sampling_params, + model_name=getattr(args, "hf_checkpoint", None), + old_logp_source=old_logp_source, + ) + + +def _record_dynamic_sampling_decision(node: Sample | list[Any], *, keep: bool, reason: str | None) -> None: + decision = {"keep": bool(keep), "reason": reason} + for sample in _iter_samples(node): + if sample.metadata is None: + sample.metadata = {} + sample.metadata["dynamic_sampling"] = decision + if sample.consistency_metadata is not None: + sample.consistency_metadata["dynamic_sampling"] = decision + _refresh_consistency_fingerprint(sample.consistency_metadata) + + def _coerce_flat_int_token_ids(ids: Any) -> list[int]: """Flatten tokenizer/processor output into ``list[int]`` for vLLM ``/inference/v1/generate``.""" if ids is None: @@ -438,6 +489,7 @@ async def generate_and_rm( with state.dp_rank_context() as _: # Check sample.generate_function_path for per-sample custom_generate_function_path (e.g., from eval dataset config) custom_func_path = getattr(sample, "generate_function_path", None) or args.custom_generate_function_path + consistency_source = "custom_generate" if custom_func_path is not None else "vllm_rollout" if custom_func_path is not None: custom_generate_func = load_function(custom_func_path) @@ -449,6 +501,13 @@ async def generate_and_rm( else: sample = await generate(args, sample, sampling_params) + _annotate_consistency_metadata( + args, + sample, + sampling_params=sampling_params, + source=consistency_source, + ) + # for the rm that need the whole group, we will not do the rm here if args.group_rm: return sample @@ -627,6 +686,12 @@ async def generate_rollout_async( all_data.append(group) dynamic_filter_output = call_dynamic_filter(dynamic_filter, args, group) + if get_consistency_mode(args) != "off": + _record_dynamic_sampling_decision( + group, + keep=dynamic_filter_output.keep, + reason=dynamic_filter_output.reason, + ) if not dynamic_filter_output.keep: metric_gatherer.on_dynamic_filter_drop(reason=dynamic_filter_output.reason) state.remaining_batch_size -= 1 diff --git a/vime/utils/consistency_metadata.py b/vime/utils/consistency_metadata.py new file mode 100644 index 00000000..3573e9bd --- /dev/null +++ b/vime/utils/consistency_metadata.py @@ -0,0 +1,612 @@ +"""Compact consistency metadata for rollout-to-training audit paths. + +The helpers in this module keep the rollout boundary observable without +shipping full token or mask payloads inside metadata records. Full tensors +continue to live in the existing training batch fields; metadata carries +stable fingerprints, active-token counts, and requested-vs-actual provenance +needed by audit/strict consistency checks. + +This is intentionally vime runtime metadata. RL-Kernel-owned alignment schemas, +A0-A5 profiles, comparators, and tolerance contracts stay behind the +``vime.backends.rl_kernel_utils`` adapter boundary. +""" + +from __future__ import annotations + +import dataclasses +import hashlib +import json +import os +from dataclasses import dataclass, field +from enum import Enum +from typing import Any + +CONSISTENCY_METADATA_SCHEMA_VERSION = 1 +CONSISTENCY_MODES = {"off", "audit", "strict"} + +REQUIRED_COMPARISON_FIELDS = ( + ("sample.rollout_id", "rollout_id_missing"), + ("tokens.response_token_ids_fingerprint", "token_ids_missing"), + ("active_mask.mask_fingerprint", "active_mask_missing"), + ("active_mask.active_token_count", "active_mask_missing"), + ("tokenizer.fingerprint", "tokenizer_missing"), + ("sampling.params_fingerprint", "sampling_config_missing"), + ("padding.side", "padding_semantics_missing"), + ("model.name", "model_name_missing"), + ("weight.version", "weight_version_missing"), + ("old_logp.source", "old_logp_source_missing"), + ("old_logp.contract_id", "logprob_contract_id_missing"), + ("provenance.actual_fingerprint", "actual_provenance_missing"), +) + +_SAMPLING_PARAM_KEYS = ( + "temperature", + "top_p", + "top_k", + "max_new_tokens", + "max_tokens", + "seed", + "stop", + "stop_token_ids", + "skip_special_tokens", +) + +_PARALLEL_ARG_KEYS = ( + "num_gpus", + "num_gpus_per_node", + "rollout_num_gpus", + "rollout_num_gpus_per_engine", + "megatron_tensor_parallel_size", + "megatron_context_parallel_size", + "megatron_expert_model_parallel_size", + "tensor_model_parallel_size", + "context_parallel_size", +) + +_ROUTER_ARG_KEYS = ( + "router_policy", + "vllm_router_ip", + "vllm_router_port", + "vllm_dp_size", + "vllm_enable_prefix_caching", + "vllm_enable_deterministic_inference", + "rollout_data_transport", +) + + +@dataclass(frozen=True) +class ConsistencyMetadataIssue: + code: str + message: str + severity: str + sample_index: int | None = None + rollout_id: int | None = None + field: str | None = None + + def to_dict(self) -> dict[str, Any]: + return { + "code": self.code, + "message": self.message, + "severity": self.severity, + "sample_index": self.sample_index, + "rollout_id": self.rollout_id, + "field": self.field, + } + + +@dataclass(frozen=True) +class ConsistencyMetadataValidation: + mode: str + active_token_count: int = 0 + zero_active_token_samples: list[dict[str, int | None]] = field(default_factory=list) + warnings: list[ConsistencyMetadataIssue] = field(default_factory=list) + failures: list[ConsistencyMetadataIssue] = field(default_factory=list) + + @property + def ok(self) -> bool: + return not self.failures + + def to_dict(self) -> dict[str, Any]: + return { + "mode": self.mode, + "ok": self.ok, + "active_token_count": self.active_token_count, + "zero_active_token_samples": self.zero_active_token_samples, + "warnings": [issue.to_dict() for issue in self.warnings], + "failures": [issue.to_dict() for issue in self.failures], + } + + +def _normalize_for_json(value: Any) -> Any: + if dataclasses.is_dataclass(value): + return _normalize_for_json(dataclasses.asdict(value)) + if isinstance(value, Enum): + return value.value + if isinstance(value, dict): + return {str(key): _normalize_for_json(value[key]) for key in sorted(value, key=str)} + if isinstance(value, (list, tuple)): + return [_normalize_for_json(item) for item in value] + if hasattr(value, "detach") and callable(value.detach): + value = value.detach().cpu() + if hasattr(value, "tolist") and callable(value.tolist) and not isinstance(value, (str, bytes)): + return _normalize_for_json(value.tolist()) + if isinstance(value, bytes): + return value.hex() + if value is None or isinstance(value, (bool, int, float, str)): + return value + return repr(value) + + +def stable_fingerprint(value: Any) -> str: + payload = json.dumps( + _normalize_for_json(value), + sort_keys=True, + separators=(",", ":"), + ensure_ascii=True, + ) + return "sha256:" + hashlib.sha256(payload.encode("utf-8")).hexdigest() + + +def get_consistency_mode(args: Any | None = None) -> str: + value = None + if args is not None: + value = ( + getattr(args, "rlk_consistency", None) + or getattr(args, "rl_kernel_consistency", None) + or getattr(args, "rl_kernel_consistency_mode", None) + ) + if value is None: + value = os.environ.get("VIME_RLK_CONSISTENCY") or os.environ.get("VIME_RL_KERNEL_CONSISTENCY") + mode = str(value or "off").lower() + if mode not in CONSISTENCY_MODES: + raise ValueError( + f"Unsupported RL-Kernel consistency mode {value!r}; expected one of {sorted(CONSISTENCY_MODES)}" + ) + return mode + + +def _to_list(value: Any | None) -> list[Any]: + if value is None: + return [] + if hasattr(value, "detach") and callable(value.detach): + value = value.detach().cpu() + if hasattr(value, "tolist") and callable(value.tolist) and not isinstance(value, (str, bytes)): + value = value.tolist() + return list(value) + + +def _to_int_mask(mask: Any | None, response_length: int) -> list[int]: + if mask is None: + return [1] * response_length + values = [int(v) for v in _to_list(mask)] + if len(values) != response_length: + raise ValueError(f"loss_mask length {len(values)} != response_length {response_length}") + return values + + +def count_active_tokens(loss_mask: Any | None, response_length: int) -> int: + return sum(1 for value in _to_int_mask(loss_mask, response_length) if value) + + +def _metadata_dict(sample: Any) -> dict[str, Any]: + metadata = getattr(sample, "metadata", None) + return metadata if isinstance(metadata, dict) else {} + + +def _sample_status_value(sample: Any) -> str | None: + status = getattr(sample, "status", None) + return getattr(status, "value", status) + + +def _compact_attrs(obj: Any | None, keys: tuple[str, ...]) -> dict[str, Any]: + if obj is None: + return {} + return {key: getattr(obj, key) for key in keys if hasattr(obj, key) and getattr(obj, key) is not None} + + +def _compact_sampling_params(sampling_params: dict[str, Any] | None) -> dict[str, Any]: + if not sampling_params: + return {} + return {key: sampling_params[key] for key in _SAMPLING_PARAM_KEYS if key in sampling_params} + + +def build_requested_actual_provenance( + *, + requested: dict[str, Any] | None = None, + actual: dict[str, Any] | None = None, +) -> dict[str, Any]: + requested = _normalize_for_json(requested or {}) + actual = _normalize_for_json(actual or {}) + mismatches = {} + for key in sorted(set(requested) | set(actual)): + if key in requested and key in actual and requested[key] != actual[key]: + mismatches[key] = {"requested": requested[key], "actual": actual[key]} + + fallback = actual.get("fallback") + requested_fallback = requested.get("fallback") + undeclared_fallback = bool(fallback) and requested_fallback is not True + + return { + "requested": requested, + "actual": actual, + "requested_fingerprint": stable_fingerprint(requested) if requested else None, + "actual_fingerprint": stable_fingerprint(actual) if actual else None, + "mismatches": mismatches, + "undeclared_fallback": undeclared_fallback, + } + + +def _extract_requested_provenance(args: Any | None) -> dict[str, Any]: + requested = _compact_attrs(args, _ROUTER_ARG_KEYS + _PARALLEL_ARG_KEYS) + requested.update(_compact_attrs(args, ("hf_checkpoint", "load", "save"))) + return requested + + +def _extract_actual_provenance(args: Any | None, *, model_name: str | None) -> dict[str, Any]: + actual = _compact_attrs(args, _ROUTER_ARG_KEYS + _PARALLEL_ARG_KEYS) + if model_name is not None: + actual["model_name"] = model_name + return actual + + +def build_rollout_consistency_metadata( + sample: Any, + *, + args: Any | None = None, + sampling_params: dict[str, Any] | None = None, + model_name: str | None = None, + old_logp_source: str | None = None, + logprob_contract_id: str | None = None, + tokenizer_fingerprint: str | None = None, + padding_side: str | None = None, + requested_provenance: dict[str, Any] | None = None, + actual_provenance: dict[str, Any] | None = None, + batch_layout: dict[str, Any] | None = None, + dynamic_sampling: dict[str, Any] | None = None, +) -> dict[str, Any]: + metadata = _metadata_dict(sample) + response_length = int(getattr(sample, "response_length", 0) or 0) + tokens = [int(token) for token in _to_list(getattr(sample, "tokens", []))] + response_tokens = tokens[-response_length:] if response_length else [] + active_mask = _to_int_mask(getattr(sample, "loss_mask", None), response_length) + active_token_count = sum(1 for value in active_mask if value) + rollout_log_probs = getattr(sample, "rollout_log_probs", None) + weight_versions = list(getattr(sample, "weight_versions", None) or []) + rollout_id = getattr(sample, "rollout_id", None) + if rollout_id is None: + rollout_id = getattr(sample, "index", None) + + sampling_summary = _compact_sampling_params(sampling_params) or metadata.get("sampling_params") or {} + model_name = model_name or metadata.get("model_name") or getattr(args, "hf_checkpoint", None) + old_logp_source = old_logp_source or metadata.get("old_logp_source") + if old_logp_source is None and rollout_log_probs is not None: + old_logp_source = "rollout_engine" + logprob_contract_id = logprob_contract_id or metadata.get("logprob_contract_id") + tokenizer_fingerprint = tokenizer_fingerprint or metadata.get("tokenizer_fingerprint") + if tokenizer_fingerprint is None and getattr(args, "hf_checkpoint", None) is not None: + tokenizer_fingerprint = stable_fingerprint({"hf_checkpoint": args.hf_checkpoint}) + padding_side = padding_side or metadata.get("padding_side") or getattr(args, "padding_side", None) + + position_cache = metadata.get("position_cache") or metadata.get("position_cache_metadata") + quantization = metadata.get("quantization") or _compact_attrs(args, ("quantization", "vllm_quantization")) + parallel_placement = metadata.get("parallel_placement") or _compact_attrs(args, _PARALLEL_ARG_KEYS) + dynamic_sampling = dynamic_sampling or metadata.get("dynamic_sampling") + + requested = requested_provenance if requested_provenance is not None else _extract_requested_provenance(args) + actual = ( + actual_provenance if actual_provenance is not None else _extract_actual_provenance(args, model_name=model_name) + ) + + record = { + "schema_version": CONSISTENCY_METADATA_SCHEMA_VERSION, + "sample": { + "group_index": getattr(sample, "group_index", None), + "index": getattr(sample, "index", None), + "rollout_id": rollout_id, + "session_id": getattr(sample, "session_id", None), + "status": _sample_status_value(sample), + }, + "tokens": { + "total_token_count": len(tokens), + "response_length": response_length, + "token_ids_fingerprint": stable_fingerprint(tokens), + "response_token_ids_fingerprint": stable_fingerprint(response_tokens), + }, + "active_mask": { + "response_length": response_length, + "active_token_count": active_token_count, + "mask_fingerprint": stable_fingerprint(active_mask), + "zero_active_tokens": active_token_count == 0, + }, + "tokenizer": {"fingerprint": tokenizer_fingerprint}, + "sampling": { + "params_fingerprint": stable_fingerprint(sampling_summary) if sampling_summary else None, + "summary": _normalize_for_json(sampling_summary), + }, + "padding": {"side": padding_side}, + "position_cache": { + "fingerprint": stable_fingerprint(position_cache) if position_cache else None, + }, + "quantization": { + "fingerprint": stable_fingerprint(quantization) if quantization else None, + "summary": _normalize_for_json(quantization), + }, + "parallel_placement": { + "fingerprint": stable_fingerprint(parallel_placement) if parallel_placement else None, + "summary": _normalize_for_json(parallel_placement), + }, + "model": {"name": model_name}, + "old_logp": { + "source": old_logp_source, + "contract_id": logprob_contract_id, + "num_values": len(rollout_log_probs) if rollout_log_probs is not None else None, + "fingerprint": stable_fingerprint(rollout_log_probs) if rollout_log_probs is not None else None, + }, + "weight": { + "version": metadata.get("weight_version") or (weight_versions[-1] if weight_versions else None), + "versions_fingerprint": stable_fingerprint(weight_versions) if weight_versions else None, + "pre_update": metadata.get("pre_update"), + }, + "provenance": build_requested_actual_provenance(requested=requested, actual=actual), + "batch_layout": batch_layout, + "dynamic_sampling": _normalize_for_json(dynamic_sampling), + } + record["fingerprint"] = stable_fingerprint(record) + return record + + +def ensure_sample_consistency_metadata( + sample: Any, + *, + args: Any | None = None, + sampling_params: dict[str, Any] | None = None, + overwrite: bool = False, + **metadata_kwargs: Any, +) -> dict[str, Any]: + existing = getattr(sample, "consistency_metadata", None) + if existing is not None and not overwrite: + return existing + record = build_rollout_consistency_metadata( + sample, + args=args, + sampling_params=sampling_params, + **metadata_kwargs, + ) + sample.consistency_metadata = record + return record + + +def sample_consistency_metadata(sample: Any) -> dict[str, Any] | None: + direct = getattr(sample, "consistency_metadata", None) + if direct is not None: + return direct + metadata = _metadata_dict(sample) + nested = metadata.get("consistency_metadata") + return nested if isinstance(nested, dict) else None + + +def _get_path(mapping: dict[str, Any], path: str) -> Any: + current: Any = mapping + for part in path.split("."): + if not isinstance(current, dict) or part not in current: + return None + current = current[part] + return current + + +def _sample_ids(sample: Any, metadata: dict[str, Any] | None) -> tuple[int | None, int | None]: + sample_index = getattr(sample, "index", None) + rollout_id = getattr(sample, "rollout_id", None) + if metadata: + sample_info = metadata.get("sample") or {} + sample_index = sample_info.get("index", sample_index) + rollout_id = sample_info.get("rollout_id", rollout_id) + if rollout_id is None: + rollout_id = sample_index + return sample_index, rollout_id + + +def _issue( + *, + mode: str, + code: str, + message: str, + sample_index: int | None, + rollout_id: int | None, + field: str | None = None, + strict_failure: bool = True, +) -> tuple[ConsistencyMetadataIssue, bool]: + is_failure = mode == "strict" and strict_failure + severity = "error" if is_failure else "warning" + return ( + ConsistencyMetadataIssue( + code=code, + message=message, + severity=severity, + sample_index=sample_index, + rollout_id=rollout_id, + field=field, + ), + is_failure, + ) + + +def validate_samples_consistency_metadata( + samples: list[Any], + *, + mode: str, + required_fields: tuple[tuple[str, str], ...] = REQUIRED_COMPARISON_FIELDS, +) -> ConsistencyMetadataValidation: + mode = str(mode).lower() + if mode not in CONSISTENCY_MODES: + raise ValueError(f"Unsupported consistency mode {mode!r}") + if mode == "off": + return ConsistencyMetadataValidation(mode=mode) + + warnings: list[ConsistencyMetadataIssue] = [] + failures: list[ConsistencyMetadataIssue] = [] + zero_active_token_samples: list[dict[str, int | None]] = [] + active_token_count = 0 + + def add_issue(issue: ConsistencyMetadataIssue, is_failure: bool) -> None: + if is_failure: + failures.append(issue) + else: + warnings.append(issue) + + for sample in samples: + metadata = sample_consistency_metadata(sample) + sample_index, rollout_id = _sample_ids(sample, metadata) + if metadata is None: + issue, is_failure = _issue( + mode=mode, + code="consistency_metadata_missing", + message="Sample is missing consistency metadata required for audit/strict comparison.", + sample_index=sample_index, + rollout_id=rollout_id, + ) + add_issue(issue, is_failure) + continue + + for path, code in required_fields: + if _get_path(metadata, path) in (None, ""): + issue, is_failure = _issue( + mode=mode, + code=code, + message=f"Consistency metadata field {path!r} is missing.", + sample_index=sample_index, + rollout_id=rollout_id, + field=path, + ) + add_issue(issue, is_failure) + + active = _get_path(metadata, "active_mask.active_token_count") + if active is not None: + active_token_count += int(active) + if int(active) == 0: + zero_active_token_samples.append({"sample_index": sample_index, "rollout_id": rollout_id}) + issue, is_failure = _issue( + mode=mode, + code="zero_active_tokens", + message="Sample has zero active response/action tokens for consistency aggregates.", + sample_index=sample_index, + rollout_id=rollout_id, + field="active_mask.active_token_count", + strict_failure=False, + ) + add_issue(issue, is_failure) + + provenance = metadata.get("provenance") if isinstance(metadata, dict) else None + if isinstance(provenance, dict): + mismatches = provenance.get("mismatches") or {} + if mismatches: + issue, is_failure = _issue( + mode=mode, + code="requested_actual_provenance_mismatch", + message="Requested-vs-actual provenance differs for compared sample.", + sample_index=sample_index, + rollout_id=rollout_id, + field="provenance.mismatches", + ) + add_issue(issue, is_failure) + if provenance.get("undeclared_fallback"): + issue, is_failure = _issue( + mode=mode, + code="undeclared_runtime_fallback", + message="Actual provenance reports fallback that was not declared in requested provenance.", + sample_index=sample_index, + rollout_id=rollout_id, + field="provenance.undeclared_fallback", + ) + add_issue(issue, is_failure) + + return ConsistencyMetadataValidation( + mode=mode, + active_token_count=active_token_count, + zero_active_token_samples=zero_active_token_samples, + warnings=warnings, + failures=failures, + ) + + +def raise_for_consistency_metadata_failures(validation: ConsistencyMetadataValidation) -> None: + if not validation.failures: + return + preview = "; ".join( + f"{issue.code}(sample_index={issue.sample_index}, rollout_id={issue.rollout_id}, field={issue.field})" + for issue in validation.failures[:5] + ) + remaining = len(validation.failures) - 5 + if remaining > 0: + preview += f"; ... {remaining} more" + raise ValueError(f"Strict consistency metadata validation failed: {preview}") + + +def build_batch_layout_fingerprints( + train_data: dict[str, Any], + *, + partitions: list[list[int]], + micro_batch_indices: list[list[list[int]]], + num_microbatches: list[int], + global_batch_sizes: list[int], +) -> list[dict[str, Any]]: + sample_count = len(train_data["tokens"]) + total_lengths = train_data["total_lengths"] + response_lengths = train_data["response_lengths"] + loss_masks = train_data["loss_masks"] + rollout_ids = train_data.get("rollout_ids", [None] * sample_count) + sample_indices = train_data.get("sample_indices", [None] * sample_count) + + packed_order = [sample_index for partition in partitions for sample_index in partition] + shared = { + "dp_size": len(partitions), + "num_microbatches": num_microbatches, + "global_batch_sizes": global_batch_sizes, + "sequence_lengths_fingerprint": stable_fingerprint(total_lengths), + "packed_order_fingerprint": stable_fingerprint(packed_order), + } + + layouts: list[dict[str, Any] | None] = [None] * sample_count + for dp_rank, partition in enumerate(partitions): + for microbatch_id, local_indices in enumerate(micro_batch_indices[dp_rank]): + for microbatch_offset, local_index in enumerate(local_indices): + global_index = partition[local_index] + response_length = int(response_lengths[global_index]) + active_tokens = count_active_tokens(loss_masks[global_index], response_length) + layout = { + "schema_version": CONSISTENCY_METADATA_SCHEMA_VERSION, + "sample_index": sample_indices[global_index], + "rollout_id": rollout_ids[global_index], + "total_length": int(total_lengths[global_index]), + "response_length": response_length, + "active_token_count": active_tokens, + "active_mask_density": active_tokens / response_length if response_length else 0.0, + "dp_rank": dp_rank, + "rank_local_index": local_index, + "microbatch_id": microbatch_id, + "microbatch_offset": microbatch_offset, + "microbatch_size": len(local_indices), + "global_batch_shape": shared, + } + layout["fingerprint"] = stable_fingerprint(layout) + layouts[global_index] = layout + + for global_index, layout in enumerate(layouts): + if layout is not None: + continue + response_length = int(response_lengths[global_index]) + active_tokens = count_active_tokens(loss_masks[global_index], response_length) + dropped_layout = { + "schema_version": CONSISTENCY_METADATA_SCHEMA_VERSION, + "sample_index": sample_indices[global_index], + "rollout_id": rollout_ids[global_index], + "total_length": int(total_lengths[global_index]), + "response_length": response_length, + "active_token_count": active_tokens, + "active_mask_density": active_tokens / response_length if response_length else 0.0, + "dropped_by_schedule": True, + "global_batch_shape": shared, + } + dropped_layout["fingerprint"] = stable_fingerprint(dropped_layout) + layouts[global_index] = dropped_layout + return [layout for layout in layouts if layout is not None] diff --git a/vime/utils/types.py b/vime/utils/types.py index 1e46c99c..d51255e3 100644 --- a/vime/utils/types.py +++ b/vime/utils/types.py @@ -119,6 +119,7 @@ class Sample: loss_mask: list[int] | None = None weight_versions: list[str] = field(default_factory=list) rollout_log_probs: list[float] | None = None # Log probabilities from rollout engine + consistency_metadata: dict[str, Any] | None = None # Ragged top-p nucleus token ids replayed from rollout sampling. For response # token i, kept ids are rollout_top_p_token_ids[offsets[i]:offsets[i + 1]]. rollout_top_p_token_ids: list[int] | torch.Tensor | None = None