Data streaming by samsja · Pull Request #4 · PrimeIntellect-ai/prime-rl

samsja · 2025-02-21T02:03:03Z

No description provided.

Signed-off-by: Sami Jaghouar <sami.jaghouar@gmail.com>

apaz-cli · 2025-02-21T05:39:51Z

+    if world_info.rank == 0:
+        files = os.listdir(path)
+
+    dist.broadcast_object_list([files], src=0)


What about dist.scatter_object_list() instead?

huh I did not know it existed 🐒 great idea

hmm I tried to use it but it require to know the size of the list before hand.
Which I dont

apaz-cli · 2025-02-21T05:45:06Z

+    def __iter__(self):
+        while True:
+            for file in self._shared_files:
+                table = pq.ParquetFile(file).read()


This holds the entire file in memory. Probably want to stream it from disk? Unless the assumption is we create a new dataloader every iteration or something, with the new files from S3.

The file are supposed to be small should not be a problem

apaz-cli · 2025-02-21T05:52:29Z

-        for _ in range(self.step):
-            next(itera)
+            self._logger.info("Waiting for new files")
+            time.sleep(0.5)


I assume the sleep is a placeholder for the actual condition to be added later.

Ex: A notification from another thread that all the files have been received.

apaz-cli · 2025-02-21T06:00:17Z

+            yield {"input_ids": input_ids, "advantages": advantages}
+
+
+class ParquetDataset(IterableDataset):


I presume this is also supposed to be an UpdateableDataset?

with Protocol we don't need to inherit. It a bit like rust trait

apaz-cli

Looks pretty good, agree on the API. Appending parquet files to the dataloader is a good idea.

Signed-off-by: Sami Jaghouar <sami.jaghouar@gmail.com>

* update verifier * update mem usage --------- Co-authored-by: Mika Senghaas <mail@mikasenghaas.de>

