nemo_gym.token_id_capture.adapters.megatron

Dependency-light extraction adapter for Megatron Inference offloaded payloads.

Module Contents

Classes

Name	Description
`MegatronCaptureAdapter`	Translate Megatron inference request/response material at the framework boundary.

Functions

Name	Description
`_field`	Read one field from a Megatron Inference payload object or an equivalent mapping.
`_log_probs`	Read a log-probability field, rejecting non-numeric elements.
`_sequence`	-
`_token_ids`	Read a token-id field, rejecting anything but plain integers.

Data

GENERATED_IDS_FIELD

GENERATED_LOGPROBS_FIELD

PREFIX_IDS_FIELD

PROMPT_IDS_FIELD

API

class nemo_gym.token_id_capture.adapters.megatron.MegatronCaptureAdapter()

Translate Megatron inference request/response material at the framework boundary.

Megatron inference offloads exact prompt ids, generated ids, and selected-token log probabilities as attributes on a payload object rather than as a chat completion dict. The adapter reads either shape so extraction failures flow through RolloutTokenCapture.complete_call_from_response and poison the call with capture_failed coordinates, the same outcome the vLLM adapter produces.

nemo_gym.token_id_capture.adapters.megatron.MegatronCaptureAdapter.enter_prefix(
    request_payload: dict[str, typing.Any],
    prefix_ids: list[int]
) -> dict[str, typing.Any]

nemo_gym.token_id_capture.adapters.megatron.MegatronCaptureAdapter.extract_extras(
    response_payload: typing.Any
) -> dict[str, typing.Any] | None

nemo_gym.token_id_capture.adapters.megatron.MegatronCaptureAdapter.extract_generation(
    response_payload: typing.Any
) -> tuple[list[int], list[float]]

nemo_gym.token_id_capture.adapters.megatron.MegatronCaptureAdapter.extract_prompt_ids(
    response_payload: typing.Any
) -> list[int]

nemo_gym.token_id_capture.adapters.megatron._field(
    payload: typing.Any,
    name: str
) -> typing.Any

Read one field from a Megatron Inference payload object or an equivalent mapping.

nemo_gym.token_id_capture.adapters.megatron._log_probs(
    payload: typing.Any,
    name: str
) -> list[float]

Read a log-probability field, rejecting non-numeric elements.

nemo_gym.token_id_capture.adapters.megatron._sequence(
    payload: typing.Any,
    name: str
) -> collections.abc.Sequence[typing.Any]

nemo_gym.token_id_capture.adapters.megatron._token_ids(
    payload: typing.Any,
    name: str
) -> list[int]

Read a token-id field, rejecting anything but plain integers.

Megatron Inference hands host-side list[int] values. A float, string, or bool element means the payload is malformed; int() would silently truncate or coerce it into a plausible-looking id.

nemo_gym.token_id_capture.adapters.megatron.GENERATED_IDS_FIELD = 'generated_token_ids'

nemo_gym.token_id_capture.adapters.megatron.GENERATED_LOGPROBS_FIELD = 'generated_log_probs'

nemo_gym.token_id_capture.adapters.megatron.PREFIX_IDS_FIELD = 'required_prefix_token_ids'

nemo_gym.token_id_capture.adapters.megatron.PROMPT_IDS_FIELD = 'prompt_token_ids'