feat: Add sampled KL support to pipeline backends

Author: vivekkalyan

src/art/dev/train.py

@@ -18,6 +18,7 @@ class TrainConfig(TypedDict, total=False):
     ]
     kimi_k2_tau: float | None
     kl_penalty_coef: float
+    kl_penalty_source: Literal["current_learner", "sample"]
     kl_ref_adapter_path: str | None
     logprob_calculation_chunk_size: int
     mask_prob_ratio: bool

src/art/local/backend.py

@@ -531,6 +531,7 @@ async def train(  # type: ignore[override]
         kl_penalty_coef: float = 0.0,
         kl_penalty_reference_step: int | None = None,
         kl_ref_adapter_path: str | None = None,
+        kl_penalty_source: Literal["current_learner", "sample"] = "current_learner",
         epsilon: float | None = None,
         epsilon_high: float | None = None,
         # Advantage computation
@@ -584,6 +585,11 @@ async def train(  # type: ignore[override]
             kl_ref_adapter_path: Direct filesystem path to a LoRA adapter
                 checkpoint to use as the KL reference. Alternative to
                 kl_penalty_reference_step.
+            kl_penalty_source: Which policy's logprobs to compare against the
+                reference when building the centered KL penalty. Use
+                "current_learner" to match the original ART implementation, or
+                "sample" to shape from the rollout policy logprobs, which is
+                usually better for async/off-policy workloads.
             epsilon: Clip epsilon for importance sampling. Defaults based on loss_fn.
             epsilon_high: Asymmetric upper clip bound. Defaults to epsilon.
             advantage_balance: Balance between negative and positive advantages
@@ -635,16 +641,20 @@ async def train(  # type: ignore[override]
             raise ValueError("LocalBackend requires normalize_advantages=True.")
         if adam_params is not None:
             raise ValueError("LocalBackend requires adam_params=None.")
+        assert kl_penalty_source in {"current_learner", "sample"}
 
         # Build config objects from explicit kwargs
         config = TrainConfig(
-            learning_rate=learning_rate, kl_penalty_coef=kl_penalty_coef
+            learning_rate=learning_rate,
+            kl_penalty_coef=kl_penalty_coef,
+            kl_penalty_source=kl_penalty_source,
         )
         dev_config: dev.TrainConfig = {
             "advantage_balance": advantage_balance,
             "allow_training_without_logprobs": allow_training_without_logprobs,
             "importance_sampling_level": importance_sampling_level,
             "kl_penalty_coef": kl_penalty_coef,
+            "kl_penalty_source": kl_penalty_source,
             "mask_prob_ratio": mask_prob_ratio,
             "plot_tensors": plot_tensors,
             "ppo": loss_fn == "ppo",

src/art/loss.py

@@ -95,7 +95,14 @@ def loss_fn(
     kl_policy_ref: torch.Tensor | None = None
     kl_penalty_coef = experimental_config.get("kl_penalty_coef", 0.0)
     if kl_penalty_coef > 0 and ref_logprobs is not None:
-        kl_per_token = (new_logprobs - ref_logprobs).detach() * assistant_mask
+        match experimental_config.get("kl_penalty_source", "current_learner"):
+            case "sample":
+                kl_source_logprobs = old_logprobs.detach()
+            case "current_learner":
+                kl_source_logprobs = new_logprobs.detach()
+            case other:
+                raise AssertionError(other)
+        kl_per_token = (kl_source_logprobs - ref_logprobs).detach() * assistant_mask
         avg_kl = kl_per_token.sum() / (assistant_mask.sum() + 1e-6)
         kl_penalty = kl_penalty_coef * (avg_kl - kl_per_token) * assistant_mask
         advantages = advantages + kl_penalty

src/art/pipeline_trainer/trainer.py

@@ -78,6 +78,8 @@ def __init__(
         loss_fn_config: dict | None = None,
         normalize_advantages: bool = True,
         adam_params: object | None = None,
+        kl_penalty_coef: float = 0.0,
+        kl_penalty_reference_step: int | None = None,
         max_steps: int | None = None,
         # Discard handling
         discard_queue_multiplier: int = 100,
@@ -129,6 +131,8 @@ def __init__(
         self.loss_fn_config = loss_fn_config
         self.normalize_advantages = normalize_advantages
         self.adam_params = adam_params
+        self.kl_penalty_coef = kl_penalty_coef
+        self.kl_penalty_reference_step = kl_penalty_reference_step
         self.max_steps = max_steps
         self._status_log_interval_seconds = log_interval_seconds
         self.eval_every_n_steps = eval_every_n_steps
@@ -452,6 +456,14 @@ async def _training_stage(self) -> None:
             if os.getenv("ART_TRAIN_STEP_LOG"):
                 print(f"[train] step {expected_step} starting (batch={len(batch)})")
             try:
+                kl_train_kwargs: dict[str, object] = {}
+                if self.kl_penalty_coef > 0.0:
+                    kl_train_kwargs["kl_penalty_coef"] = self.kl_penalty_coef
+                    kl_train_kwargs["kl_penalty_source"] = "sample"
+                    if self.kl_penalty_reference_step is not None:
+                        kl_train_kwargs["kl_penalty_reference_step"] = (
+                            self.kl_penalty_reference_step
+                        )
                 result = await self.backend.train(
                     self.model,
                     batch,
@@ -461,6 +473,7 @@ async def _training_stage(self) -> None:
                     normalize_advantages=self.normalize_advantages,
                     save_checkpoint=should_checkpoint,
                     adam_params=self.adam_params,
+                    **kl_train_kwargs,
                 )
             except Exception:
                 self._status.note_training_end()

src/art/test/test_kl_advantage.py

@@ -2,7 +2,7 @@
 
 import torch
 
-from art.loss import Loss, loss_fn
+from art.loss import loss_fn, shift_tensor
 
 
 def _make_inputs(
@@ -114,3 +114,50 @@ def test_kl_advantage_does_not_affect_when_no_ref():
 
     loss = loss_fn(inputs, new_logprobs, None, None, {"kl_penalty_coef": 0.5})
     assert loss.kl_policy_ref is None
+
+
+def test_kl_advantage_can_use_sample_logprobs() -> None:
+    """Sample-source KL should use stored rollout logprobs rather than learner logprobs."""
+    inputs = _make_inputs(seq_len=8)
+    inputs["logprobs"] = torch.tensor(
+        [[0.0, -0.2, -0.4, -0.6, -0.8, -1.0, -1.2, -1.4]], dtype=torch.float32
+    )
+    new_logprobs = torch.tensor(
+        [[0.0, -1.2, -1.1, -1.0, -0.9, -0.8, -0.7, -0.6]], dtype=torch.float32
+    )
+    ref_logprobs = torch.full((1, 8), -0.5)
+    assistant_mask = shift_tensor(inputs["assistant_mask"], False).to(
+        new_logprobs.dtype
+    )
+    sampled_logprobs = torch.where(
+        torch.isnan(shift_tensor(inputs["logprobs"], float("nan"))),
+        new_logprobs.detach(),
+        shift_tensor(inputs["logprobs"], float("nan")),
+    )
+    expected_sample_kl = ((sampled_logprobs - ref_logprobs) * assistant_mask).sum() / (
+        assistant_mask.sum() + 1e-6
+    )
+    expected_current_kl = ((new_logprobs - ref_logprobs) * assistant_mask).sum() / (
+        assistant_mask.sum() + 1e-6
+    )
+
+    sample_loss = loss_fn(
+        inputs,
+        new_logprobs,
+        ref_logprobs,
+        None,
+        {"kl_penalty_coef": 0.5, "kl_penalty_source": "sample"},
+    )
+    learner_loss = loss_fn(
+        inputs,
+        new_logprobs,
+        ref_logprobs,
+        None,
+        {"kl_penalty_coef": 0.5, "kl_penalty_source": "current_learner"},
+    )
+
+    assert sample_loss.kl_policy_ref is not None
+    assert learner_loss.kl_policy_ref is not None
+    assert torch.isclose(sample_loss.kl_policy_ref, expected_sample_kl)
+    assert torch.isclose(learner_loss.kl_policy_ref, expected_current_kl)
+    assert not torch.isclose(sample_loss.kl_policy_ref, learner_loss.kl_policy_ref)

src/art/tinker_native/backend.py

@@ -24,6 +24,7 @@
 from openai.types.chat.completion_create_params import CompletionCreateParams
 from openai.types.completion_usage import CompletionUsage
 import tinker
+import torch
 import uvicorn
 
 from art.tinker.cookbook_v import renderers, tokenizer_utils
@@ -82,6 +83,76 @@ def _canonicalize_upstream_metric_key(metric: str) -> str:
     return _UPSTREAM_TRAIN_METRIC_KEYS.get(metric, metric)
 
 
+async def _apply_kl_penalty(
+    datums: list[tinker.Datum],
+    reference_sampling_client: tinker.SamplingClient,
+    kl_penalty_coef: float,
+) -> dict[str, float]:
+    assert datums
+    assert kl_penalty_coef > 0.0
+
+    full_sequences: list[tinker.ModelInput] = []
+    sampled_logprobs_by_datum: list[torch.Tensor] = []
+    masks_by_datum: list[torch.Tensor] = []
+    advantages_by_datum: list[torch.Tensor] = []
+    for datum in datums:
+        target_tokens = datum.loss_fn_inputs["target_tokens"].to_torch()
+        assert target_tokens.numel() > 0
+        full_sequences.append(
+            datum.model_input.append_int(int(target_tokens[-1].item()))
+        )
+        sampled_logprobs_by_datum.append(datum.loss_fn_inputs["logprobs"].to_torch())
+        masks_by_datum.append(datum.loss_fn_inputs["mask"].to_torch().float())
+        advantages_by_datum.append(datum.loss_fn_inputs["advantages"].to_torch())
+
+    reference_logprobs_by_datum = await asyncio.gather(
+        *[
+            reference_sampling_client.compute_logprobs_async(full_sequence)
+            for full_sequence in full_sequences
+        ]
+    )
+
+    logprob_diffs_by_datum: list[torch.Tensor] = []
+    for reference_logprobs, sampled_logprobs, mask in zip(
+        reference_logprobs_by_datum,
+        sampled_logprobs_by_datum,
+        masks_by_datum,
+        strict=True,
+    ):
+        reference_values = reference_logprobs[1:]
+        assert len(reference_values) == sampled_logprobs.numel()
+        assert all(value is not None for value in reference_values)
+        reference_logprobs_tensor = torch.tensor(
+            reference_values,
+            dtype=sampled_logprobs.dtype,
+        )
+        logprob_diffs_by_datum.append(
+            (sampled_logprobs - reference_logprobs_tensor) * mask
+        )
+
+    total_tokens = torch.stack([mask.sum() for mask in masks_by_datum]).sum()
+    assert total_tokens.item() > 0
+    avg_logprob_diff = (
+        torch.stack(
+            [logprob_diff.sum() for logprob_diff in logprob_diffs_by_datum]
+        ).sum()
+        / total_tokens
+    )
+
+    for datum, advantages, mask, logprob_diff in zip(
+        datums,
+        advantages_by_datum,
+        masks_by_datum,
+        logprob_diffs_by_datum,
+        strict=True,
+    ):
+        datum.loss_fn_inputs["advantages"] = tinker.TensorData.from_torch(
+            advantages + kl_penalty_coef * (avg_logprob_diff - logprob_diff) * mask
+        )
+
+    return {"loss/kl_policy_ref": float(avg_logprob_diff)}
+
+
 @dataclass
 class ModelState:
     service_client: tinker.ServiceClient
@@ -239,6 +310,9 @@ async def train(  # type: ignore[override]
         save_checkpoint: bool = False,
         loss_fn_config: dict | None = None,
         adam_params: tinker.AdamParams | None = None,
+        kl_penalty_coef: float = 0.0,
+        kl_penalty_reference_step: int | None = None,
+        kl_penalty_source: Literal["sample"] = "sample",
     ) -> TrainResult:
         state = self._model_state[model.name]
         groups_list = list(trajectory_groups)
@@ -259,6 +333,10 @@ async def train(  # type: ignore[override]
             "data/step_num_datums": float(len(datums)),
         }
 
+        assert kl_penalty_source == "sample", (
+            "TinkerNativeBackend only supports kl_penalty_source='sample'."
+        )
+
         if not datums:
             return TrainResult(step=state.current_step, metrics=metrics)
 
@@ -273,6 +351,23 @@ async def train(  # type: ignore[override]
             )
         trainer_started = time.monotonic()
 
+        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,
+                    ),
+                )
+            )
+
         if adam_params is None:
             adam_params = tinker.AdamParams(
                 learning_rate=learning_rate,
@@ -697,6 +792,19 @@ async def _get_sampler_client(
         state.sampler_clients[actual_step] = sampler_client
         return sampler_client
 
+    async def _get_kl_reference_sampling_client(
+        self,
+        state: ModelState,
+        base_model: str,
+        step: int | None,
+    ) -> tinker.SamplingClient:
+        if step is not None:
+            return await self._get_sampler_client(state, step)
+        return await self._tinker_sample_call(
+            "create_sampling_client_async",
+            state.service_client.create_sampling_client_async(base_model=base_model),
+        )
+
     def _normalize_messages(self, messages: Iterable[Any]) -> list[dict[str, Any]]:
         normalized: list[dict[str, Any]] = []
         for message in messages:

src/art/types.py

@@ -17,6 +17,7 @@
 class TrainConfig(pydantic.BaseModel):
     learning_rate: float = 5e-6
     kl_penalty_coef: float = 0.0
+    kl_penalty_source: Literal["current_learner", "sample"] = "current_learner"
 
 
 class TrainSFTConfig(pydantic.BaseModel):