fix: recover per-token ids and logprobs from content[] to stop RL KLD corruption (#458)

Dranoxgithub · cursoragent · web-flow · commit 2c0867156a72 · 2026-07-06T12:57:04.000-07:00
* fix: request exact token ids and fail loud in FireworksV1CompletionsClient

The /v1/completions sampling path corrupted RL training via retokenization
drift. The request never set return_token_ids, so choices[].token_ids came
back empty and the client silently re-encoded decoded text. Retokenization
drops the trailing end-of-turn/EOS token, making completion_ids one shorter
than token_logprobs and misaligning every per-token logprob (inference KLD
spiked to ~60 vs ~0.028 with exact ids).

- Default return_token_ids=True in the request payload (overridable via
  request_params).
- Remove the tokenizer.encode() re-encode fallback; raise when exact ids are
  absent instead of silently returning corrupted data.
- Assert len(completion_token_ids) == len(completion_logprobs) at the boundary
  to catch any residual drift.

Co-authored-by: Cursor &lt;cursoragent@cursor.com&gt;

* refactor: recover per-token ids and logprobs from content[] logprobs shape

This client always requests boolean logprobs=True, so every response uses the
new content[] logprobs shape. Read the exact per-token ids and logprobs
together from content[] (token_id + sampling_logprob) so they are aligned by
construction.

Bug 1 (alignment / KLD corruption) — the real fix: previously ids were read
only from the top-level choices[].token_ids (populated by return_token_ids,
which this client never requested) and, when absent, silently re-derived via
tokenizer.encode(decode(text)). Retokenization drops the trailing
end-of-turn/EOS token, making completion_ids one shorter than the logprobs and
misaligning every per-token logprob (inference_kld ~60 vs ~0.028). Since
content[] already carries token_id per entry, reading ids there removes the
drift entirely — no return_token_ids, no re-encode fallback.

Bug 2 (precision) — more of a feature: prefer content[].sampling_logprob (the
exact value the sampler drew with) over content[].logprob (rounded).

- Drop legacy token_logprobs / top-level token_ids / raw_output id sources.
- Fail loud when content[] is absent or a content[] entry lacks token_id,
  instead of silently returning corrupted data.

Co-authored-by: Cursor &lt;cursoragent@cursor.com&gt;

* fix: require sampling_logprob per content[] entry, no fallback

Always read the per-token logprob from content[].sampling_logprob (the exact
value the sampler drew with). Fail loud if any content[] entry lacks
sampling_logprob instead of substituting the rounded content[].logprob or 0.0,
so silently degraded logprobs never reach RL training. Removes the now-unused
_extract_entry_logprob helper.

Co-authored-by: Cursor &lt;cursoragent@cursor.com&gt;

* test: use pytest.approx for logprob list comparison

Co-authored-by: Cursor &lt;cursoragent@cursor.com&gt;

---------

Co-authored-by: Cursor &lt;cursoragent@cursor.com&gt;
diff --git a/eval_protocol/integrations/fireworks_v1_completions_client.py b/eval_protocol/integrations/fireworks_v1_completions_client.py
@@ -393,35 +393,50 @@ async def create_completion_from_prompt_ids(
         finish_reason = str(choice.get("finish_reason") or "unknown")
 
         raw_output = choice.get("raw_output") if isinstance(choice.get("raw_output"), dict) else {}
-        completion_token_ids = _normalize_token_id_sequence(
-            choice.get("token_ids") or raw_output.get("completion_token_ids") or []
-        )
         choice_prompt_token_ids = _normalize_token_id_sequence(
             choice.get("prompt_token_ids") or raw_output.get("prompt_token_ids") or normalized_prompt_token_ids
         )
 
+        # -- Extract per-token ids and logprobs together --------------------
+        # Both come from the same ``content[]`` array entry-by-entry, so they are
+        # inherently the same length and aligned. Reading ids from a different
+        # source (top-level token_ids) and re-encoding decoded text when it is
+        # absent drops the trailing end-of-turn token and misaligns per-token
+        # logprobs, corrupting inference KLD — so we never do that.
+        choice_logprobs = choice.get("logprobs")
+        content_entries = choice_logprobs.get("content") if isinstance(choice_logprobs, dict) else None
+        if not content_entries:
+            raise RuntimeError(
+                "Fireworks /v1/completions returned no content[] logprobs entries. "
+                "This client requires the boolean logprobs=True (content) shape to "
+                "recover exact per-token ids and sampling logprobs. Refusing to "
+                "re-encode decoded text: retokenization drops the end-of-turn token "
+                "and misaligns per-token logprobs, corrupting inference KLD. "
+                f"choice keys={list(choice.keys())}"
+            )
+
+        completion_token_ids: List[int] = []
+        completion_logprobs: List[float] = []
+        for index, entry in enumerate(content_entries):
+            if not isinstance(entry, dict) or entry.get("token_id") is None:
+                raise RuntimeError(
+                    "Fireworks /v1/completions content[] entry is missing token_id "
+                    f"at index {index}; cannot align per-token logprobs to token ids "
+                    "without re-encoding. Refusing to return corrupted data."
+                )
+            if entry.get("sampling_logprob") is None:
+                raise RuntimeError(
+                    "Fireworks /v1/completions content[] entry is missing "
+                    f"sampling_logprob at index {index}. The sampling logprob is the "
+                    "exact value the sampler drew with and is required for correct "
+                    "inference KLD; refusing to substitute the rounded logprob or 0.0."
+                )
+            completion_token_ids.append(int(entry["token_id"]))
+            completion_logprobs.append(float(entry["sampling_logprob"]))
+
         completion_text = self.decode_token_ids(token_ids=completion_token_ids)
         if not completion_text:
             completion_text = str(choice.get("text") or "")
-        if not completion_token_ids and completion_text:
-            tokenizer = self._get_tokenizer()
-            completion_token_ids = list(tokenizer.encode(completion_text, add_special_tokens=False))
-
-        # -- Extract logprobs -----------------------------------------------
-        completion_logprobs: List[float] = []
-        choice_logprobs = choice.get("logprobs")
-        if isinstance(choice_logprobs, dict):
-            token_logprobs = choice_logprobs.get("token_logprobs") or []
-            if token_logprobs:
-                completion_logprobs = [float(lp) if lp is not None else 0.0 for lp in token_logprobs]
-            else:
-                content_logprobs = choice_logprobs.get("content") or []
-                completion_logprobs = [
-                    float(entry.get("logprob", 0.0)) if isinstance(entry, dict) else 0.0
-                    for entry in content_logprobs
-                ]
-        elif isinstance(choice_logprobs, list):
-            completion_logprobs = [float(lp) if lp is not None else 0.0 for lp in choice_logprobs]
 
         # -- Build message via parser or raw --------------------------------
         if self.tool_call_parser is not None:
diff --git a/tests/test_fireworks_v1_completions_client.py b/tests/test_fireworks_v1_completions_client.py
@@ -124,6 +124,169 @@ def encode(self, text, add_special_tokens=False):
     asyncio.run(client.close())
 
 
+class _FakeResponse:
+    def __init__(self, payload: Dict[str, Any]):
+        self._payload = payload
+
+    def model_dump(self) -> Dict[str, Any]:
+        return self._payload
+
+
+def _install_fake_completion(client, monkeypatch, payload):
+    captured: Dict[str, Any] = {}
+
+    async def fake_create(**kwargs):
+        captured.update(kwargs)
+        return _FakeResponse(payload)
+
+    class _FakeCompletions:
+        create = staticmethod(fake_create)
+
+    class _FakeClient:
+        completions = _FakeCompletions()
+
+        async def close(self):
+            return None
+
+    monkeypatch.setattr(client, "_client", _FakeClient())
+    return captured
+
+
+def test_reads_ids_and_sampling_logprobs_from_content(monkeypatch):
+    client = FireworksV1CompletionsClient(
+        model_id="test-model",
+        tokenizer_name_or_path="Qwen/Qwen3-0.6B",
+    )
+    monkeypatch.setattr(client, "decode_token_ids", lambda token_ids: "text")
+
+    def _fail_tokenizer():
+        raise AssertionError("tokenizer must not be used to re-encode completion text")
+
+    monkeypatch.setattr(client, "_get_tokenizer", lambda: _fail_tokenizer())
+    captured = _install_fake_completion(
+        client,
+        monkeypatch,
+        {
+            "choices": [
+                {
+                    "finish_reason": "stop",
+                    "logprobs": {
+                        "content": [
+                            {"token_id": 271, "sampling_logprob": -0.05483185, "logprob": -0.0548313},
+                            {"token_id": 248068, "sampling_logprob": -0.0014, "logprob": -0.0014},
+                            {"token_id": 26108, "sampling_logprob": -1.0, "logprob": -0.99},
+                        ]
+                    },
+                }
+            ],
+        },
+    )
+    result = asyncio.run(client.create_completion_from_prompt_ids(prompt_token_ids=[1, 2]))
+    assert "return_token_ids" not in captured
+    assert result["completion_ids"] == [271, 248068, 26108]
+    assert result["completion_logprobs"] == pytest.approx([-0.05483185, -0.0014, -1.0])
+    assert len(result["completion_ids"]) == len(result["completion_logprobs"])
+    asyncio.run(client.close())
+
+
+def test_raises_when_content_entry_missing_sampling_logprob(monkeypatch):
+    client = FireworksV1CompletionsClient(
+        model_id="test-model",
+        tokenizer_name_or_path="Qwen/Qwen3-0.6B",
+    )
+    monkeypatch.setattr(client, "decode_token_ids", lambda token_ids: "text")
+    _install_fake_completion(
+        client,
+        monkeypatch,
+        {
+            "choices": [
+                {
+                    "finish_reason": "stop",
+                    "logprobs": {
+                        "content": [
+                            {"token_id": 1, "sampling_logprob": -0.1},
+                            {"token_id": 2, "logprob": -0.2},
+                        ]
+                    },
+                }
+            ],
+        },
+    )
+    with pytest.raises(RuntimeError, match="missing .*sampling_logprob"):
+        asyncio.run(client.create_completion_from_prompt_ids(prompt_token_ids=[1]))
+    asyncio.run(client.close())
+
+
+def test_raises_when_no_content_logprobs(monkeypatch):
+    client = FireworksV1CompletionsClient(
+        model_id="test-model",
+        tokenizer_name_or_path="Qwen/Qwen3-0.6B",
+    )
+    _install_fake_completion(
+        client,
+        monkeypatch,
+        {"choices": [{"text": "hello world", "finish_reason": "stop"}]},
+    )
+    with pytest.raises(RuntimeError, match="no content\\[\\] logprobs entries"):
+        asyncio.run(client.create_completion_from_prompt_ids(prompt_token_ids=[1, 2]))
+    asyncio.run(client.close())
+
+
+def test_ignores_legacy_token_logprobs_shape(monkeypatch):
+    """The legacy token_logprobs shape has no content[]; the client must not use it."""
+    client = FireworksV1CompletionsClient(
+        model_id="test-model",
+        tokenizer_name_or_path="Qwen/Qwen3-0.6B",
+    )
+    _install_fake_completion(
+        client,
+        monkeypatch,
+        {
+            "choices": [
+                {
+                    "finish_reason": "stop",
+                    "token_ids": [1, 2, 3],
+                    "logprobs": {
+                        "token_ids": [1, 2, 3],
+                        "token_logprobs": [-0.1, -0.2, -0.3],
+                    },
+                }
+            ],
+        },
+    )
+    with pytest.raises(RuntimeError, match="no content\\[\\] logprobs entries"):
+        asyncio.run(client.create_completion_from_prompt_ids(prompt_token_ids=[1]))
+    asyncio.run(client.close())
+
+
+def test_raises_when_content_entry_missing_token_id(monkeypatch):
+    client = FireworksV1CompletionsClient(
+        model_id="test-model",
+        tokenizer_name_or_path="Qwen/Qwen3-0.6B",
+    )
+    monkeypatch.setattr(client, "decode_token_ids", lambda token_ids: "text")
+    _install_fake_completion(
+        client,
+        monkeypatch,
+        {
+            "choices": [
+                {
+                    "finish_reason": "stop",
+                    "logprobs": {
+                        "content": [
+                            {"token_id": 1, "sampling_logprob": -0.1},
+                            {"sampling_logprob": -0.2},
+                        ]
+                    },
+                }
+            ],
+        },
+    )
+    with pytest.raises(RuntimeError, match="missing token_id"):
+        asyncio.run(client.create_completion_from_prompt_ids(prompt_token_ids=[1]))
+    asyncio.run(client.close())
+
+
 def test_thinking_kwargs_respects_enable_thinking():
     client_none = FireworksV1CompletionsClient(
         model_id="test", tokenizer_name_or_path="Qwen/Qwen3-0.6B",
