From 0ea526a382ec6026b5c97b228655c09fe69397bb Mon Sep 17 00:00:00 2001 From: inaniloquentee <3051000145@qq.com> Date: Wed, 15 Jul 2026 17:56:26 +0800 Subject: [PATCH 1/2] Add audit-only dlogp diagnostics Co-authored-by: OpenAI Codex Signed-off-by: inaniloquentee <3051000145@qq.com> --- tests/utils/test_dlogp_diagnostics.py | 337 ++++++++++++++++++++ vime/backends/megatron_utils/data.py | 1 + vime/backends/megatron_utils/loss.py | 27 ++ vime/ray/rollout.py | 1 + vime/utils/arguments.py | 11 + vime/utils/dlogp_diagnostics.py | 433 ++++++++++++++++++++++++++ 6 files changed, 810 insertions(+) create mode 100644 tests/utils/test_dlogp_diagnostics.py create mode 100644 vime/utils/dlogp_diagnostics.py diff --git a/tests/utils/test_dlogp_diagnostics.py b/tests/utils/test_dlogp_diagnostics.py new file mode 100644 index 00000000..507ea8ed --- /dev/null +++ b/tests/utils/test_dlogp_diagnostics.py @@ -0,0 +1,337 @@ +from __future__ import annotations + +import importlib +import math +import sys +from argparse import Namespace +from pathlib import Path + +import pytest +import torch + +_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 + +from vime.utils.dlogp_diagnostics import compute_dlogp_diagnostics, get_rlk_consistency_mode, is_dlogp_audit_enabled + +NUM_GPUS = 0 + + +FULL_METADATA = { + "model_name": "unit-model", + "backend_id": "megatron", + "contract_id": "contract-v1", + "batch_layout_fingerprint": "layout-abc", + "provenance_fingerprint": "prov-def", +} + +LOSS_MODULE_PATH = "vime.backends.megatron_utils.loss" +MEGATRON_STUB_MODULES = ( + "megatron", + "megatron.core", + "megatron.core.parallel_state", + "megatron.core.transformer", + "megatron.core.transformer.transformer_layer", + LOSS_MODULE_PATH, +) + + +def _report_with_full_metadata(**kwargs): + return compute_dlogp_diagnostics( + model_name=FULL_METADATA["model_name"], + backend_id=FULL_METADATA["backend_id"], + contract_id=FULL_METADATA["contract_id"], + batch_layout_fingerprint=FULL_METADATA["batch_layout_fingerprint"], + provenance_fingerprint=FULL_METADATA["provenance_fingerprint"], + **kwargs, + ) + + +@pytest.fixture() +def megatron_loss_module(): + saved = _unit_stubs.save_sys_modules(MEGATRON_STUB_MODULES) + for module_name in MEGATRON_STUB_MODULES: + sys.modules.pop(module_name, None) + _unit_stubs.install_megatron_mpu_stub() + try: + yield importlib.import_module(LOSS_MODULE_PATH) + finally: + _unit_stubs.restore_sys_modules(saved) + + +def _policy_args(mode: str) -> Namespace: + return Namespace( + use_rollout_logprobs=False, + rollout_top_p=1.0, + use_opsm=False, + advantage_estimator="grpo", + eps_clip=0.2, + eps_clip_high=0.2, + get_mismatch_metrics=False, + use_tis=False, + custom_pg_loss_reducer_function_path=None, + entropy_coef=0.0, + use_kl_loss=False, + rlk_consistency_mode=mode, + model_name=FULL_METADATA["model_name"], + train_backend=FULL_METADATA["backend_id"], + rlk_contract_id=FULL_METADATA["contract_id"], + rlk_batch_layout_fingerprint=FULL_METADATA["batch_layout_fingerprint"], + rlk_provenance_fingerprint=FULL_METADATA["provenance_fingerprint"], + ) + + +def _policy_batch() -> dict: + return { + "advantages": [torch.ones(2)], + "log_probs": [torch.zeros(2)], + "rollout_log_probs": [torch.zeros(2)], + "response_lengths": [2], + "total_lengths": [2], + "loss_masks": [torch.ones(2, dtype=torch.int32)], + "unconcat_tokens": [torch.tensor([1, 2])], + "sample_indices": [42], + "rollout_ids": [7], + } + + +@pytest.mark.unit +def test_dlogp_metrics_use_active_tokens_only_and_identify_worst_token(): + train_log_probs = [ + torch.tensor([-1.0, -2.0, -3.0]), + torch.tensor([-0.1, -0.2, 10.0]), + ] + rollout_log_probs = [ + torch.tensor([-1.0, -1.5, -4.0]), + torch.tensor([-0.6, -0.2, -10.0]), + ] + loss_masks = [torch.tensor([1, 0, 1]), torch.tensor([1, 1, 0])] + + report = _report_with_full_metadata( + train_log_probs=train_log_probs, + rollout_log_probs=rollout_log_probs, + loss_masks=loss_masks, + sample_indices=torch.tensor([10, 11]), + rollout_ids=torch.tensor([3, 4]), + rank=7, + eps_clip=0.2, + ) + + # Active dlogp values are [0.0, 1.0, 0.5, 0.0]. The masked 20.0 delta must not win. + metrics = report.metrics + assert metrics["rlk_audit_active_token_count"].item() == pytest.approx(4.0) + assert metrics["rlk_audit_mask_coverage"].item() == pytest.approx(4.0 / 6.0) + assert metrics["rlk_audit_dlogp_abs_mean"].item() == pytest.approx(0.375) + assert metrics["rlk_audit_dlogp_abs_max"].item() == pytest.approx(1.0) + assert metrics["rlk_audit_dlogp_abs_p50"].item() == pytest.approx(0.25) + assert metrics["rlk_audit_dlogp_abs_p90"].item() == pytest.approx(0.85) + assert metrics["rlk_audit_dlogp_abs_p99"].item() == pytest.approx(0.985) + assert metrics["rlk_audit_warning_count"].item() == pytest.approx(0.0) + + assert report.worst_token == { + "abs_dlogp": pytest.approx(1.0), + "dlogp": pytest.approx(1.0), + "sample_position": 0, + "sample_id": 10, + "sample_index": 10, + "rollout_id": 3, + "token_position": 2, + "rank": 7, + **FULL_METADATA, + } + assert metrics["rlk_audit_worst_sample_index"].item() == pytest.approx(10.0) + assert metrics["rlk_audit_worst_rollout_id"].item() == pytest.approx(3.0) + assert metrics["rlk_audit_worst_token_position"].item() == pytest.approx(2.0) + assert metrics["rlk_audit_worst_rank"].item() == pytest.approx(7.0) + + +@pytest.mark.unit +def test_dlogp_ratio_clipfrac_and_approx_kl_follow_issue_formulas(): + dlogp = torch.tensor([0.0, math.log(1.5), math.log(0.75), math.log(1.1)]) + report = _report_with_full_metadata( + train_log_probs=[dlogp], + rollout_log_probs=[torch.zeros_like(dlogp)], + loss_masks=[torch.ones_like(dlogp)], + eps_clip=0.2, + ) + + ratio0 = dlogp.exp() + expected_clipfrac = ((ratio0 - 1.0).abs() > 0.2).float().mean() + expected_approx_kl = (ratio0 - 1.0 - dlogp).mean() + + assert report.metrics["rlk_audit_ratio0_mean"].item() == pytest.approx(ratio0.mean().item()) + assert report.metrics["rlk_audit_clipfrac0"].item() == pytest.approx(expected_clipfrac.item()) + assert report.metrics["rlk_audit_approx_kl0"].item() == pytest.approx(expected_approx_kl.item()) + + +@pytest.mark.unit +def test_dlogp_zero_active_sample_is_reported_without_worst_token(): + report = _report_with_full_metadata( + train_log_probs=[torch.tensor([1.0, 2.0])], + rollout_log_probs=[torch.tensor([1.0, 2.0])], + loss_masks=[torch.tensor([0, 0])], + ) + + assert report.metrics["rlk_audit_active_token_count"].item() == pytest.approx(0.0) + assert report.metrics["rlk_audit_mask_coverage"].item() == pytest.approx(0.0) + assert report.metrics["rlk_audit_zero_active_sample_count"].item() == pytest.approx(1.0) + assert report.metrics["rlk_audit_worst_token_position"].item() == pytest.approx(-1.0) + assert report.worst_token is None + assert [warning.code for warning in report.warnings] == ["zero_active_tokens"] + + +@pytest.mark.unit +def test_dlogp_shape_mismatch_is_a_structured_warning_and_skips_sample(): + report = _report_with_full_metadata( + train_log_probs=[torch.tensor([1.0, 2.0])], + rollout_log_probs=[torch.tensor([1.0])], + loss_masks=[torch.tensor([1, 1])], + sample_indices=[123], + ) + + assert report.metrics["rlk_audit_active_token_count"].item() == pytest.approx(0.0) + assert report.metrics["rlk_audit_warning_count"].item() == pytest.approx(1.0) + assert report.warnings[0].code == "shape_mismatch" + assert report.warnings[0].sample_id == 123 + + +@pytest.mark.unit +def test_dlogp_missing_rollout_log_probs_is_a_structured_warning(): + report = _report_with_full_metadata( + train_log_probs=[torch.tensor([1.0, 2.0])], + rollout_log_probs=None, + loss_masks=[torch.tensor([1, 1])], + ) + + assert report.metrics["rlk_audit_active_token_count"].item() == pytest.approx(0.0) + assert report.warnings[0].code == "missing_rollout_log_probs" + assert report.warnings[0].field == "rollout_log_probs" + + +@pytest.mark.unit +def test_dlogp_missing_metadata_produces_structured_warnings(): + report = compute_dlogp_diagnostics( + train_log_probs=[torch.tensor([1.0])], + rollout_log_probs=[torch.tensor([0.0])], + loss_masks=[torch.tensor([1])], + ) + + warning_fields = {warning.field for warning in report.warnings if warning.code == "missing_metadata"} + assert warning_fields == { + "model_name", + "backend_id", + "contract_id", + "batch_layout_fingerprint", + "provenance_fingerprint", + } + assert report.metrics["rlk_audit_warning_count"].item() == pytest.approx(5.0) + + +@pytest.mark.unit +def test_dlogp_sample_metadata_can_supply_context_fields(): + metadata = [ + { + "model_name": "sample-model", + "backend_id": "sample-backend", + "contract_id": "sample-contract", + "batch_layout_fingerprint": "sample-layout", + "provenance_fingerprint": "sample-provenance", + } + ] + + report = compute_dlogp_diagnostics( + train_log_probs=[torch.tensor([2.0])], + rollout_log_probs=[torch.tensor([0.0])], + loss_masks=[torch.tensor([1])], + metadata=metadata, + ) + + assert not report.warnings + assert report.worst_token is not None + assert report.worst_token["model_name"] == "sample-model" + assert report.worst_token["backend_id"] == "sample-backend" + assert report.worst_token["contract_id"] == "sample-contract" + assert report.worst_token["batch_layout_fingerprint"] == "sample-layout" + assert report.worst_token["provenance_fingerprint"] == "sample-provenance" + + +@pytest.mark.unit +def test_dlogp_diagnostics_are_read_only_and_detached(): + train = torch.tensor([1.0, 2.0], requires_grad=True) + rollout = torch.tensor([0.5, 1.5]) + mask = torch.tensor([1, 0]) + train_before = train.detach().clone() + rollout_before = rollout.clone() + mask_before = mask.clone() + + report = _report_with_full_metadata( + train_log_probs=[train], + rollout_log_probs=[rollout], + loss_masks=[mask], + ) + + torch.testing.assert_close(train.detach(), train_before) + torch.testing.assert_close(rollout, rollout_before) + torch.testing.assert_close(mask, mask_before) + assert all(not metric.requires_grad for metric in report.metrics.values()) + + +@pytest.mark.unit +def test_rlk_consistency_mode_defaults_to_off_and_supports_args_env_and_aliases(): + assert get_rlk_consistency_mode(None, environ={}) == "off" + assert not is_dlogp_audit_enabled(Namespace(), environ={}) + + assert get_rlk_consistency_mode(Namespace(rlk_consistency_mode="audit"), environ={}) == "audit" + assert is_dlogp_audit_enabled(Namespace(rlk_consistency_mode="strict"), environ={}) + assert is_dlogp_audit_enabled(Namespace(rlk_consistency_mode="audit"), environ={"VIME_RLK_CONSISTENCY": "off"}) + assert get_rlk_consistency_mode(Namespace(), environ={"VIME_RLK_CONSISTENCY": "audit-only"}) == "audit" + assert get_rlk_consistency_mode(Namespace(), environ={"VIME_RL_KERNEL_CONSISTENCY": "true"}) == "audit" + + with pytest.raises(ValueError, match="Unsupported RL-Kernel consistency mode"): + get_rlk_consistency_mode(Namespace(rlk_consistency_mode="fast"), environ={}) + + +@pytest.mark.unit +def test_policy_loss_adds_audit_metrics_only_when_enabled_without_changing_loss(monkeypatch, megatron_loss_module): + train_log_probs = [torch.tensor([0.1, 0.3])] + + def fake_get_log_probs_and_entropy(*args, **kwargs): + return None, {"log_probs": train_log_probs, "entropy": [torch.zeros(2)]} + + def fake_compute_policy_loss(ppo_kl, advantages, eps_clip, eps_clip_high): + del advantages, eps_clip, eps_clip_high + return torch.ones_like(ppo_kl), torch.zeros_like(ppo_kl) + + monkeypatch.setattr(megatron_loss_module, "get_log_probs_and_entropy", fake_get_log_probs_and_entropy) + monkeypatch.setattr(megatron_loss_module, "compute_policy_loss", fake_compute_policy_loss) + + def reducer(tensor): + return tensor.mean() + + logits = torch.zeros(1, 2, 4) + off_loss, off_metrics = megatron_loss_module.policy_loss_function( + _policy_args("off"), + _policy_batch(), + logits, + reducer, + ) + audit_loss, audit_metrics = megatron_loss_module.policy_loss_function( + _policy_args("audit"), + _policy_batch(), + logits, + reducer, + ) + + torch.testing.assert_close(audit_loss, off_loss) + assert not any(key.startswith("rlk_audit_") for key in off_metrics) + assert audit_metrics["rlk_audit_active_token_count"].item() == pytest.approx(2.0) + assert audit_metrics["rlk_audit_dlogp_abs_mean"].item() == pytest.approx(0.2) + assert audit_metrics["rlk_audit_worst_sample_index"].item() == pytest.approx(42.0) + assert audit_metrics["rlk_audit_worst_rollout_id"].item() == pytest.approx(7.0) + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__])) diff --git a/vime/backends/megatron_utils/data.py b/vime/backends/megatron_utils/data.py index b5ec48f6..c1f4cdf7 100644 --- a/vime/backends/megatron_utils/data.py +++ b/vime/backends/megatron_utils/data.py @@ -290,6 +290,7 @@ def log_rollout_data( "global_batch_sizes", "num_microbatches", "micro_batch_indices", + "metadata", "source_names", ]: continue diff --git a/vime/backends/megatron_utils/loss.py b/vime/backends/megatron_utils/loss.py index 2003dd59..c24b5323 100644 --- a/vime/backends/megatron_utils/loss.py +++ b/vime/backends/megatron_utils/loss.py @@ -10,6 +10,7 @@ from torch.utils.checkpoint import checkpoint from vime.utils.distributed_utils import distributed_masked_whiten +from vime.utils.dlogp_diagnostics import compute_dlogp_diagnostics, is_dlogp_audit_enabled from vime.utils.misc import load_function from vime.utils.ppo_utils import ( calculate_log_probs_and_entropy, @@ -38,6 +39,12 @@ ) +def _get_dist_rank_or_none() -> int | None: + if not dist.is_available() or not dist.is_initialized(): + return None + return dist.get_rank() + + def get_rollout_top_p_logprob_kwargs(args: Namespace, batch: dict[str, Any]) -> dict[str, Any]: if args.rollout_top_p == 1.0: return {} @@ -932,6 +939,7 @@ def policy_loss_function( ) log_probs = log_probs_and_entropy["log_probs"] + audit_train_log_probs = log_probs if not args.use_rollout_logprobs and not old_log_probs: old_log_probs = [log_prob.detach() for log_prob in log_probs] train_log_probs_for_tis = batch.get("log_probs") @@ -1083,6 +1091,24 @@ def policy_loss_function( log_probs_to_compare = log_probs if args.use_rollout_logprobs else old_log_probs train_rollout_logprob_abs_diff = sum_of_sample_mean((log_probs_to_compare - rollout_log_probs).abs()) + dlogp_audit_metrics = {} + if is_dlogp_audit_enabled(args): + dlogp_audit_metrics = compute_dlogp_diagnostics( + audit_train_log_probs, + batch.get("rollout_log_probs"), + batch["loss_masks"], + sample_indices=batch.get("sample_indices"), + rollout_ids=batch.get("rollout_ids"), + metadata=batch.get("metadata"), + rank=_get_dist_rank_or_none(), + model_name=getattr(args, "model_name", None), + backend_id=getattr(args, "train_backend", "megatron"), + contract_id=getattr(args, "rlk_contract_id", None), + batch_layout_fingerprint=getattr(args, "rlk_batch_layout_fingerprint", None), + provenance_fingerprint=getattr(args, "rlk_provenance_fingerprint", None), + eps_clip=getattr(args, "eps_clip", 0.2), + ).metrics + reported_loss = { "loss": loss.clone().detach(), "pg_loss": pg_loss.clone().detach(), @@ -1093,6 +1119,7 @@ def policy_loss_function( if train_rollout_logprob_abs_diff is not None: reported_loss["train_rollout_logprob_abs_diff"] = train_rollout_logprob_abs_diff.clone().detach() + reported_loss.update(dlogp_audit_metrics) if args.use_kl_loss: reported_loss["kl_loss"] = kl_loss.clone().detach() diff --git a/vime/ray/rollout.py b/vime/ray/rollout.py index ec88fbfd..7fb938c4 100644 --- a/vime/ray/rollout.py +++ b/vime/ray/rollout.py @@ -872,6 +872,7 @@ def _split_train_data_by_dp(self, data): "rollout_top_p_token_ids", "rollout_top_p_token_offsets", "rollout_routed_experts", + "metadata", "source_names", "prompt", "teacher_log_probs", diff --git a/vime/utils/arguments.py b/vime/utils/arguments.py index 25fd9cc9..0b462fc2 100644 --- a/vime/utils/arguments.py +++ b/vime/utils/arguments.py @@ -1047,6 +1047,17 @@ def add_algo_arguments(parser): "If not set, we will use the logprobs from the actor model." ), ) + parser.add_argument( + "--rlk-consistency-mode", + "--rl-kernel-consistency-mode", + dest="rlk_consistency_mode", + choices=["off", "audit", "strict"], + default=None, + help=( + "RL-Kernel consistency diagnostics mode. 'audit' records native read-only dlogp telemetry " + "without enabling RL-Kernel fast operators; 'strict' currently records the same diagnostics." + ), + ) # Off-Policy Correction using Importance Sampling: https://fengyao.notion.site/off-policy-rl parser.add_argument( "--use-tis", diff --git a/vime/utils/dlogp_diagnostics.py b/vime/utils/dlogp_diagnostics.py new file mode 100644 index 00000000..7ac3ec7a --- /dev/null +++ b/vime/utils/dlogp_diagnostics.py @@ -0,0 +1,433 @@ +from __future__ import annotations + +import os +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from typing import Any + +import torch + +RLK_CONSISTENCY_MODES = ("off", "audit", "strict") +RLK_CONSISTENCY_MODE_ATTRS = ( + "rlk_consistency_mode", + "rlk_consistency", + "rl_kernel_consistency_mode", + "rl_kernel_consistency", +) +RLK_CONSISTENCY_MODE_ENVS = ( + "VIME_RLK_CONSISTENCY_MODE", + "VIME_RLK_CONSISTENCY", + "VIME_RL_KERNEL_CONSISTENCY_MODE", + "VIME_RL_KERNEL_CONSISTENCY", +) + +MISSING_METADATA_FIELDS = ( + "model_name", + "backend_id", + "contract_id", + "batch_layout_fingerprint", + "provenance_fingerprint", +) + + +@dataclass(frozen=True) +class DlogpAuditWarning: + code: str + message: str + sample_position: int | None = None + sample_id: int | str | None = None + field: str | None = None + + +@dataclass(frozen=True) +class DlogpAuditReport: + metrics: dict[str, torch.Tensor] + warnings: tuple[DlogpAuditWarning, ...] + worst_token: dict[str, Any] | None + + +def get_rlk_consistency_mode(args: Any | None, environ: Mapping[str, str] | None = None) -> str: + """Return the requested RL-Kernel consistency mode. + + The helper deliberately accepts several attribute/env names so this + audit-only feature can coexist with older or stacked branches that used + slightly different flag names. + """ + + raw_mode: Any | None = None + if args is not None: + for attr in RLK_CONSISTENCY_MODE_ATTRS: + value = getattr(args, attr, None) + if value is not None: + raw_mode = value + break + + if raw_mode is None: + env = os.environ if environ is None else environ + for env_name in RLK_CONSISTENCY_MODE_ENVS: + value = env.get(env_name) + if value is not None: + raw_mode = value + break + + if raw_mode is None: + return "off" + + mode = str(raw_mode).strip().lower().replace("_", "-") + aliases = { + "0": "off", + "false": "off", + "no": "off", + "none": "off", + "1": "audit", + "true": "audit", + "yes": "audit", + "audit-only": "audit", + } + mode = aliases.get(mode, mode) + if mode not in RLK_CONSISTENCY_MODES: + raise ValueError( + f"Unsupported RL-Kernel consistency mode {raw_mode!r}; " + f"expected one of {', '.join(RLK_CONSISTENCY_MODES)}." + ) + return mode + + +def is_dlogp_audit_enabled(args: Any | None, environ: Mapping[str, str] | None = None) -> bool: + return get_rlk_consistency_mode(args, environ=environ) in {"audit", "strict"} + + +def compute_dlogp_diagnostics( + train_log_probs: Sequence[torch.Tensor] | torch.Tensor, + rollout_log_probs: Sequence[torch.Tensor] | torch.Tensor | None, + loss_masks: Sequence[torch.Tensor] | torch.Tensor, + *, + sample_indices: Sequence[int | torch.Tensor] | torch.Tensor | None = None, + rollout_ids: Sequence[int | torch.Tensor] | torch.Tensor | None = None, + metadata: Sequence[Mapping[str, Any] | None] | None = None, + rank: int | None = None, + model_name: str | None = None, + backend_id: str | None = None, + contract_id: str | None = None, + batch_layout_fingerprint: str | None = None, + provenance_fingerprint: str | None = None, + eps_clip: float = 0.2, + prefix: str = "rlk_audit_", +) -> DlogpAuditReport: + """Compute read-only rollout-training dlogp diagnostics over active tokens.""" + + tensors = _as_tensor_list(train_log_probs) + masks = _as_tensor_list(loss_masks) + rollout_tensors = _as_tensor_list(rollout_log_probs) if rollout_log_probs is not None else [] + device = _first_device(tensors, masks, rollout_tensors) + dtype = _first_floating_dtype(tensors, rollout_tensors) + + warnings: list[DlogpAuditWarning] = [] + if not rollout_tensors: + warnings.append( + DlogpAuditWarning( + code="missing_rollout_log_probs", + message="rollout_log_probs are required for dlogp diagnostics.", + field="rollout_log_probs", + ) + ) + + if len(tensors) != len(masks): + warnings.append( + DlogpAuditWarning( + code="sample_count_mismatch", + message=f"train_log_probs has {len(tensors)} samples but loss_masks has {len(masks)}.", + ) + ) + if rollout_tensors and len(tensors) != len(rollout_tensors): + warnings.append( + DlogpAuditWarning( + code="sample_count_mismatch", + message=f"train_log_probs has {len(tensors)} samples but rollout_log_probs has {len(rollout_tensors)}.", + field="rollout_log_probs", + ) + ) + + context_metadata = { + "model_name": model_name, + "backend_id": backend_id, + "contract_id": contract_id, + "batch_layout_fingerprint": batch_layout_fingerprint, + "provenance_fingerprint": provenance_fingerprint, + } + for field, value in context_metadata.items(): + if value is None and not _metadata_field_available(metadata, field): + warnings.append( + DlogpAuditWarning( + code="missing_metadata", + message=f"{field} is not available for dlogp diagnostics.", + field=field, + ) + ) + + dlogp_parts: list[torch.Tensor] = [] + abs_parts: list[torch.Tensor] = [] + ratio_parts: list[torch.Tensor] = [] + clip_parts: list[torch.Tensor] = [] + approx_kl_parts: list[torch.Tensor] = [] + total_token_count = 0 + active_token_count = 0 + zero_active_sample_count = 0 + worst_token: dict[str, Any] | None = None + worst_abs_value: torch.Tensor | None = None + + sample_count = min(len(tensors), len(masks), len(rollout_tensors)) + with torch.no_grad(): + for sample_position in range(sample_count): + train = tensors[sample_position].detach().flatten() + rollout = rollout_tensors[sample_position].detach().flatten() + mask = masks[sample_position].detach().flatten() + sample_id = _value_at(sample_indices, sample_position) + rollout_id = _value_at(rollout_ids, sample_position) + sample_meta = ( + metadata[sample_position] if metadata is not None and sample_position < len(metadata) else None + ) + + if train.numel() != rollout.numel() or train.numel() != mask.numel(): + warnings.append( + DlogpAuditWarning( + code="shape_mismatch", + message=( + f"sample {sample_position} has train_log_probs={train.numel()}, " + f"rollout_log_probs={rollout.numel()}, loss_masks={mask.numel()}." + ), + sample_position=sample_position, + sample_id=sample_id, + ) + ) + continue + + total_token_count += int(mask.numel()) + active_mask = mask.to(dtype=torch.bool) + sample_active_count = int(active_mask.sum().item()) + active_token_count += sample_active_count + if sample_active_count == 0: + zero_active_sample_count += 1 + warnings.append( + DlogpAuditWarning( + code="zero_active_tokens", + message=f"sample {sample_position} has no active response/action tokens.", + sample_position=sample_position, + sample_id=sample_id, + ) + ) + continue + + active_positions = active_mask.nonzero(as_tuple=False).flatten() + sample_dlogp = (train.to(dtype=dtype) - rollout.to(dtype=dtype))[active_mask] + finite_mask = torch.isfinite(sample_dlogp) + if not bool(finite_mask.all().item()): + dropped = int((~finite_mask).sum().item()) + warnings.append( + DlogpAuditWarning( + code="non_finite_dlogp", + message=f"sample {sample_position} has {dropped} non-finite active dlogp values.", + sample_position=sample_position, + sample_id=sample_id, + ) + ) + active_positions = active_positions[finite_mask] + sample_dlogp = sample_dlogp[finite_mask] + if sample_dlogp.numel() == 0: + continue + + sample_abs = sample_dlogp.abs() + sample_ratio = sample_dlogp.exp() + sample_clip = ((sample_ratio - 1.0).abs() > eps_clip).to(dtype=dtype) + sample_approx_kl = sample_ratio - 1.0 - sample_dlogp + + dlogp_parts.append(sample_dlogp) + abs_parts.append(sample_abs) + ratio_parts.append(sample_ratio) + clip_parts.append(sample_clip) + approx_kl_parts.append(sample_approx_kl) + + sample_worst_abs, sample_worst_active_index = sample_abs.max(dim=0) + if worst_abs_value is None or bool((sample_worst_abs > worst_abs_value).item()): + token_position = int(active_positions[int(sample_worst_active_index.item())].item()) + worst_abs_value = sample_worst_abs + worst_token = { + "abs_dlogp": float(sample_worst_abs.item()), + "dlogp": float(sample_dlogp[int(sample_worst_active_index.item())].item()), + "sample_position": sample_position, + "sample_id": sample_id, + "sample_index": sample_id, + "rollout_id": rollout_id, + "token_position": token_position, + "rank": rank, + "model_name": _metadata_value("model_name", sample_meta, model_name), + "backend_id": _metadata_value("backend_id", sample_meta, backend_id), + "contract_id": _metadata_value("contract_id", sample_meta, contract_id), + "batch_layout_fingerprint": _metadata_value( + "batch_layout_fingerprint", + sample_meta, + batch_layout_fingerprint, + ), + "provenance_fingerprint": _metadata_value( + "provenance_fingerprint", + sample_meta, + provenance_fingerprint, + ), + } + + if dlogp_parts: + abs_dlogp = torch.cat(abs_parts).to(device=device, dtype=dtype) + ratio0 = torch.cat(ratio_parts).to(device=device, dtype=dtype) + clipfrac0_values = torch.cat(clip_parts).to(device=device, dtype=dtype) + approx_kl0_values = torch.cat(approx_kl_parts).to(device=device, dtype=dtype) + quantiles = torch.quantile(abs_dlogp, torch.tensor([0.5, 0.9, 0.99], device=device, dtype=dtype)) + metrics = { + f"{prefix}dlogp_abs_mean": abs_dlogp.mean(), + f"{prefix}dlogp_abs_max": abs_dlogp.max(), + f"{prefix}dlogp_abs_p50": quantiles[0], + f"{prefix}dlogp_abs_p90": quantiles[1], + f"{prefix}dlogp_abs_p99": quantiles[2], + f"{prefix}ratio0_mean": ratio0.mean(), + f"{prefix}clipfrac0": clipfrac0_values.mean(), + f"{prefix}approx_kl0": approx_kl0_values.mean(), + } + else: + metrics = { + f"{prefix}dlogp_abs_mean": _zero(device, dtype), + f"{prefix}dlogp_abs_max": _zero(device, dtype), + f"{prefix}dlogp_abs_p50": _zero(device, dtype), + f"{prefix}dlogp_abs_p90": _zero(device, dtype), + f"{prefix}dlogp_abs_p99": _zero(device, dtype), + f"{prefix}ratio0_mean": _zero(device, dtype), + f"{prefix}clipfrac0": _zero(device, dtype), + f"{prefix}approx_kl0": _zero(device, dtype), + } + + metrics.update( + { + f"{prefix}active_token_count": torch.tensor(float(active_token_count), device=device, dtype=dtype), + f"{prefix}mask_coverage": torch.tensor( + float(active_token_count / total_token_count) if total_token_count else 0.0, + device=device, + dtype=dtype, + ), + f"{prefix}zero_active_sample_count": torch.tensor( + float(zero_active_sample_count), device=device, dtype=dtype + ), + f"{prefix}warning_count": torch.tensor(float(len(warnings)), device=device, dtype=dtype), + } + ) + + if worst_token is not None: + metrics.update( + { + f"{prefix}worst_abs_dlogp": torch.tensor(worst_token["abs_dlogp"], device=device, dtype=dtype), + f"{prefix}worst_dlogp": torch.tensor(worst_token["dlogp"], device=device, dtype=dtype), + f"{prefix}worst_sample_position": torch.tensor( + float(worst_token["sample_position"]), + device=device, + dtype=dtype, + ), + f"{prefix}worst_sample_index": torch.tensor( + float(worst_token["sample_index"]) if _is_number(worst_token["sample_index"]) else -1.0, + device=device, + dtype=dtype, + ), + f"{prefix}worst_rollout_id": torch.tensor( + float(worst_token["rollout_id"]) if _is_number(worst_token["rollout_id"]) else -1.0, + device=device, + dtype=dtype, + ), + f"{prefix}worst_token_position": torch.tensor( + float(worst_token["token_position"]), + device=device, + dtype=dtype, + ), + f"{prefix}worst_rank": torch.tensor( + float(rank) if rank is not None else -1.0, + device=device, + dtype=dtype, + ), + } + ) + else: + metrics.update( + { + f"{prefix}worst_abs_dlogp": _zero(device, dtype), + f"{prefix}worst_dlogp": _zero(device, dtype), + f"{prefix}worst_sample_position": torch.tensor(-1.0, device=device, dtype=dtype), + f"{prefix}worst_sample_index": torch.tensor(-1.0, device=device, dtype=dtype), + f"{prefix}worst_rollout_id": torch.tensor(-1.0, device=device, dtype=dtype), + f"{prefix}worst_token_position": torch.tensor(-1.0, device=device, dtype=dtype), + f"{prefix}worst_rank": torch.tensor( + float(rank) if rank is not None else -1.0, device=device, dtype=dtype + ), + } + ) + + metrics = {key: value.clone().detach() for key, value in metrics.items()} + return DlogpAuditReport(metrics=metrics, warnings=tuple(warnings), worst_token=worst_token) + + +def _as_tensor_list(value: Sequence[torch.Tensor] | torch.Tensor | None) -> list[torch.Tensor]: + if value is None: + return [] + if isinstance(value, torch.Tensor): + return [value] + return [torch.as_tensor(item) for item in value] + + +def _first_device(*groups: Sequence[torch.Tensor]) -> torch.device: + for group in groups: + for tensor in group: + return tensor.device + return torch.device("cpu") + + +def _first_floating_dtype(*groups: Sequence[torch.Tensor]) -> torch.dtype: + for group in groups: + for tensor in group: + if tensor.is_floating_point(): + return tensor.dtype + return torch.float32 + + +def _zero(device: torch.device, dtype: torch.dtype) -> torch.Tensor: + return torch.tensor(0.0, device=device, dtype=dtype) + + +def _value_at(values: Sequence[Any] | torch.Tensor | None, index: int) -> Any: + if values is None: + return None + if isinstance(values, torch.Tensor): + if index >= values.numel(): + return None + return _python_scalar(values.flatten()[index]) + if index >= len(values): + return None + value = values[index] + if isinstance(value, torch.Tensor): + return _python_scalar(value) + return value + + +def _python_scalar(value: torch.Tensor) -> Any: + if value.numel() != 1: + return value.detach().cpu().tolist() + return value.detach().cpu().item() + + +def _metadata_value(field: str, sample_meta: Mapping[str, Any] | None, fallback: Any) -> Any: + if sample_meta is not None and field in sample_meta: + return sample_meta[field] + return fallback + + +def _metadata_field_available(metadata: Sequence[Mapping[str, Any] | None] | None, field: str) -> bool: + if metadata is None: + return False + return any(sample_meta is not None and sample_meta.get(field) is not None for sample_meta in metadata) + + +def _is_number(value: Any) -> bool: + return isinstance(value, int | float) and not isinstance(value, bool) From bd1f73ffd12b6d1f6402d8df4fff0c838fcd530a Mon Sep 17 00:00:00 2001 From: inaniloquentee <3051000145@qq.com> Date: Sun, 26 Jul 2026 19:24:14 +0800 Subject: [PATCH 2/2] test(rl-kernel): guard dlogp diagnostics boundary --- tests/utils/test_dlogp_diagnostics.py | 25 +++++++++++++++++++++++++ vime/utils/dlogp_diagnostics.py | 7 +++++++ 2 files changed, 32 insertions(+) diff --git a/tests/utils/test_dlogp_diagnostics.py b/tests/utils/test_dlogp_diagnostics.py index 507ea8ed..c24f44a7 100644 --- a/tests/utils/test_dlogp_diagnostics.py +++ b/tests/utils/test_dlogp_diagnostics.py @@ -39,6 +39,12 @@ ) +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 _report_with_full_metadata(**kwargs): return compute_dlogp_diagnostics( model_name=FULL_METADATA["model_name"], @@ -50,6 +56,25 @@ def _report_with_full_metadata(**kwargs): ) +@pytest.mark.unit +def test_dlogp_diagnostics_do_not_import_rl_engine(): + _drop_rl_engine_modules() + module = importlib.reload(importlib.import_module("vime.utils.dlogp_diagnostics")) + + report = module.compute_dlogp_diagnostics( + train_log_probs=[torch.tensor([-1.0, -2.0])], + rollout_log_probs=[torch.tensor([-1.0, -2.1])], + loss_masks=[torch.tensor([1, 1])], + metadata=[FULL_METADATA], + ) + + assert report.metrics["rlk_audit_active_token_count"].item() == pytest.approx(2.0) + assert report.worst_token is not None + assert "alignment_profile" not in report.worst_token + assert "ScoreArtifact" not in report.worst_token + assert not any(name == "rl_engine" or name.startswith("rl_engine.") for name in sys.modules) + + @pytest.fixture() def megatron_loss_module(): saved = _unit_stubs.save_sys_modules(MEGATRON_STUB_MODULES) diff --git a/vime/utils/dlogp_diagnostics.py b/vime/utils/dlogp_diagnostics.py index 7ac3ec7a..4ed44dc5 100644 --- a/vime/utils/dlogp_diagnostics.py +++ b/vime/utils/dlogp_diagnostics.py @@ -1,3 +1,10 @@ +"""Read-only dlogp audit diagnostics for vime runtime paths. + +These helpers intentionally stop at tensor and metadata telemetry. RL-Kernel +owned ScoreArtifact schemas, A0-A5 profiles, comparators, and tolerance +contracts stay behind adapter imports instead of being reimplemented here. +""" + from __future__ import annotations import os