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

[Dranoxgithub]

Summary

FireworksV1CompletionsClient.create_completion_from_prompt_ids is the sampling path for RL rollouts against Fireworks /v1/completions. Downstream, the trainer aligns each per-token inference logprob to its completion token id position-by-position, so completion_ids[i] and the per-token logprob at i must be the same length and the same tokens.

This client always requests boolean logprobs=True, so every response uses the new content[] logprobs shape (choices[0].logprobs.content[], where each entry carries token_id, token, bytes, logprob, and sampling_logprob). The legacy integer logprobs=1 shape (token_logprobs + logprobs.token_ids) is not requested and is intentionally not supported here. This PR makes the client read everything from that new-mode content[] array.

Bug 1 — alignment / KLD corruption (the actual bug)

Per-token ids and logprobs were read from two independent places:

• ids: only from the top-level choices[].token_ids, which is populated only when return_token_ids=True — a flag this client never sent. When absent, ids were silently re-derived via tokenizer.encode(decode(text), add_special_tokens=False).
• logprobs: from content[].logprob, always full length.

Retokenization does not reproduce the model's exact sequence — the trailing end-of-turn/EOS token decodes to empty text and is not re-added — so completion_ids = N while logprobs = N+1. The +1 shift misaligns every per-token logprob against the wrong token, and per-token KLD becomes garbage.

Observed on a GLM-family LoRA RFT deployment (MTP speculative decoding): completion_ids was one shorter than the logprobs on all 120 rows (diff=+1) and inference_kld climbed to ~60; aligning ids correctly made diff=0 and inference_kld flat ~0.028.

Fix: the content[] array already carries token_id per entry, so we read ids and logprobs from the same array, entry-by-entry. They are the same length and aligned by construction — no return_token_ids, no re-encode fallback.

Bug 2 — sampling precision (more of a feature)

In the content[] shape each entry exposes two numbers: sampling_logprob (the exact, full-precision value the sampler actually drew with) and logprob (a rounded/verification value). The client read the rounded logprob. We now prefer sampling_logprob (falling back to logprob only if absent), which is the correct inference-side value for KLD.

Changes

• Read per-token completion_ids from content[].token_id and logprobs from content[].sampling_logprob in a single pass, so they are aligned by construction.
• Drop the tokenizer.encode(...) re-encode fallback and the legacy id/logprob sources (token_logprobs, top-level token_ids, raw_output.completion_token_ids).
• No longer send return_token_ids; rely on the content[] shape that boolean logprobs=True already returns.
• Fail loud when content[] is absent (e.g. logprobs disabled or legacy shape returned) or when a content[] entry lacks token_id, instead of silently returning corrupted data.

Test plan

• Unit: content[] with token_id + sampling_logprob -> completion_ids from token_id, logprobs from sampling_logprob (falls back to logprob), lengths equal, tokenizer never used to re-encode.
• Unit: no content[] -> raises (no silent re-encode).
• Unit: legacy token_logprobs shape (no content[]) -> raises (legacy not supported).
• Unit: content[] entry missing token_id -> raises.
• pytest tests/test_fireworks_v1_completions_client.py passes (12 tests).

---

Note

Medium Risk
Changes the RL sampling contract on a critical path (token/logprob alignment); correct for intended Fireworks responses but will error on legacy or incomplete logprob payloads that previously returned silently wrong data.

Overview
Fixes RL rollout corruption in FireworksV1CompletionsClient.create_completion_from_prompt_ids by building completion_ids and completion_logprobs in one pass from choices[0].logprobs.content[] (token_id + sampling_logprob), so lengths stay matched for per-token inference KLD.

Removes top-level token_ids, legacy token_logprobs / rounded logprob paths, and tokenizer re-encode when ids were missing (that path dropped EOS and shifted logprobs). Requires the boolean logprobs=True content[] shape and raises if content[], token_id, or sampling_logprob is missing instead of returning misaligned data.

Adds unit tests with a fake completions client covering the happy path, missing fields, absent content[], and legacy logprob shapes.

^{Reviewed by Cursor Bugbot for commit 03facac. Bugbot is set up for automated code reviews on this repo. Configure here.}

[Dranoxgithub]

Superseded by #458 (same commits, moved from a fork branch to an in-repo branch so CI runs with repo secrets).