What data crosses each hop when Megatron captures tokens — and is it the same as vLLM?

NVIDIA-NeMo/RL #4129 · Gym #2823 pinned at 9fc05c0 · Megatron-LM #7015 pinned at c711dc0 · RL permalinks pin head 0f38893.

Before
Only vLLM could record the exact token ids of each model call. Token capture with policy.generation.backend=megatron failed at setup.
This PR
Adds the same capture for Megatron inference, using two hooks inside Megatron's engine.
Status
In review — the data matches vLLM at every hop. The open review items are about tests and speed, not about what data moves where.

The idea is the same for both backends. Gym sends a chat request plus a small admission record that says which earlier call this one continues. The inference worker runs the model and writes the heavy data (token ids, mask, logprobs) to TransferQueue (TQ). Then it returns only the text plus a small receipt, ng_commit_coords, that points at the TQ row. Gym keeps a token-free ledger of those receipts. Before training, a finalizer reads the rows back by key and joins them into one sample.

One call, followed across every hop

Turn 2 of a rollout. Turn 1 is already in TQ under key k1 (7 tokens). The chat template adds 5 new prompt tokens, the model writes 4, and the policy weights are version 3.

heavy: token ids / logprobs light: ids, lengths, hashes sameboth backends do it the same way ≠differs (see table) Gym model server Inference worker TransferQueue Finalizer same in both? A · request: chat text + admission external_capture.py:271-316 ≠ field location B · read k1 → 7 token ids ChainPrefixCache.fetch same code prompt = 7 + 5 = 12generate 4 tokens ≠ where it runs C · write row k2: 9 tokens TQTokenSink.stage same code D · reply: text + receipt for k2 same shape E · ledger rowno tokens, only after D same code F · read k1, k2 → 16 tokens fetch_for_finalization same code

What each message holds, for both backends

HopWhat crosses (this example)vLLMMegatron
A request chat messages (text); admission = {rollout r, call c2, parent c1, prev_len 7, mode token_in, staging_chain [k1], parent_chain_hash}; logprobs=true. No prefix token ids go over HTTP. admission at top-level ng_capture (:276) ≠ admission at offload_params.ng_capture; must be n=1, text only (:296). The endpoint also adds the re-rendered turn-1 tokens and the EOS id to offload_params (chat_completions.py:602).
B prefix read TQ get [k1] → the 7 exact ids turn 1 produced. Spliced in front of the 5 new template tokens with replace_prefix_tokens. in a worker thread (:898) ≠ inside the engine step, via prepare_prompt (tq_token_sink.py:431)
C row write key k2: token_ids_delta 9 ids (5 prompt + 4 generated), token_mask_delta [0,0,0,0,0,1,1,1,1], generation_logprobs_delta 9 floats (0.0 on prompt), prev_len 7 · delta_len 9 · cum_len 16, weight_version 3, digests (digest.py:281) version = worker value when the call began same row; version = epoch the engine stamped at admission (:563). Routed-expert columns only from vLLM: setup rejects router replay on Megatron.
D reply assistant text + ng_commit_coords = {staging_key k2, prev_len 7, delta_len 9, cum_len 16, weight_version 3, digests}. No token ids, logprobs, or routes. worker strips them itself (:694) same shape: the engine marks the reply offloaded and drops them (dynamic_engine.py:1564, inference_request.py:932)
E ledger CallRecord with lengths, staging_key k2, hashes, response id — no tokens. Written only after D, so a call can't become a parent before its row is safe in TQ. Gym also strips any leftover heavy fields again before the agent sees the reply (:60). same Gym code for both (:178)
F finalize get [k1, k2] → 7 + 9 = 16 tokens, one training sample. Rejects the rollout if a row is missing or fails its checks. same code; never knows the backend

Where they really differ

DifferenceDoes it matter?
Where the admission sits in the request (A)No. Same record, different field name.
Where the TQ read and write run (B, C)Speed. vLLM moves them off its event loop. Megatron runs them inside the engine step, so the whole model-parallel group waits for each one. Megatron-LM needs the receipt before it replies, so RL can't move them itself. The review asks for a capture-on vs capture-off throughput number.
Where the weight version comes from (C)No. Both stamp the version in effect when the call began, even if a refit lands mid-generation.
Routed experts (C)A missing feature, not a silent bug: Megatron capture with router replay fails at setup.
Stager fails (C)Same outcome. Megatron sends the heavy fields back with no receipt. Gym strips them and marks the rollout bad, as it does for vLLM.
So what. The data model holds: only light data crosses HTTP in both directions, heavy data goes to TQ exactly once per call, and the row format and the code that writes it are shared. The risk left in this PR isn't the shape of the data. It's that the Megatron end-to-end test has not run in CI yet, and the TQ calls add waiting time inside the Megatron engine step.