Repository navigation
perf: return completion logprobs as a float32 array - #201
Merged
Merged
Conversation
prime-rl's vLLM server sends sampled logprobs as a base64 float32 buffer. generate() decoded it, widened it to float64 and built a Python list, which verifiers then narrowed back into a float32 array. generate() now returns the decoded buffer as a float32 array without a copy. The JSON and compact forms return float32 arrays too: vLLM samples in float32, so the narrowing is lossless. Callers that compared completion_logprobs to a list must compare its tolist() or use pytest.approx. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This was referenced Oct 23, 2026
ApprovabilityVerdict: Needs Human Approval Macroscope's review found this PR requires human approval — The PR improves allocation efficiency by returning float32 NumPy arrays, but changes the existing public Notes:
You can add or adjust custom eligibility rules. Learn more. |
samsja
approved these changes
Oct 24, 2026
hallerite
approved these changes
Oct 24, 2026
mikasenghaas
added a commit
to PrimeIntellect-ai/verifiers
that referenced
this pull request
Oct 24, 2026
renderers 0.1.13.dev17 returns completion logprobs as a float32 array (PrimeIntellect-ai/renderers#201). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
mikasenghaas
added a commit
to PrimeIntellect-ai/verifiers
that referenced
this pull request
Oct 24, 2026
## Summary A `MessageNode`'s per-token streams were Python lists, with one object per token. These streams are `token_ids`, `mask`, `is_content` and `logprobs`, plus the consumer-stamped `advantages`, `reference_logprobs`, `trainer_logprobs`, `entropies` and `loss_weights`. The cyclic GC walks each object on every full collection, and msgpack packs them one by one. An RL orchestrator that holds thousands of long agentic episodes spent minutes per full collection, with the GIL held. This PR makes every stream a typed numpy array, from the renderer client to the wire. - **The field types carry their dtype.** `IntStream` is an int32 array, `BoolStream` a bool array and `FloatStream` a float32 array. Validation accepts a list, an array or a raw-bytes `__nd__` dict, the encoding `routed_experts` already uses. Field defaults are typed empty arrays. - The types are plain `np.ndarray` with a JSON schema attached. A `list[int] | np.ndarray` union makes pydantic try the list branch element by element: about 200 ms per 200k-token node. - Serialization stays a `@field_serializer`, because `serialize_as_any` (the wire dump) skips serializers in `Annotated` types. - **`TurnTokens` streams are typed arrays too:** `prompt_ids`, `completion_ids`, `completion_logprobs`, `renderer_prompt_ids` and `is_content` use the stream types. - The config tree imported `verifiers.v1.types` for only `ID`, `DialectName` and `SamplingConfig`, and must load without numpy. `ID` and `DialectName` now live in `verifiers.v1.configs.harness`, and `SamplingConfig` (with its `Sampling` alias) in the new `verifiers.v1.configs.sampling`. `verifiers.v1.types` no longer has them; every importer and the `vf.*` lazy table use the new modules. - `verifiers.v1.types` now imports numpy and holds the stream types and the ndarray codec, which moved from `graph.py`. The config tree still loads without numpy and without `verifiers.v1.types`. - **The env server builds nodes from arrays.** - The renderer client converts the turn's token lists to arrays once (`TurnTokens`). - The turn's commit copies each node's slice. A view would keep the turn's full-context prompt alive. - `renderer_token_ids` is set only after multimodal expansion. Before, every node held a second copy of `token_ids` there. - The multimodal prompt projection is vectorised. Token-span comparisons use `np.array_equal`. - Renderer bridging still takes lists, so `previous_renderer_token_ids` converts the context once per turn. - **The env server always streams node streams as raw bytes** (`__nd__` dicts), not msgpack arrays. A server that predates this PR still sends lists, which validate the same way. - **`encode_ndarray` / `decode_ndarray` are public**, so consumers of raw deltas (prime-rl's live trace files) can decode them. - **JSON records are unchanged.** They are still lists, with logprob, advantage and entropy streams rounded to `float_decimals`. Loss weights stay exact. Python dumps without the wire context also still give lists. - **`Branch` returns numpy arrays:** `token_ids` int32, `sampled_mask` bool, the spread streams and `loss_weights` float32, `mm_token_type_ids` int32. `spread` is vectorised. Part of a set with PrimeIntellect-ai/renderers#201 (merged: logprobs decode to a float32 array without a list round trip) and PrimeIntellect-ai/prime-rl#4226 (the consumer). This PR requires `renderers>=0.1.13.dev17`, the first dev release with #201. ## Breaking - **Env server and client must both include this PR.** An older client cannot validate the raw-bytes node streams, so its rollouts fail validation. A newer client with an older server still works. - `Branch.token_ids`, `sampled_mask`, `logprobs`, `advantages`, `reference_logprobs`, `trainer_logprobs`, `entropies`, `loss_weights(name)` and `mm_token_type_ids` return numpy arrays, not lists. `MessageNode`'s stream fields hold arrays too. Callers that compare with `==`, rely on truthiness or concatenate with `+` must use `.tolist()`, `len()`, `.any()` or `np.concatenate`. - `MessageNode.renderer_token_ids` is `None` unless multimodal expansion changed the prompt. Read `logical_ids` to get the renderer tokens in every case. - `TurnTokens.prompt_ids`, `completion_ids`, `completion_logprobs`, `renderer_prompt_ids` and `is_content` are typed numpy arrays. Lists still validate, into arrays. ## Tests - `tests/unit/test_token_arrays.py`: - typed arrays after validation, and typed empty defaults; - `__nd__` wire round trip, including `loss_weights`, and lists without the context; - JSON records stay lists and round correctly; - committed env-server nodes hold owned typed arrays and no duplicate renderer ids; - over 300 randomised branches mixing list and array nodes, vectorised `spread` matches the previous list implementation exactly. - Graph and semantic tests compare streams via `list(...)` or `.tolist()`. - `tests/unit` passes locally, except three tests that fail the same way on this PR's base: two `runtimes/test_private_network.py` VM tests and `test_api.py::test_config_tree_imports_without_the_runtime`. The test host's venv has a `.pth` hook that loads `nvidia_cutlass_dsl` into every process. Run by hand, the config tree loads neither numpy nor `verifiers.v1.types`. - `ty check verifiers` passes. ## Measured A recorded batch of 2048 agentic episodes (393k nodes, 194M tokens, branches up to 262k tokens) was replayed through the production env-server and wire code, on `main` and on this PR (with PrimeIntellect-ai/prime-rl#4226 as the client): | | `main` (lists) | This PR (arrays) | |---|---|---| | Env server: memory for the held traces | 22.6 GB | 5.7 GB | | Env server: stream encode (dump + pack) | 26.2 s | 9.3 s | | Wire size | 3.72 GB | 2.99 GB | | Env server: full GC with the traces held | 13.5 s | 6.8 s | | Client: decode (unpack + assemble + validate) | 19.4 s | 6.9 s | | Client: memory for the held episodes | 21.2 GB | 6.0 GB | The replay is single-process and CPU-only. GC is disabled during timed phases and timed separately. Memory is RSS growth. 🤖 Generated with [Claude Code](https://claude.com/claude-code) <!-- CURSOR_SUMMARY --> --- > [!NOTE] > **Medium Risk** > Breaking public API (`Branch`/`MessageNode` streams are arrays) and env-server node wire format requires a matching client; core trace/graph/training paths are affected but covered by targeted tests. > > **Overview** > **Per-token trace data is stored as typed NumPy arrays end-to-end**, replacing Python lists on `MessageNode`, `TurnTokens`, and `Branch` aggregations (`token_ids`, masks, logprobs, advantages, loss weights, etc.). Validation accepts lists, arrays, or `__nd__` wire dicts; JSON records still serialize as lists. > > **Config types move out of `verifiers.v1.types`** so the config tree can load without NumPy: `SamplingConfig` / `Sampling` → `verifiers.v1.configs.sampling`, `ID` and `DialectName` → `verifiers.v1.configs.harness`, with imports and lazy `vf.*` exports updated accordingly. > > **Turn commit and training client paths** convert renderer output to arrays once, copy per-node slices (no views), vectorize multimodal attribution, and only set `renderer_token_ids` when the prompt is multimodal-expanded. **Env-server deltas** stream new nodes with `token_arrays` so token streams use raw-bytes encoding via public `encode_ndarray` / `decode_ndarray`. > > Bumps **`renderers[multimodal]`** to `>=0.1.13.dev17` and adds **`tests/unit/test_token_arrays.py`** (wire round-trip, spread parity, committed-node ownership). Existing graph/semantic tests compare streams via `.tolist()` / `list(...)`. > > <sup>Reviewed by [Cursor Bugbot](https://cursor.com/bugbot) for commit 058db32. Bugbot is set up for automated code reviews on this repo. Configure [here](https://www.cursor.com/dashboard/bugbot).</sup> <!-- /CURSOR_SUMMARY --> <!-- Macroscope's pull request summary starts here --> <!-- Macroscope will only edit the content between these invisible markers, and the markers themselves will not be visible in the GitHub rendered markdown. --> <!-- If you delete either of the start / end markers from your PR's description, Macroscope will append its summary at the bottom of the description. --> > [!NOTE] > ### Store per-token node streams as typed NumPy arrays in `MessageNode` > - Adds typed stream aliases (`IntStream`, `BoolStream`, `FloatStream`) and a `_token_stream` validator in [graph.py](https://github.com/PrimeIntellect-ai/verifiers/pull/2961/files#diff-64c9d635eb215902a512c9e8ab316fd9b9967c36db03d4049aacc812cdfbd740) so token IDs, masks, logprobs, advantages, entropies, and loss weights on `MessageNode` become int32/bool/float32 arrays. > - Rewrites graph commit, prefix reconciliation, and multimodal attribution helpers in [graph.py](https://github.com/PrimeIntellect-ai/verifiers/pull/2961/files#diff-64c9d635eb215902a512c9e8ab316fd9b9967c36db03d4049aacc812cdfbd740) to use array operations and a new `same_tokens` equality helper instead of list comparisons. > - `Branch` properties and `spread` in [trace.py](https://github.com/PrimeIntellect-ai/verifiers/pull/2961/files#diff-834820c3aa80d87ee1c1c57fb54c74f09c66ac727036d964e6005c8613c00b5e) now return concatenated typed arrays instead of Python lists; JSON dumps stay list-based, and an opt-in token-array wire format uses public `encode_ndarray`/`decode_ndarray` (also used by delta streaming in [delta.py](https://github.com/PrimeIntellect-ai/verifiers/pull/2961/files#diff-106a565d15f6f5c86f7357854d512768b3f8a4775026a0eae32c69b828f5da31)). > - [train.py](https://github.com/PrimeIntellect-ai/verifiers/pull/2961/files#diff-c9bcd6c691890b1de6942fe012d2e74678460b22bf335721a8bba5e939d139c3) converts generated streams to arrays before building `TurnTokens`, and [types.py](https://github.com/PrimeIntellect-ai/verifiers/pull/2961/files#diff-88a4e4ee0c58ac5dc7344161a011374082a3767eda2ad32af97c942dca982112) loosens those annotations to `Any` to avoid a NumPy import. > - Behavioral Change: `Branch.token_ids`, `sampled_mask`, `logprobs`, `advantages`, `mm_token_type_ids` and related properties return NumPy arrays instead of lists; `Branch.spread` returns a float32 array. Out-of-tree code comparing these with `==` or expecting lists will behave differently. > > <!-- Macroscope's changelog starts here --> > #### Changes since #2961 opened > > - Updated `renderers[multimodal]` dependency version constraint [0d44cb4] > - Introduced typed numpy array stream infrastructure for token data in `verifiers.v1.types` module, including `encode_ndarray`, `decode_ndarray`, `as_list` utilities, `_token_stream` factory with pydantic `BeforeValidator` and `WithJsonSchema`, typed stream aliases `IntStream`, `BoolStream`, `FloatStream`, and `empty_stream` default factory [6a17126] > - Converted `TurnTokens` model fields to use typed numpy array streams with automatic serialization in `verifiers.v1.types` module [6a17126] > - Refactored `MessageNode` class in `verifiers.v1.graph` module to use centralized token stream utilities from `verifiers.v1.types` [6a17126] > - Created `verifiers.v1.configs.types` module defining `DialectName`, `SamplingConfig`, `Sampling` alias, and `ID` type with validation logic [6a17126] > - Updated imports across `verifiers.v1.configs.agent`, `verifiers.v1.configs.cli.eval`, `verifiers.v1.configs.env`, `verifiers.v1.configs.harness`, `verifiers.v1.configs.judge`, `verifiers.v1.configs.taskset`, and `verifiers.v1.utils.generic` modules to source types from `verifiers.v1.configs.types` [6a17126] > - Renamed `verifiers.v1.configs.types` module to `verifiers.v1.configs.sampling` and redistributed type definitions by moving `ID` and `DialectName` type aliases to `verifiers.v1.configs.harness` [61c8062] > - Updated import paths for `Sampling` and `SamplingConfig` from `verifiers.v1.types` to `verifiers.v1.configs.sampling` across all consuming modules [61c8062] > - Updated import paths for `ID` and `DialectName` from `verifiers.v1.types` to `verifiers.v1.configs.harness` across all consuming modules [61c8062] > - Removed re-exports of `ID`, `DialectName`, `Sampling`, and `SamplingConfig` from `verifiers.v1.types` module and updated lazy-loading configuration in `verifiers.v1.__init__` [61c8062] > - Reordered and deduplicated imports across all `verifiers.v1` modules, test files, and related configuration modules [83e11cd] > - Reflowed and rewrapped the docstring text for the `_token_stream` function in the `verifiers/v1` module [058db32] > <!-- Macroscope's changelog ends here --> > > <!-- Macroscope's review summary starts here --> > > <details> > <summary>📊 <a href="https://app.macroscope.com">Macroscope</a> summarized 49ae611. 3 files reviewed, 1 issue evaluated, 0 issues filtered, 1 comment posted</summary> > > ### 🗂️ Filtered Issues > > </details><!-- Macroscope's review summary ends here --> > <!-- Macroscope's pull request summary ends here --> --------- Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
mikasenghaas
added a commit
to PrimeIntellect-ai/prime-rl
that referenced
this pull request
Oct 24, 2026
…#4226) ## Summary The orchestrator held every in-flight and buffered episode, and every pending `TrainingSample`, as Python lists with one entry per token. That made memory about 3.5x larger and put hundreds of millions of references in every full GC. It also made the per-token Python loops expensive: advantage assignment, echo weighting and zero-advantage pruning. This PR keeps per-token streams as typed arrays (int32 ids, bool masks, float32 values) from the vLLM response to the trainer. - **vLLM to the env server:** `deps/renderers` (PrimeIntellect-ai/renderers#201) decodes sampled logprobs to a float32 array, without a float64 list round trip. `deps/verifiers` (PrimeIntellect-ai/verifiers#2961) builds env-server nodes as owned arrays. - **Env server to the orchestrator:** env servers always stream node streams as raw-bytes arrays. In-flight assemblies hold the bytes, and validated episodes hold arrays that view them. - **Algorithms** work on arrays, vectorised: - `assign_advantages` and `assign_reference_logprobs`; - `scalar_advantage` and `is_trainable`; - `stamp_loss_routing`; - grpo `length_weighted_baseline` and `prompt` loss aggregation; - echo observation weights; - the advantage-range curriculum gate; - `_prune_zero_advantages`, which also ran per token on the event loop. - **Samples and micro batches:** `TrainingSample` and `MicroBatch` streams are numpy arrays. - Packing keeps `main`'s shape: each sample's arrays are appended and concatenated once per bin, and padding uses `np.pad`. - On the wire, each `MicroBatch` stream is a raw buffer of a fixed dtype (`MICRO_BATCH_DTYPES`, shipped by the batch encoder's `enc_hook`). The per-micro-batch alignment assert also checks each stream's dtype. - **Trainer:** each stream becomes a tensor through one numpy copy that also widens the dtype (`np.frombuffer` + `astype`), not `torch.tensor(list)`. - **`MicroBatch.env_names`** holds one name per sequence, not one per token. A new `padding` field holds the padding count. The per-env entropy and mismatch-KL metrics expand the names over `sequence_lengths` and select tokens with vectorised masks. The annotation export strips padding by count. - **Live trace files** decode the raw-bytes streams with verifiers' public `decode_ndarray` before writing JSON. - **opd/opsd** send token ids to the scorer as lists. Depends on PrimeIntellect-ai/verifiers#2961. `deps/verifiers` points at that PR's commit, so this stays a draft until it merges and the pin moves to `main`. PrimeIntellect-ai/renderers#201 has merged; `deps/renderers` points at `main`. Companion to #4225, which stops automatic full GCs. This PR makes each full GC cheaper. ## Breaking - **The `MicroBatch` wire format changes.** Per-token streams are raw buffers, not msgpack lists, and `env_names` is per sequence. The orchestrator and the trainer must run the same prime-rl commit. - **Env servers and the orchestrator must both include verifiers#2961.** Node streams always ride the env-server wire as raw bytes, which an older client cannot validate. ## Verification - **Replay at scale.** A recorded batch of 2048 agentic episodes (147 groups, 393k nodes, 194M tokens, branches up to 262k tokens) went through the production env-server stream, decode, GRPO scoring, sample building, `prepare_batch` (`seq_len=262144`, 8 ranks), the batch encoder, the trainer decoder and `_micro_batch_to_tensor`. Both trees produced the same 1802 samples, 167M sample tokens and 648 micro batches. - **Disk and upload formats on the same 2048 episodes,** this branch against `main`: - Byte-identical for every episode: the file monitor's trace records (default `float_decimals`), live trace files, the Prime Traces payload, the samples projection and the training-viewer parquet rows. - Full-precision records (`float_decimals=None`) and annotation advantages print some floats differently, as the float32 value's shortest form. Every differing number is the same float32 value. - **Orchestrator equivalence.** On an earlier recorded batch (115 samples, 7.8M tokens), every packed `MicroBatch` field and every scalar advantage matched `main` (float32 compare). That held for mean and `prompt` loss aggregation, with and without truncation. Every trainer tensor also matched `main` on 104 micro batches with padding and dummy batches. - **Unit tests:** - New `tests/unit/orchestrator/test_episode_records.py` sends a wire-decoded episode through the file monitor's episode, live and annotation writers and the Prime upload projections. It checks that each output holds JSON lists. - `tests/unit/orchestrator`, `transports`, `eval`, `utils` and the dashboard tests pass locally. One exception: `test_prefill_logprobs_uses_inference_generate` fails at import, because the local vLLM lacks `GenerateTokensResponse`. - **Not yet run:** a GPU training run. ## Measured The 2048-episode replay. It is single-process and CPU-only; GC is disabled during timed phases and timed separately; memory is RSS growth. Trainer figures add up all 8 ranks. | | `main` | This PR | |---|---|---| | Decode (unpack + assemble + validate) | 19.4 s | 6.9 s | | Memory for the held episodes | 21.2 GB | 6.0 GB | | Score + build samples | 23.0 s | 7.8 s | | Memory for the pending samples | 6.8 GB | 3.1 GB | | Full GC, episodes + samples held | 14.8 s | 7.6 s | | Pack (`prepare_batch`) | 9.5 s | 3.6 s | | Batch wire size | 6.79 GB | 3.51 GB | | Batch encode | 5.3 s | 1.1 s | | Trainer decode | 20.4 s | 0.003 s | | Trainer tensor conversion | 25.4 s | 0.20 s | | Peak orchestrator RSS | 85.3 GB | 17.5 GB | The remaining full-GC time is pydantic's per-model objects (about 8M for this batch), not token data. #4225's scheduled-only full collection keeps it off the hot path. 🤖 Generated with [Claude Code](https://claude.com/claude-code) <!-- CURSOR_SUMMARY --> --- > [!NOTE] > **High Risk** > Breaking batch wire format and env-server ndarray streams require lockstep orchestrator/trainer/verifiers upgrades; training correctness depends on dtype alignment and per-sequence `env_names`/`padding` semantics. > > **Overview** > Moves **per-token training data** from Python lists to **typed numpy arrays** end-to-end (orchestrator samples → packed `MicroBatch` → trainer tensors), cutting memory, GC cost, and Python-loop overhead on advantage assignment, echo/GRPO weighting, pruning, and packing. > > **Wire format (breaking):** `MicroBatch` per-token fields are **raw byte buffers** (`MICRO_BATCH_DTYPES`) via `encode_array`, not msgpack lists. **`env_names`** is **one entry per sequence** (not per token); new **`padding`** tracks trailing pad tokens for annotation export and metrics. Trainer builds tensors with **`np.frombuffer` + `astype`** instead of `torch.tensor(list)`. > > **Algorithms & sink** use vectorized mask/weight ops; OPD/OPSD pass **`token_ids.tolist()`** to scorers. **Live trace files** decode `__nd__` payloads with **`decode_ndarray`** before JSON. Disk/upload paths still emit plain JSON lists (new `test_episode_records.py`). > > <sup>Reviewed by [Cursor Bugbot](https://cursor.com/bugbot) for commit 9652f05. Bugbot is set up for automated code reviews on this repo. Configure [here](https://www.cursor.com/dashboard/bugbot).</sup> <!-- /CURSOR_SUMMARY --> --------- Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
generate()returnscompletion_logprobsas a float32 numpy array, not a list of Python floats.generate()widened it to float64 and built a list. verifiers then narrowed that list back into a float32 array.Part of a set with PrimeIntellect-ai/verifiers#2961 and PrimeIntellect-ai/prime-rl#4226. verifiers accepts either a list or an array, so the PRs can merge in any order.
Breaking
generate()["completion_logprobs"]is anp.ndarray(float32), not alist[float]. Code that compares it to a list with==, or tests its truthiness, must use.tolist(),len()orpytest.approx. Indexing,len()and iteration work as before.Tests
tests/test_client.pyandtests/replay.pycompare the array through.tolist()orpytest.approx, and check the float32 dtype.tests/test_client.pyandtests/test_incremental.pypass locally.test_multimodal.pyInkling processor-parity cases. They render prompts and don't callgenerate().🤖 Generated with Claude Code
Note
Medium Risk
Breaking API change for
generate()["completion_logprobs"](list → ndarray); downstream code that assumed list equality or truthiness may fail until updated.Overview
generate()now returnscompletion_logprobsas a float32numpyarray instead of alist[float], avoiding extra widen-to-float64 / list materialization on the hot path (especially for prime-rl’s base64 packed float32 buffer, which can decode without a copy).The logprob parsers (
_parse_packed_completion_logprobs,_parse_compact_logprobs,_parse_completion_logprobs) and compact sampling output now build or returnnp.ndarraydtypefloat32;generate()’s docstring documents the new shape. Breaking: callers that compared the field to a plain list with==need.tolist(),pytest.approx, or element-wise checks; indexing,len(), and iteration still work.Tests in
tests/test_client.pyandtests/replay.pyassert dtype and compare via.tolist()/pytest.approx.Reviewed by Cursor Bugbot for commit faf9e2e. Bugbot is set up for automated code reviews on this repo. Configure here.
Note
Return completion logprobs from
renderersas float32 NumPy arrays_parse_packed_completion_logprobs,_parse_compact_logprobs, and_parse_completion_logprobsin renderers/client.py now return float32 NumPy arrays instead of Python float listspytest.approxor list conversion, and adds dtype checkscompletion_logprobsis now a float32 NumPy array with one entry per completion token, so consumers expecting Python floats lose precision and list semanticsMacroscope summarized faf9e2e.