[mikasenghaas]

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 (perf: return completion logprobs as a float32 array renderers#201) decodes sampled logprobs to a float32 array, without a float64 list round trip. deps/verifiers (perf(v1): per-token node streams as numpy arrays 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

---

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).

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