From 6a2308df34e1243345a89cb314588a2cc681f2f4 Mon Sep 17 00:00:00 2001 From: lamentropetion <3051000145@qq.com> Date: Thu, 17 Sep 2026 23:59:48 +0800 Subject: [PATCH 1/2] fix(vime): avoid double temperature scaling for reused logits Signed-off-by: lamentropetion <3051000145@qq.com> --- .../integrations/vime/linear_logp_provider.py | 13 +++++-- tests/test_vime_linear_logp_provider.py | 36 +++++++++++++++++++ 2 files changed, 47 insertions(+), 2 deletions(-) diff --git a/rl_engine/integrations/vime/linear_logp_provider.py b/rl_engine/integrations/vime/linear_logp_provider.py index 476ba5db..e80a1e04 100644 --- a/rl_engine/integrations/vime/linear_logp_provider.py +++ b/rl_engine/integrations/vime/linear_logp_provider.py @@ -220,6 +220,14 @@ def _is_identity_temperature(value: Any) -> bool: return value is None or (not isinstance(value, torch.Tensor) and float(value) == 1.0) +def _local_logits_temperature(request: Any) -> Any: + """Return scaling still required for logits supplied by Vime.""" + + if _metadata(request).get("logits_are_temperature_scaled") is True: + return None + return getattr(request, "temperature", None) + + @torch.no_grad() def _metric_entropy_from_strict_lse( local_logits: torch.Tensor, @@ -299,11 +307,12 @@ def _provider_impl(request: Any, *, linear_logp: Any = None) -> LinearLogpResult reuse_local_logits = True with_entropy = bool(getattr(request, "with_entropy", False)) with_entropy_grad = bool(getattr(request, "with_entropy_grad", False)) + local_logits_temperature = _local_logits_temperature(request) fast_metric_entropy = ( reuse_local_logits and with_entropy and not with_entropy_grad - and _is_identity_temperature(getattr(request, "temperature", None)) + and _is_identity_temperature(local_logits_temperature) ) strict_lse = None if reuse_local_logits: @@ -325,7 +334,7 @@ def _provider_impl(request: Any, *, linear_logp: Any = None) -> LinearLogpResult global_vocab_size=int(partition.padded_size), real_vocab_size=int(partition.real_size), target="training", - temperature=getattr(request, "temperature", None), + temperature=local_logits_temperature, return_lse=fast_metric_entropy, diagnostics_hidden=hidden, diagnostics_lm_head_weight=projection.weight, diff --git a/tests/test_vime_linear_logp_provider.py b/tests/test_vime_linear_logp_provider.py index d8071a6f..3a5d48bf 100644 --- a/tests/test_vime_linear_logp_provider.py +++ b/tests/test_vime_linear_logp_provider.py @@ -152,6 +152,42 @@ def from_local_logits(self, local_logits, target_ids, **_kwargs): assert result.provenance["execution"]["role"] == "vime_training_linear_logp" +@pytest.mark.parametrize( + ("logits_are_temperature_scaled", "expected_temperature"), + ((True, None), (False, 0.7)), +) +def test_provider_scales_reused_local_logits_at_most_once( + monkeypatch, logits_are_temperature_scaled, expected_temperature +): + monkeypatch.setenv("VIME_RL_KERNEL_STRICT", "1") + observed = {} + + class FakeLinearLogp: + backend_id = "fake-linear-logp" + provenance = {"actual_backend": "fake-linear-logp"} + + def from_local_logits(self, local_logits, target_ids, **kwargs): + observed["temperature"] = kwargs["temperature"] + return torch.log_softmax(local_logits[:, :7], dim=-1)[ + torch.arange(target_ids.size(0)), target_ids + ] + + import rl_engine.integrations.vime.linear_logp_provider as provider_module + + monkeypatch.setattr(provider_module.torch.version, "hip", "6.0") + monkeypatch.setattr( + provider_module, "_default_strict_linear_logp", lambda: FakeLinearLogp() + ) + request = _structural_request() + request.temperature = 0.7 + request.metadata["logits_are_temperature_scaled"] = logits_are_temperature_scaled + + result = provider(request) + + assert result.logp.shape == (3, 1) + assert observed["temperature"] == expected_temperature + + def test_megatron_adapter_forwards_structured_context(monkeypatch): request = _structural_request() observed = {} From 9d13d2b9cb2098eab966b747f5676e6e7bce0a2a Mon Sep 17 00:00:00 2001 From: lamentropetion <3051000145@qq.com> Date: Fri, 18 Sep 2026 00:21:16 +0800 Subject: [PATCH 2/2] test(vime): cover legacy and current temperature contracts Signed-off-by: lamentropetion <3051000145@qq.com> --- .../integrations/vime/linear_logp_provider.py | 12 ++- tests/test_vime_linear_logp_provider.py | 86 ++++++++++++++----- 2 files changed, 73 insertions(+), 25 deletions(-) diff --git a/rl_engine/integrations/vime/linear_logp_provider.py b/rl_engine/integrations/vime/linear_logp_provider.py index e80a1e04..9a181dff 100644 --- a/rl_engine/integrations/vime/linear_logp_provider.py +++ b/rl_engine/integrations/vime/linear_logp_provider.py @@ -221,7 +221,14 @@ def _is_identity_temperature(value: Any) -> bool: def _local_logits_temperature(request: Any) -> Any: - """Return scaling still required for logits supplied by Vime.""" + """Return scaling still required by the Vime local-logits contract. + + Older Vime revisions scale ``request.logits`` before provider dispatch and + advertise that fact with ``logits_are_temperature_scaled``. Current Vime + revisions dispatch unscaled logits and leave the marker absent. Only the + reused-local-logits path consumes this distinction; recomputation from + hidden states always starts from unscaled values. + """ if _metadata(request).get("logits_are_temperature_scaled") is True: return None @@ -298,7 +305,8 @@ def _provider_impl(request: Any, *, linear_logp: Any = None) -> LinearLogpResult and isinstance(request_logits, torch.Tensor) and request_logits.ndim == 2 and request_logits.dtype in (torch.bfloat16, torch.float16, torch.float32) - and request_logits.shape == ( + and request_logits.shape + == ( hidden.size(0), projection.weight.size(0), ) diff --git a/tests/test_vime_linear_logp_provider.py b/tests/test_vime_linear_logp_provider.py index 3a5d48bf..c97dadd1 100644 --- a/tests/test_vime_linear_logp_provider.py +++ b/tests/test_vime_linear_logp_provider.py @@ -11,13 +11,10 @@ import torch from rl_engine.integrations import framework_operators -from rl_engine.integrations.framework_operators import MegatronLogpOperator from rl_engine.integrations.ablation import IntegrationPlan +from rl_engine.integrations.framework_operators import MegatronLogpOperator from rl_engine.integrations.megatron import MegatronIntegration -from rl_engine.integrations.state import ( - clear_active_integration, - set_active_integration, -) +from rl_engine.integrations.state import clear_active_integration, set_active_integration from rl_engine.integrations.vime.linear_logp_provider import ( LinearLogpProviderUnavailable, LinearLogpResult, @@ -106,9 +103,7 @@ def test_provider_entropy_preserves_vime_semantics_and_autograd(): result = provider(request) reference_logits = request.logits.detach().clone().requires_grad_(True) log_probs = torch.log_softmax(reference_logits[:, :7], dim=-1) - reference_logp = log_probs[ - torch.arange(reference_logits.size(0)), request.target_ids - ] + reference_logp = log_probs[torch.arange(reference_logits.size(0)), request.target_ids] reference_entropy = -(log_probs.exp() * log_probs).sum(dim=-1) torch.testing.assert_close(result.logp.squeeze(-1), reference_logp) @@ -141,9 +136,7 @@ def from_local_logits(self, local_logits, target_ids, **_kwargs): import rl_engine.integrations.vime.linear_logp_provider as provider_module - monkeypatch.setattr( - provider_module, "_default_strict_linear_logp", lambda: FakeLinearLogp() - ) + monkeypatch.setattr(provider_module, "_default_strict_linear_logp", lambda: FakeLinearLogp()) request = _structural_request() result = provider(request) @@ -153,11 +146,12 @@ def from_local_logits(self, local_logits, target_ids, **_kwargs): @pytest.mark.parametrize( - ("logits_are_temperature_scaled", "expected_temperature"), + ("pre_scaled", "expected_temperature"), ((True, None), (False, 0.7)), + ids=("legacy-vime-pre-scaled", "current-vime-unscaled"), ) -def test_provider_scales_reused_local_logits_at_most_once( - monkeypatch, logits_are_temperature_scaled, expected_temperature +def test_provider_applies_temperature_once_for_both_vime_contracts( + monkeypatch, pre_scaled, expected_temperature ): monkeypatch.setenv("VIME_RL_KERNEL_STRICT", "1") observed = {} @@ -167,25 +161,73 @@ class FakeLinearLogp: provenance = {"actual_backend": "fake-linear-logp"} def from_local_logits(self, local_logits, target_ids, **kwargs): - observed["temperature"] = kwargs["temperature"] - return torch.log_softmax(local_logits[:, :7], dim=-1)[ + temperature = kwargs["temperature"] + observed["temperature"] = temperature + effective_logits = local_logits.float() + if temperature is not None: + effective_logits = effective_logits / temperature + return torch.log_softmax(effective_logits[:, :7], dim=-1)[ torch.arange(target_ids.size(0)), target_ids ] import rl_engine.integrations.vime.linear_logp_provider as provider_module monkeypatch.setattr(provider_module.torch.version, "hip", "6.0") - monkeypatch.setattr( - provider_module, "_default_strict_linear_logp", lambda: FakeLinearLogp() - ) + monkeypatch.setattr(provider_module, "_default_strict_linear_logp", lambda: FakeLinearLogp()) request = _structural_request() request.temperature = 0.7 - request.metadata["logits_are_temperature_scaled"] = logits_are_temperature_scaled + unscaled_logits = request.logits.detach().clone() + if pre_scaled: + # Vime c80200e (used by the PR #400 run) divides before provider + # dispatch and publishes this marker. Vime PRs #423/#424 dispatch + # unscaled logits and omit the marker. + request.logits = request.logits / request.temperature + request.metadata["logits_are_temperature_scaled"] = True result = provider(request) + expected = torch.log_softmax(unscaled_logits[:, :7] / request.temperature, dim=-1)[ + torch.arange(request.target_ids.size(0)), request.target_ids + ] assert result.logp.shape == (3, 1) assert observed["temperature"] == expected_temperature + torch.testing.assert_close(result.logp.squeeze(-1), expected) + + +def test_pre_scaled_marker_does_not_change_hidden_recomputation_temperature( + monkeypatch, +): + """The compatibility marker applies only when the supplied logits are reused.""" + + monkeypatch.setenv("VIME_RL_KERNEL_STRICT", "1") + observed = {} + + class FakeLinearLogp: + backend_id = "fake-linear-logp" + provenance = {"actual_backend": "fake-linear-logp"} + + def __call__(self, hidden, weight, target_ids, bias, **kwargs): + observed["temperature"] = kwargs["temperature"] + logits = hidden @ weight.transpose(0, 1) + if bias is not None: + logits = logits + bias + logits = logits / kwargs["temperature"] + return torch.log_softmax(logits[:, :7], dim=-1)[ + torch.arange(target_ids.size(0)), target_ids + ] + + import rl_engine.integrations.vime.linear_logp_provider as provider_module + + monkeypatch.setattr(provider_module.torch.version, "hip", None) + monkeypatch.setattr(provider_module, "_default_strict_linear_logp", lambda: FakeLinearLogp()) + request = _structural_request() + request.temperature = 0.7 + request.metadata["logits_are_temperature_scaled"] = True + + result = provider(request) + + assert result.logp.shape == (3, 1) + assert observed["temperature"] == 0.7 def test_megatron_adapter_forwards_structured_context(monkeypatch): @@ -204,9 +246,7 @@ def provider(actual_request, *, linear_logp): ) wrapper = SimpleNamespace(backend_id="fake-linear-logp", provenance={}) - monkeypatch.setattr( - framework_operators, "_require_nvidia_cuda", lambda *_args: None - ) + monkeypatch.setattr(framework_operators, "_require_nvidia_cuda", lambda *_args: None) result = MegatronLogpOperator(provider, linear_logp=wrapper)(request) assert observed["context"] is request.context