fix: Preserve sampled KL metric in TinkerNativeBackend

Author: vivekkalyan

src/art/tinker_native/backend.py

@@ -350,23 +350,23 @@ async def train(  # type: ignore[override]
                 train_tokens, pricing
             )
         trainer_started = time.monotonic()
+        sampled_kl_policy_ref: float | None = None
 
         if kl_penalty_coef > 0:
-            reference_sampling_client = await self._get_kl_reference_sampling_client(
-                state,
-                model.base_model,
-                kl_penalty_reference_step,
-            )
-            metrics.update(
-                await self._tinker_sample_call(
-                    "apply_kl_penalty",
-                    _apply_kl_penalty(
-                        datums,
-                        reference_sampling_client,
-                        kl_penalty_coef,
+            kl_metrics = await self._tinker_sample_call(
+                "apply_kl_penalty",
+                _apply_kl_penalty(
+                    datums,
+                    await self._get_kl_reference_sampling_client(
+                        state,
+                        model.base_model,
+                        kl_penalty_reference_step,
                     ),
-                )
+                    kl_penalty_coef,
+                ),
             )
+            sampled_kl_policy_ref = kl_metrics["loss/kl_policy_ref"]
+            metrics.update(kl_metrics)
 
         if adam_params is None:
             adam_params = tinker.AdamParams(
@@ -405,13 +405,23 @@ def remove_mask(datum: tinker.Datum) -> tinker.Datum:
                 if value is None:
                     continue
                 canonical_key = _canonicalize_upstream_metric_key(key)
+                if (
+                    sampled_kl_policy_ref is not None
+                    and canonical_key == "loss/kl_policy_ref"
+                ):
+                    continue
                 if canonical_key:
                     metrics[canonical_key] = float(value)
         if optim_output.metrics:
             for key, value in optim_output.metrics.items():
                 if value is None:
                     continue
                 canonical_key = _canonicalize_upstream_metric_key(key)
+                if (
+                    sampled_kl_policy_ref is not None
+                    and canonical_key == "loss/kl_policy_ref"
+                ):
+                    continue
                 if canonical_key:
                     metrics[canonical_key] = float(value)
 

tests/integration/test_tinker_native_backend.py

@@ -9,6 +9,8 @@
 
 import art
 from art.tinker_native import TinkerNativeBackend
+from art.tinker_native.backend import _apply_kl_penalty
+from art.tinker_native.data import trajectory_groups_to_datums
 
 DEFAULT_BASE_MODEL = "Qwen/Qwen3-30B-A3B-Instruct-2507"
 
@@ -37,6 +39,8 @@ async def simple_rollout(
         max_tokens=10,
         timeout=60,
         temperature=1,
+        logprobs=True,
+        top_logprobs=0,
     )
     choice = chat_completion.choices[0]
     content = (choice.message.content or "").lower()
@@ -115,6 +119,85 @@ async def make_group(prompt: str) -> art.TrajectoryGroup:
             await backend.close()
 
 
+@pytest.mark.skipif(
+    "TINKER_API_KEY" not in os.environ,
+    reason="TINKER_API_KEY not set - skipping TinkerNativeBackend KL test",
+)
+async def test_tinker_native_backend_kl_identity_metric():
+    model_name = f"test-tinker-native-kl-{uuid.uuid4().hex[:8]}"
+    with tempfile.TemporaryDirectory() as tmpdir:
+        backend = TinkerNativeBackend(path=tmpdir)
+        model = art.TrainableModel(
+            name=model_name,
+            project="integration-tests",
+            base_model=get_base_model(),
+        )
+        try:
+            await model.register(backend)
+
+            openai_client = model.openai_client()
+            current_step = await model.get_step()
+            model_name_step = model.get_inference_name(step=current_step)
+            prompts = ["Say yes", "Say no", "Say maybe"]
+
+            async def make_group(prompt: str) -> art.TrajectoryGroup:
+                import asyncio
+
+                trajectories = await asyncio.gather(
+                    *[
+                        simple_rollout(openai_client, model_name_step, prompt)
+                        for _ in range(2)
+                    ]
+                )
+                return art.TrajectoryGroup(trajectories)  # type: ignore[attr-defined]
+
+            train_groups = await art.gather_trajectory_groups(  # type: ignore[attr-defined]
+                [make_group(prompt) for prompt in prompts]
+            )
+            ensure_reward_variance(train_groups)
+
+            state = backend._model_state[model.name]
+            datums = trajectory_groups_to_datums(
+                train_groups,
+                state.renderer,
+                state.tokenizer,
+            )
+            assert datums
+
+            reference_sampling_client = await backend._get_kl_reference_sampling_client(
+                state,
+                model.base_model,
+                current_step,
+            )
+            expected_kl = (
+                await _apply_kl_penalty(
+                    trajectory_groups_to_datums(
+                        train_groups,
+                        state.renderer,
+                        state.tokenizer,
+                    ),
+                    reference_sampling_client,
+                    kl_penalty_coef=0.25,
+                )
+            )["loss/kl_policy_ref"]
+
+            result = await backend.train(
+                model,
+                train_groups,
+                learning_rate=1e-5,
+                kl_penalty_coef=0.25,
+                kl_penalty_reference_step=current_step,
+            )
+
+            assert result.metrics["loss/kl_policy_ref"] == pytest.approx(
+                expected_kl,
+                abs=0.05,
+            )
+            assert result.metrics["loss/kl_policy_ref"] == pytest.approx(0.0, abs=0.05)
+        finally:
+            await backend.close()
+
+
 @pytest.mark.skipif(
     "TINKER_API_KEY" not in os.environ,
     reason="TINKER_API_KEY not set - skipping TinkerNativeBackend fork test",