fix: request exact token ids and fail loud in FireworksV1CompletionsClient

[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, which requires completion_ids[i] and token_logprobs[i] to be the same length and same tokens. Two combining defects corrupted that alignment on every rollout:

• Defect 1 — flags never sent. The request only set logprobs; it never set return_token_ids, so choices[].token_ids came back empty by design.
• Defect 2 — silent lossy re-encode. With exact ids absent, the client re-derived them via tokenizer.encode(decoded_text, add_special_tokens=False). Retokenization drops the trailing end-of-turn/EOS token (which decodes to empty text), so completion_ids = N while token_logprobs = N+1. The +1 shift misaligns every logprob against the wrong token.

Observed impact on a GLM-family LoRA RFT deployment (MTP speculative decoding): without return_token_ids, completion_ids was one shorter than token_logprobs on all 120 rows (diff=+1) and inference_kld climbed to ~60; requesting exact ids made diff=0 and inference_kld flat ~0.028.

Changes

• Default return_token_ids=True in the request payload via setdefault, so it stays overridable through request_params.
• Remove the tokenizer.encode(...) re-encode fallback. If exact ids are absent after reading choices[].token_ids / raw_output.completion_token_ids, raise RuntimeError — a missing/contract failure must not silently reach training.
• Assert len(completion_token_ids) == len(completion_logprobs) (when logprobs are present) so any residual drift is caught at the boundary rather than as a downstream KLD anomaly.

This mirrors the upstream Fireworks monorepo fix (fw-ai/fireworks#33729), bringing the base eval_protocol client in line so unmodified callers (e.g. examples/multihop_qa/multihop_qa_rollout.py) are no longer silently corrupted.

Test plan

• Unit: response with choices[0].token_ids present → completion_ids equals it (tokenizer not used to re-encode) and len(completion_ids) == len(completion_logprobs).
• Unit: return_token_ids=True sent by default; overridable via request_params.
• Unit: no token_ids and no raw_output.completion_token_ids → raises (no silent re-encode).
• Unit: id/logprob length mismatch → raises.
• pytest tests/test_fireworks_v1_completions_client.py passes (13 tests).

Made with Cursor

[Dranoxgithub]

Background / how the root cause was pinned down (adding here since it wasn't captured in the PR):

The two runs that matter — same GLM-5.1 multihop-QA RL config, 256 rows, identical code except the return_token_ids request flag:

run	return_token_ids	token_logprobs − completion_ids	inference_kld
28745665391	not sent	+1 on all 120 rows	spiked to ~60
28762217001	True	0 on all 120 rows	flat ~0.028

That A/B is the whole proof: flipping only the flag moved every row from diff=+1 to diff=0 and collapsed the KLD spikes.

How we got there (the misleading trail):

1. Spikes were rare (~1–2% of rows) and clustered on mojibake / mixed-script byte tokens (©´¶, Korean/Cyrillic fragments), which first looked like genuine sampler-vs-trainer divergence on degenerate generations — i.e. "bad luck," not a bug.
2. A per-token loss dump showed old ≈ pi ≠ inf (both trainer forwards agreed; only the stored sampling logprob was the outlier), pointing at the stored logprobs, not the model.
3. An earlier alignment audit reported first_mismatch=None and was dismissed as harmless — but it compared full_tokens vs completion_ids (both from the same N-length array) and never checked len(completion_ids) vs len(token_logprobs). That blind spot hid the +1.
4. Logging the two lengths per request revealed token_logprobs = N+1 while completion_ids = N, uniformly.

Live confirmation of the mechanism (Qwen 3.6 deployment, direct /v1/completions, logprobs=1, natural stop): with return_token_ids=true, choice.token_ids is present and its last id is <|endoftext|> (the EOS), length == token_logprobs. With it false, choice.token_ids is absent, so the client re-encodes the decoded text; re-encoding "…Four" with add_special_tokens=False yields 6 ids vs token_logprobs 7 — exactly the dropped trailing EOS → the +1. So the failure mode is deployment-dependent: Qwen omits token_ids (→ lossy re-encode), the GLM rollout deployment returns a short array; both are fixed by requesting exact ids and refusing to re-encode.

Monorepo counterpart (e2e vendored client + fail-loud guard): fw-ai/fireworks#33733.