Token capture finalizer: batch the group staging fetch, and give the pool a placement policy + a splittable timer
- Dominant language
- Python
- Stars
- 2k
- Forks
- 561
- Avg merge
- 4d 5h
- Merged PRs (30d)
- 145
Description
Follow-ups from review of #3837 (gate-authoritative token capture). Neither blocks that PR — both were raised there as non-blocking comments and are filed here so they are not lost on merge.
Both concern the same component, the CPU finalizer pool, and they interact: you cannot pick between them without the timer split in (2).
```
SingleController ── group ──► FinalizerActor pool (default 2, unplaced)
│
├─ 1 fetch per rollout ──► TQ (1)
└─ per-token verify in Python (2)
row_assembly/rollouts_ms measures BOTH ──► cannot tell which dominates
```
---
## 1. `finalize_group` issues N staging fetches where 1 would do
Ref: https://github.com/NVIDIA-NeMo/RL/pull/3837#discussion_r3938568784
`finalize_group` calls `finalize_rollout` per rollout, and each does its own
`fetch_for_finalization` → `kv_batch_get`. All N rollouts' staging keys are known before
the first fetch (the receipts are already in hand, and receipt → `staging_keys` is pure
local work), so one batched read over the union would serve the group.
```
now for rollout in group: fetch(that rollout's keys) 2N messages
proposed phase 1 parse N receipts -> all keys (no I/O)
phase 2 ONE fetch over the union 2 messages
phase 3 per-rollout verify from a key -> row map
```
Each `kv_batch_get` is constant in key count, but its `kv_retrieve_meta` hop lands on the
TQ controller's single request thread, and the pool is bounded, so the N calls serialize
inside one actor. `finalize/calls_per_rollout == 1.0` in observed runs, meaning each
current call fetches a *single* key — per-call overhead dominates payload.
Not measured. Byte volume is unchanged; the claim is `(N-1)` fewer serialized round trips
per group.
**Two guards have to become per-key first** — this is the actual work, not the batching:
- the row-count check and `_row_to_base_snapshot` are both request-scoped today, so one
missing or malformed row would reject all N rollouts instead of one. Rows carry their own
identity (`staging_key` derived from `rollout_id_utf8` + `model_call_id_utf8`, already
asserted against the requested key), so both can be made per-key.
- direct mode's extras probe falls back to the base schema on any exception, and
`kv_batch_get` readiness is all-or-nothing across a batch — so one row that staged no
routes degrades every row in the request. Reading `STAGING_FIELDS` first, then
`[ROUTED_EXPERTS_FIELD]` for only the keys with `encoding != 0`, makes readiness hold by
construction. This one has value independent of the batching.
---
## 2. Finalizer pool has no placement policy and no way to size itself
Ref: https://github.com/NVIDIA-NeMo/RL/pull/3837#discussion_r3938814571
### 2a. Placement
```
actor scheduling_strategy
───────────────────────── ────────────────────────────────
GPU worker groups PlacementGroupSchedulingStrategy
GenerationRouterActor NodeAffinity(driver, soft=False)
NemoGym actor NodeAffinity(driver, soft=True)
FinalizerActor — none —
```
`FinalizerActor.remote(...)` takes no `.options()`, so `num_cpus=1` is the whole policy and
Ray places the pool on whatever node has a free logical CPU. Correctness is unaffected (no
GPU state; the actor connects rather than bootstraps), but network locality to the storage
units is left to luck, and under Mooncake each TQ client mounts its own segment — so N
finalizers reserve N × `global_segment_size` of DRAM on unpredictable nodes, which
`num_cpus=1` does not express.
Fix: match `GenerationRouterActor`, with `soft=True` since the pool scales with
`num_finalizer_workers` and must be allowed to spill.
### 2b. Sizing — and why this is CPU-bound, not I/O-bound
This is the part that decides whether more actors help at all, so stating the evidence
rather than asserting it.
**The architecture already concedes it.** The pool is N separate Ray *processes*, each
reserving a whole CPU. If the work were I/O-bound, a thread pool or an inline
`asyncio.to_thread` on the controller would have been the cheaper answer. Separate
processes buy exactly one thing — separate GILs. `num_cpus=1` is then the correct
reservation, since a GIL-bound actor cannot use more than about one core anyway.
**The work is per-token pure Python, repeated per staged call:**
```
per staged call, over delta_len tokens
tq_token_sink.py:527 token_ids_delta [int(t) for t in ....tolist()]
:529-530 token_mask_delta [float(m) for m in ....tolist()]
:531-533 logprobs_delta [float(p) for p in ....tolist()]
:479 id/text fields bytes(int(v) for v in ....tolist())
Gym verify_and_linearize digest recompute + validity scan
× calls/rollout × rollouts/group × groups/step
```
Every one of those is `.tolist()` (allocates a Python list) followed by an interpreter loop
over it. Only the digest's hashing leaves Python.
**Two honest caveats:**
- The digest and validity passes now live in Gym's `verify_and_linearize`
(`blackbox_finalizer.py:222-223` — "All base token/digest/lineage/terminal semantics
belong to Gym; the finalizer never re-verifies them"). They still run on the finalizer
actor's CPU, but I could not read that implementation from my workspace, so its cost is
inferred from its inputs. The three list conversions above are directly verifiable in RL.
- The digest is *recomputed*, not trusted — that is the ledger's security property, so it
cannot be skipped. It does mean the token data is hashed a second time after the writer
already hashed it once.
**Why the default is untested rather than wrong.** `num_finalizer_workers: 2` is set by no
recipe, and the one functional test runs at `num_generations_per_prompt=2` /
`max_buffered_rollouts=4` — far too small to saturate. Single-call math workloads
(`finalize/calls_per_rollout == 1.0`, short deltas) will be fine at 2. The multi-turn
agentic chains this feature exists for multiply every term above, and that is the regime
nothing has exercised.
The obstacle to picking a number is one merged timer: `row_assembly/rollouts_ms` spans the blocking fetch *and*
the Python passes.
```
fetch_ms dominant -> do item 1, don't add actors
verify_ms dominant -> add actors, ceiling = free CPUs
```
Splitting it into `row_assembly/fetch_ms` and a derived `verify_ms` is ~3 lines and makes the
choice decidable. `finalize/queue_depth` and `finalize/queue_wait_ms` already exist, so rising
queue depth is the undersize signal once the split lands.
---
**Suggested order:** the timer split first (cheapest, and it decides the rest), then the
placement fix, then the group-fetch batching if `fetch_ms` turns out to dominate.
Contributor guide
Assessment
This issue has not been assessed yet.