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.

_init_inference_engine_state() → None#

Reset all inference-engine attributes to their uninitialized state.

_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() → None#

Wind down a generation cycle.

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

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

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 “nvshmem”).

preinit_nvshmem_collective() → None#

Initialize NVShmem collectively before any weight transfer.

Must be called on ALL participating ranks (training + inference) simultaneously, after prepare_for_generation() has completed and the CG has been recorded. The NVSHMEMCopyService lazy init 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.

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.