fix(vime): avoid double temperature scaling for reused logits

rl_engine/integrations/vime/linear_logp_provider.py

@@ -220,6 +220,21 @@ 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 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
+    return getattr(request, "temperature", None)
+
+
 @torch.no_grad()
 def _metric_entropy_from_strict_lse(
     local_logits: torch.Tensor,
@@ -290,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),
             )
@@ -299,11 +315,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 +342,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,

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)
 
@@ -152,6 +145,91 @@ def from_local_logits(self, local_logits, target_ids, **_kwargs):
     assert result.provenance["execution"]["role"] == "vime_training_linear_logp"
 
 
+@pytest.mark.parametrize(
+    ("pre_scaled", "expected_temperature"),
+    ((True, None), (False, 0.7)),
+    ids=("legacy-vime-pre-scaled", "current-vime-unscaled"),
+)
+def test_provider_applies_temperature_once_for_both_vime_contracts(
+    monkeypatch, pre_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):
+            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())
+    request = _structural_request()
+    request.temperature = 0.7
+    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):
     request = _structural_request()
     observed = {}
@@ -168,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