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 read sampling_logprob (the correct inference-side value for KLD) and fail loud if it is missing.

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 or sampling_logprob, instead of silently returning corrupted data.

Mirrors the intent of the upstream Fireworks monorepo fix (fw-ai/fireworks#33729): callers that use the base eval_protocol client unmodified (e.g. examples/multihop_qa/multihop_qa_rollout.py) are no longer silently corrupted.

Test plan

• Unit: content[] with token_id + sampling_logprob -> completion_ids from token_id, logprobs from sampling_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.
• Unit: content[] entry missing sampling_logprob -> raises.
• pytest tests/test_fireworks_v1_completions_client.py passes (13 tests).

---

Supersedes #457 (same commits, moved from a fork branch to an in-repo branch so CI runs with secrets). Previously approved by @websterbei.

Made with Cursor