nemo_rl.models.generation.megatron.token_capture#

Megatron Inference (MInf) hooks for NeMo-Gym token capture.

The Megatron generation worker installs these two adapters on the dynamic inference engine of the model-parallel coordinator:

  • TQMegatronPromptPreparer resolves a Gym-authorized staging_chain prefix from TransferQueue and splices it into the rendered prompt before the engine admits the request.

  • TQMegatronTokenStager canonicalizes the finished completion through Gym’s capture core and writes the same TQ row the vLLM worker writes.

Both reach TransferQueue only through the backend-neutral TQTokenSink / TQTokenSource in nemo_rl.data_plane.tq_token_sink; this module is the Megatron analog of the capture glue in vllm_worker_async.py.

Module Contents#

Classes#

TQMegatronPromptPreparer

Resolve a Gym-authorized staged prefix before MInf admits a request.

TQMegatronTokenStager

Canonicalize one admitted MInf completion through Gym’s capture core.

API#

class nemo_rl.models.generation.megatron.token_capture.TQMegatronPromptPreparer(
source: nemo_rl.data_plane.tq_token_sink.TQTokenSource,
)#

Resolve a Gym-authorized staged prefix before MInf admits a request.

Mirrors the vLLM worker: prepare_prompt resolves the prefix through the shared resolve_admission_prefix / ChainPrefixCache pair, then splices it with the shared replace_prefix_tokens using the rendered prior-turn tokens and EOS id the Megatron endpoint carried in offload_params.

Initialization

prepare_prompt(
prompt: str | list[int] | torch.Tensor,
*,
offload_params: dict[str, Any] | None = None,
) → megatron.core.inference.inference_request.RequestPromptPreparationResult#

Fetch a chained prefix, splice it into the prompt, and update admission.

class nemo_rl.models.generation.megatron.token_capture.TQMegatronTokenStager(
sink: nemo_rl.data_plane.tq_token_sink.TQTokenSink,
)#

Canonicalize one admitted MInf completion through Gym’s capture core.

MInf owns the exact prompt/output material and its per-request policy epoch. Gym owns the lineage admission carried opaquely as ng_capture. This adapter joins them before the response leaves MInf, writes the same canonical TQ row as vLLM, and returns lightweight commit coordinates.

Initialization

property epoch_span_count: int#

Number of staged calls whose generation spanned more than one policy epoch.

_weight_version(finished_metadata: Any) → int#

Stamp the policy epoch the request was admitted under.

The engine records policy_epoch as (token_index, epoch) boundaries: one at admission, plus one appended on every set_generation_epoch while the request is active, so a request that straddles a refit carries several. vLLM stamps the version in effect at begin_call and never re-checks, so the admission epoch (first boundary) is the matching choice here. Spans are counted and logged rather than masked; _abort_stale_inflight is skipped on the Gym path (#2625), so they are routine under async rollouts.

stage(
uid: str,
payload: Any,
*,
finished_metadata: Any,
offload_params: dict[str, Any] | None = None,
) → megatron.core.inference.inference_request.RequestPayloadStageResult | None#

Stage an admitted request, or decline ordinary non-capture traffic.

_stage_admitted(
payload: Any,
*,
capture_payload: Any,
finished_metadata: Any,
) → megatron.core.inference.inference_request.RequestPayloadStageResult#

Validate and stage traffic that carries a Gym capture admission.