The segmented dispatch in train.py used to skip the entropy append entirely for compaction micro-batches, because segmented_forward's per-segment backward path never surfaced a full-sequence entropy tensor. As a result, W&B entropy/mean reported only the non-compaction partition of each batch — a misleading subset metric during compaction training. This patch computes entropy inside the per-segment loss closure on the already temperature-scaled, boundary-trimmed seg_logits_effective the closure receives, masks with seg_mask (loss_mask sliced to the segment's owned target range), and accumulates the per-segment 1-D results in accumulated_entropy_masked. After segmented_forward returns, the list is concatenated into a single 1-D tensor and appended to tensors["entropy"] — matching the shape convention the standard path's out["entropy"][loss_mask] produces, so downstream compute_stats aggregation stays uniform. Correctness by construction: each segment's owned logit range is disjoint from its neighbors (seg_end-1 trim on non-final, boundary- overlap re-feed on the next segment's seg_start), so concatenating per-segment masked entropies across segments counts every completion token exactly once. The closure's effective_logit_end trim on the final segment drops the unresolvable seq_len-1 logit, matching the standard path's shift_tensor_right behavior. Degenerate case: if every segment short-circuits to the zero-loss tail branch (all owned ranges trimmed away), the list is empty and we append a shape-(0,) bfloat16 tensor so Tensors.compute_stats's torch.cat across micro-batches doesn't fail on a dtype mismatch with the standard path's bf16 entries. Debug log guard relaxed from `if not use_segmented` to `if tensors["entropy"][-1].numel() > 0` so the per-micro-step entropy line prints in both branches. Cost: one compute_entropy call per segment, running under torch.no_grad (compute_entropy's own decorator). Negligible next to the model forward. Zero extra GPU memory beyond the momentary entropy tensor which is immediately detached and moved to CPU. Resolves deferred item D2 from plans/phase3_training_integration.md. Verified end-to-end by Smoke PrimeIntellect-ai#4 (5-step stability run, batch=64) where entropy/mean = 0.28735 across all steps, with meaningful variation across the ▁▇█▄▅ sparkline that wasn't present before. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

First fix (kv-eviction e216849) added module-level verifiers monkey-patches in kv_eviction/env.py that make compaction_events flow from vLLM ChatCompletion responses through verifiers' client adapter and MultiTurnEnv.add_model_response into the trajectory step's extras dict. The __init__.py was updated to import env as a side effect, so patches would apply on any kv_eviction.* import. BUT the orchestrator never imports kv_eviction. Its only reference to compaction types is `from prime_rl.transport.types import CompactionEventWire`, and transport/types.py defines its own re-export rather than importing kv_eviction. So kv_eviction.__init__.py never runs in the orchestrator process and the monkey-patches never apply. AND the env server runs in a `mp.get_context("spawn").Process` subprocess with target=ZMQEnvServer.run_server. Spawn starts a fresh interpreter, so even if the parent imports kv_eviction, the subprocess (which is where rollouts, get_model_response, and add_model_response actually run) doesn't inherit the patched classes. Fix: 1. Add `import kv_eviction` at the top of prime_rl.orchestrator.envs — ensures the orchestrator main process patches verifiers before any env is loaded. 2. Add `_env_server_subprocess_entrypoint` that imports kv_eviction before delegating to ZMQEnvServer.run_server. Swap the subprocess target from ZMQEnvServer.run_server to this wrapper. This ensures the spawned subprocess also applies the patches before any rollout runs in it. Verified: post-fix smoke PrimeIntellect-ai#4 run (pre-commit, no events still) had 0/64 compaction events per step — confirms the patches weren't firing. With this commit, kv_eviction's monkey-patches fire in both the orchestrator process and the env-server subprocess, and compaction_events should now flow end-to-end from vLLM into TrainingExample.compaction_events. Expected impact on Mismatch KL: should drop from ~0.04 (the genuine divergence between full-context trainer reforward and evicted-context inference) to ~0.001 (the offline kernel floor measured by compare_segforward_modes.py). Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Smoke PrimeIntellect-ai#4 v4 stalled all 4 DP ranks at step 0 for 25+ minutes. Per-rank stderr logs show the trainer stuck right after the Inductor "Online softmax is disabled" warning — which is emitted by prime-rl's compute_entropy (trainer/rl/loss.py:52) decorated with @torch.compile(dynamic=True). Per-segment entropy aggregation calls compute_entropy once per segment. Segments produce [1, N, vocab] tensors with N varying (300-1000 depending on compaction boundaries). torch.compile's dynamic-shape handling does not generalize across the sizes segmented_forward produces — Inductor appears to recompile or specialize per new shape, which produces the observed stall. Fix: inline the entropy math (logsumexp - sum(softmax * logits)) in _segment_loss_fn under torch.no_grad. Identical semantics to compute_entropy, but runs as an eager kernel that doesn't hit Inductor's compile cache at all. Single-sample offline verification: yields the same entropy values as compute_entropy within float precision on a [1, 512, 151936] tensor. Cost: marginally higher per-segment wall-clock (~1 ms vs ~0.3 ms from compiled path), dominated by the softmax/logsumexp kernels themselves not by launch overhead. Negligible next to model forward (~100 ms per segment). Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Smoke PrimeIntellect-ai#4 v3 exposed a latent incompatibility between two features that a user can enable together in good faith: KV cache compaction (trainer.compaction.window_size > 0) and per-block activation checkpointing (trainer.model.ac != None). The failure mode at step 0 was: torch.utils.checkpoint.CheckpointError: Recomputed values for the following tensors have different metadata than during the forward pass. tensor at position 29: saved metadata: shape [1, 4098, 8, 128] recomputed metadata: shape [1, 8196, 8, 128] Mechanism: prime-rl's apply_ac wraps every decoder layer with torch.utils.checkpoint(preserve_rng_state=False) — non-reentrant mode. Non-reentrant re-runs each layer's forward during backward to recompute activations. Segmented_forward threads past_key_values through each segment via a DynamicCache, and each layer's forward calls DynamicCache.update() which appends the new K/V to the stored sequence. On the recompute call during backward, the SAME cache object is updated AGAIN, doubling the stored length. The checkpoint metadata guard catches the mismatch and raises. Per-segment backward mode does NOT sidestep this — the double-update happens inside window_loss.backward() within a single segment, before any subsequent segment starts. probe_ac_cache_mutation.py reproduces it on a minimal 1024-token single-segment case. Fix: add a config-level validator that rejects the combination with an actionable error message pointing the user at removing the [trainer.model.ac] section. Memory budget is still fine with AC off because segmented_forward's per-segment backward bounds activation memory to O(1 segment), empirically tighter than per-block AC on the full sequence. (Smoke PrimeIntellect-ai#4 v5 hit peak 43.7 GiB without AC vs the AC-enabled baseline's 59 GiB.) Tested locally with TrainerConfig construction: good config (compaction + no AC) is accepted, bad config (compaction + AC) is rejected with the informative error. This closes the "why does my trainer crash the moment compaction events start flowing" footgun permanently — the next person who enables ac.freq=1 alongside compaction gets the explanation upfront instead of burning a 2-node interactive allocation to rediscover it. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

In a compaction training run (compaction.window_size > 0) every sample is now routed through segmented_forward, including rollouts whose inference never triggered a compaction event. Samples with events run as multi-segment forwards; event-less samples run as single-segment forwards (numerically identical to a plain text forward on the same un-packed sample, same flash_attention_2 kernel). Threaded compaction_enabled through BasePacker/SinglePacker/MultiPacker/ setup_packer/DataLoader/prepare_batch/prepare_sample/ packed_samples_into_micro_bs/pad_micro_batch, and flipped the per-micro-batch dispatch from "sample has events?" to "config.compaction.window_size > 0". When enabled, every non-multimodal sample becomes its own un-packed micro-batch with continuous position_ids and prompt_len set. Non-compaction runs (window_size == 0) are unaffected. Root cause of smoke PrimeIntellect-ai#4's step-1 OOM was NOT the rank-level collective divergence the previous plan section hypothesized. Actual cause: each rank processed its compaction micro-batches first, then its packed text micro-batches. The per-segment allocation pattern left the CUDA caching allocator holding ~11 GiB of cached blocks sized for per-layer/per-segment tensors that didn't fit the contiguous shapes the 16k packed text forward needed, pushing over 80 GiB on step 1 once the optimizer state allocated. Validated by an FSDP2 4-rank probe that reproduced the exact OOM at matched seq_len, and by the allocator-reserved delta between text-only (68.5 GB) and compaction-then-text (79.7 GB) at seq_len=8192. Smoke PrimeIntellect-ai#4 v6 validates the fix end-to-end. All 5 steps pass. Peak memory flat at 45.9 GiB across steps 1-4 (vs v5 OOM at 79+ on step 1). Mismatch KL stays at kernel floor (0.0009-0.0010). Loss, entropy, grad norm, and reward all track normally. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

- renderers fe67f9f (PR #4 squash-merge on main) - verifiers 8e56263 (PR #1298 rebased on main) Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Adds two new fields to RendererConfig: - preserve_all_thinking - preserve_thinking_between_tool_calls The orchestrator forwards them in two places so train and infer stay consistent: 1. create_renderer() — bound at construction on the training-side renderer used by build_trajectory_step / render_ids. 2. setup_inference_pool() → setup_clients() → vf.ClientConfig — the verifiers RendererClient picks them up and forwards to its create_renderer_pool, so every inference render carries the same thinking-preservation behaviour. Both flags are off by default → zero behaviour change for existing configs. Setting either without orchestrator.use_renderer=True is rejected by validate_renderer_args, matching the existing renderer knobs. Bumps source pins: - verifiers 3b77145 → a7516a1 (PR #1298: ClientConfig propagation) - renderers (now a separate repo): pinned to fe67f9f (PR #4: construction-time preserve flags) Re-applies #2433, which was merged into feat/unify-inference-generate by mistake. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

…2436) * feat(renderer): expose preserve_*_thinking flags in RendererConfig Adds two new fields to RendererConfig: - preserve_all_thinking - preserve_thinking_between_tool_calls The orchestrator forwards them in two places so train and infer stay consistent: 1. create_renderer() — bound at construction on the training-side renderer used by build_trajectory_step / render_ids. 2. setup_inference_pool() → setup_clients() → vf.ClientConfig — the verifiers RendererClient picks them up and forwards to its create_renderer_pool, so every inference render carries the same thinking-preservation behaviour. Both flags are off by default → zero behaviour change for existing configs. Setting either without orchestrator.use_renderer=True is rejected by validate_renderer_args, matching the existing renderer knobs. Bumps source pins: - verifiers 3b77145 → a7516a1 (PR #1298: ClientConfig propagation) - renderers (now a separate repo): pinned to fe67f9f (PR #4: construction-time preserve flags) Re-applies #2433, which was merged into feat/unify-inference-generate by mistake. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> * fix(sft): forward preserve_*_thinking flags through SFT renderer path The flags were added to the shared RendererConfig but only wired up in the orchestrator. SFT also constructs a renderer via create_renderer for training-side tokenization, so it must forward both flags or silently ignore them. Also tighten validate_renderer_args to reject either flag when use_renderer=False, matching the orchestrator validator. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> --------- Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

…ht resolver - #1 reserve loss-term names {sft, opd} (any term) and `rl` (non-primary), so an overlay can't silently overwrite a training_mode dispatch core or the rl primary in the trainer registry. - #5 resolve the primary's advantage weight orchestrator-side: scale the per-token advantage by the advantage-weight's tau in process_group and drop adv_tau from the dppo_kl core / RLLossConfig. Now *any* primary core (dppo_kl or custom) gets the resolved advantage × tau — no per-core special-case. Bit-identical for the default tau=1.0. - #4 overlay trainability = non-None AND non-zero, so a zero-weight overlay (e.g. advantage-weighted with zero advantage) no longer keeps an otherwise-empty batch alive past the empty-batch guard. - #6 custom overlay weight resolver is group-aware: it now receives `WeightInputs{sample, rollouts}` (the full GRPO group) instead of a lone sample, so it can compute group-relative weights. - #10 document the overlay_mask/overlay_weight token-export columns (schema v2) in the configs skill. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

samsja added 2 commits February 20, 2025 23:38

add schema for infernece

1bcb7c5

Signed-off-by: Sami Jaghouar <sami.jaghouar@gmail.com>

fix inference schema

23c5adf

Signed-off-by: Sami Jaghouar <sami.jaghouar@gmail.com>

samsja force-pushed the data-streaming branch 2 times, most recently from 5214445 to 39c1c66 Compare February 21, 2025 04:50

add updatable datasets

bf9ce67

Signed-off-by: Sami Jaghouar <sami.jaghouar@gmail.com>

samsja force-pushed the data-streaming branch from 39c1c66 to bf9ce67 Compare February 21, 2025 04:52

remove kl for now

56f0cf5

Signed-off-by: Sami Jaghouar <sami.jaghouar@gmail.com>

samsja marked this pull request as draft February 21, 2025 05:13

apaz-cli reviewed Feb 21, 2025

View reviewed changes

samsja added 5 commits February 21, 2025 21:43

add beartype

e2a787d

Signed-off-by: Sami Jaghouar <sami.jaghouar@gmail.com>

fix data stream loading

44c1aaa

Signed-off-by: Sami Jaghouar <sami.jaghouar@gmail.com>

fix grpoype hint

74ac92c

Signed-off-by: Sami Jaghouar <sami.jaghouar@gmail.com>

add tests for no debug data

2d73b9c

Signed-off-by: Sami Jaghouar <sami.jaghouar@gmail.com>

add tests for pq datasets

63a4544

Signed-off-by: Sami Jaghouar <sami.jaghouar@gmail.com>

samsja marked this pull request as ready for review February 21, 2025 23:49

samsja requested a review from Jackmin801 February 21, 2025 23:50

samsja added 2 commits February 22, 2025 00:31

fix loss

114d0bc

Signed-off-by: Sami Jaghouar <sami.jaghouar@gmail.com>

fix inference

c4bbc5c

Signed-off-by: Sami Jaghouar <sami.jaghouar@gmail.com>

samsja closed this Feb 25, 2025

samsja added a commit that referenced this pull request Mar 30, 2026

update 12 26 (#4)

bf9315b

* update verifier * update mem usage --------- Co-authored-by: Mika Senghaas <mail@mikasenghaas.de>

hallerite mentioned this pull request May 6, 2026

feat(renderer): expose preserve_*_thinking flags in RendererConfig #2433

Merged

4 tasks

hallerite mentioned this pull request May 7, 2026

feat(renderer): expose preserve_*_thinking flags in RendererConfig #2436

Merged

sfc-gh-kganesan mentioned this pull request May 22, 2026

feat: integrate Arctic RL backend (opt-in via [arctic]) sfc-gh-kganesan/prime-rl#1

Draft

mikasenghaas mentioned this pull request May 24, 2026

perf(orchestrator): reduce event-loop lag #2609

Merged

hallerite mentioned this pull request Jul 29, 2026

chore(configs)!: remove every config shim #3156

Merged

mikasenghaas deleted the data-streaming branch August 5, 2026 04:26

Navigation Menu

Uh oh!

Data streaming - #4

Data streaming#4
samsja wants to merge 11 commits into
mainfrom
data-streaming

samsja commented Feb 21, 2025

Uh oh!

apaz-cli Feb 21, 2025

Uh oh!

samsja Feb 21, 2025

Uh oh!

samsja Feb 21, 2025

Uh oh!

apaz-cli Feb 21, 2025

Uh oh!

samsja Feb 21, 2025

Uh oh!

apaz-cli Feb 21, 2025

Uh oh!

apaz-cli Feb 21, 2025

Uh oh!

samsja Feb 21, 2025

Uh oh!

apaz-cli left a comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

		yield {"input_ids": input_ids, "advantages": advantages}


		class ParquetDataset(IterableDataset):

Uh oh!

Conversation

samsja commented Feb 21, 2025

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

apaz-cli left a comment

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants