nemo_rl.models.generation.megatron.megatron_worker

Module Contents

Classes

MegatronGenerationMixin	Engine lifecycle, coordinator, HTTP server, and finish-generation machinery.
MegatronGenerationRefitMixin	Refit collective, weight transfer, and engine suspend/resume around refits.

API

class nemo_rl.models.generation.megatron.megatron_worker.MegatronGenerationMixin

Engine lifecycle, coordinator, HTTP server, and finish-generation machinery.

The host class must provide:

• model: the megatron module.

• cfg: policy config (TypedDict).

• rank: global rank (used for logging).

• tokenizer: HF tokenizer.

• megatron_tokenizer: tokenizer for inference.

• is_generation_colocated: Whether colocated or distributed.

• _reserved_http_server_socket: driver-reserved server socket, or None.

inference_model

None

_colocated_reshard_plan

None

_gen_model() → megatron.core.transformer.MegatronModule

The model the inference engine wraps.

Returns the dedicated inference-layout model when one exists (colocated reshard), otherwise the shared training model.

_init_inference_engine_state() → None

Reset all inference-engine attributes to their uninitialized state.

_setup_colocated_cuda_graph_managers() → None

Create inference CUDA-graph managers for shared-model colocated generation.

Colocated policies build the shared training model without CUDA graphs. But CUDA graphs configs MUST be in-place at initialization.

Must run at worker init: colocated reshard is detected via the pending reshard plan (its dedicated inference model is built later, with graphs enabled); the plan is consumed on the first generation cycle.

get_inference_cuda_graph_capture_count() → int

Inference CUDA graphs captured in this worker process (0 = eager decode).

_initialize_inference_engine(mcore_generation_config: dict) → None

Initialize the persistent inference engine and client.

async _start_inference_coordinator()

Start the inference coordinator and engine loop.

_sleep() → None

Pause + suspend the engine. No-op if already asleep.

async _sleep_engine()

_wake() → None

Resume + unpause the engine. No-op if already awake.

async _wake_engine()

_start_inference_loop_thread()

Start a background thread with a persistent event loop for inference.

_setup_openai_api_server() → str

Start the OpenAI-compatible HTTP server on this worker.

_run_async_coordinator_start()

Start the coordinator and engine loop in the background thread.

finish_generation(*, release_gpu: bool = True) → None

Wind down a generation cycle.

Parameters:

release_gpu – the caller needs the GPUs for itself (a training step or a checkpoint save), so even a colocated engine must fully stand down. Pass False between generation phases (e.g. after validation) to let a colocated engine keep serving; tearing it down there would discard KV/prefix caches and CUDA graphs for no reason. For non-colocated engines this body reduces to a rotary-cache clear (their engine pause happens in suspend_for_refit).

prepare_for_generation(tags=None, **kwargs) → None

Enter inference mode and start (or wake) the inference engine.

Idempotent wake: a plain call (no tags) on an already-serving engine returns immediately. Refit-protocol calls (tags) are never skipped.

Called in both colocated and non-colocated setups. Even in non-colocated mode, Megatron’s engine has to be intentionally paused before a refit (and its weights are not detachable), so we have to switch modes around every refit.

report_dp_openai_server_base_url() → Optional[str]

Return this worker’s OpenAI server base URL (None if not the leader).

_build_sampling_params(

greedy: bool,

stop_words: Optional[list[str]],

) → megatron.core.inference.sampling_params.SamplingParams

Build mcore SamplingParams for a single request.

_merge_stop_strings(

batch_stop_strings: Optional[list[Optional[list[str]]]],

) → Optional[list[str]]

Union the config’s stop_strings with the given per-sample stop strings.

_prepare_data_for_generation(

data: nemo_rl.distributed.batched_data_dict.BatchedDataDict[nemo_rl.models.generation.interfaces.GenerationDatumSpec],

greedy: bool = False,

) → tuple[torch.Tensor, torch.Tensor, list[megatron.core.inference.sampling_params.SamplingParams]]

Build the prompt tensors and a per-request SamplingParams for each sample.

_parse_result_to_batched_data_dict(

data: nemo_rl.distributed.batched_data_dict.BatchedDataDict[nemo_rl.models.generation.interfaces.GenerationDatumSpec],

result: list,

) → nemo_rl.distributed.batched_data_dict.BatchedDataDict[nemo_rl.models.generation.interfaces.GenerationOutputSpec]

Pack DynamicInferenceRequest results into a GenerationOutputSpec batch.

generate(

*,

data: nemo_rl.distributed.batched_data_dict.BatchedDataDict[nemo_rl.models.generation.interfaces.GenerationDatumSpec],

greedy: bool = False,

) → nemo_rl.distributed.batched_data_dict.BatchedDataDict[nemo_rl.models.generation.interfaces.GenerationOutputSpec]

Synchronous batched generation via the mcore data-parallel coordinator.

Parameters:

• data – BatchedDataDict containing input_ids and input_lengths tensors

• greedy – Whether to use greedy decoding instead of sampling

Returns:

• output_ids: input + generated token IDs with proper padding

• logprobs: Log probabilities for tokens

• generation_lengths: Lengths of each response

• unpadded_sequence_lengths: Lengths of each input + generated sequence

Return type:

BatchedDataDict conforming to GenerationOutputSpec

async generate_async(

data: nemo_rl.distributed.batched_data_dict.BatchedDataDict[nemo_rl.models.generation.interfaces.GenerationDatumSpec],

greedy: bool = False,

) → AsyncGenerator[tuple[int, nemo_rl.distributed.batched_data_dict.BatchedDataDict[nemo_rl.models.generation.interfaces.GenerationOutputSpec]], None]

Streaming generation: yield (index, batch) tuples as they complete.

Parameters:

• data – BatchedDataDict with input_ids and input_lengths

• greedy – Whether to use greedy decoding instead of sampling

Yields:

Tuple of (original_index, BatchedDataDict conforming to GenerationOutputSpec for the single sequence)

async _generate_with_persistent_engine(

prompt_tokens_tensor: torch.Tensor,

prompt_lengths_tensor: torch.Tensor,

sampling_params: list[megatron.core.inference.sampling_params.SamplingParams],

) → list

Submit requests through the persistent inference client (rank 0 only).

class nemo_rl.models.generation.megatron.megatron_worker.MegatronGenerationRefitMixin

Refit collective, weight transfer, and engine suspend/resume around refits.

init_collective_mcore_generation(

ip: str,

port: int,

world_size: int,

rank_offset: int,

refit_backend: str = 'gloo',

) → None

Initialize the refit collective for non-colocated weight transfer.

Parameters:

• ip – IP address for the process group rendezvous.

• port – Port for the process group rendezvous.

• world_size – Total world size (train + inference workers).

• rank_offset – Offset for this side’s ranks (train_world_size for inference).

• refit_backend – Copy-service backend (“gloo” or “nccl”; “nvshmem” is currently broken, see the issue below).

preinit_nvshmem_collective() → None

Initialize NVShmem collectively before any weight transfer.

Must be called on ALL participating ranks (training + inference) simultaneously, outside CUDA graph capture. Lazy initialization during graph recording or replay can corrupt CUDA graph state.

swap_weights_via_reshard(is_source: bool) → bool

Transfer weights using Megatron’s swap_model_weights API.

Parameters:

is_source – True for training workers (senders), False for inference workers (receivers).

Returns:

True on success.

_onload_inference_model() → None

Restore the colocated inference weights to GPU before resharding / generation.

_offload_inference_model() → None

Offload the colocated inference weights to CPU while training runs.

_reshard_into_inference_model() → None

Reshard current training weights into the colocated inference-layout model.

suspend_for_refit() → None

Pause+suspend the inference engine before a weight refit.

resume_after_refit() → None

Resume+unpause the inference engine after a weight refit.