nemo_rl.models.generation.megatron.megatron_generation#

Module Contents#

Classes#

MegatronGeneration

Generation interface backed by Megatron (colocated or non-colocated).

API#

class nemo_rl.models.generation.megatron.megatron_generation.MegatronGeneration(
config: nemo_rl.models.policy.PolicyConfig,
tokenizer: transformers.tokenization_utils_base.PreTrainedTokenizerBase,
cluster: Optional[nemo_rl.distributed.virtual_cluster.RayVirtualCluster] = None,
policy: Optional[nemo_rl.models.policy.lm_policy.Policy] = None,
name_prefix: str = 'megatron_generation',
processor: Optional[transformers.AutoProcessor] = None,
weights_path: Optional[str] = None,
)#

Bases: nemo_rl.models.generation.interfaces.GenerationInterface

Generation interface backed by Megatron (colocated or non-colocated).

Initialization

Initialize a MegatronGeneration instance.

Exactly one of cluster or policy must be provided.

Parameters:

  • config – PolicyConfig for the Megatron model.

  • tokenizer – The tokenizer for the model.

  • cluster – Cluster to deploy a dedicated inference Policy on.

  • policy – Existing training Policy to reuse for generation.

  • name_prefix – Prefix for naming the worker group (non-colocated only).

  • processor – Optional processor for VLMs (non-colocated only).

  • weights_path – Optional path to model weights (non-colocated only).

static init_cluster_placement_groups(
cluster: nemo_rl.distributed.virtual_cluster.RayVirtualCluster,
config: nemo_rl.models.policy.PolicyConfig,
) → None#

Pre-initialize the inference cluster’s placement groups.

Parameters:

  • cluster – The inference RayVirtualCluster.

  • config – The full PolicyConfig (megatron parallelism + colocation).

init_collective(
ip: str,
port: int,
world_size: int,
*,
train_world_size: int,
refit_backend: str = 'gloo',
) → list[ray.ObjectRef]#

Initialize the refit collective for weight synchronization.

Parameters:

  • ip – IP address for the process group rendezvous.

  • port – Port for the process group rendezvous.

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

  • train_world_size – Number of training workers (used to offset ranks).

  • refit_backend – Copy service backend (“gloo” or “nvshmem”).

Returns:

List of Ray ObjectRefs for the collective init futures.

update_weights_from_collective() → list[ray.ObjectRef]#

Receive updated weights from the training cluster via collective communication.

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]#

Generate a batch of data using the Megatron generation backend.

mcore’s data-parallel coordinator only accepts requests from DP rank 0 — the other workers’ engine loops drain the coordinator queue but never receive a Python-side call. So we dispatch straight to worker 0.

Parameters:

  • data – BatchedDataDict containing input_ids and input_lengths.

  • greedy – Whether to use greedy decoding.

Returns:

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]#

Generate asynchronously, yielding (index, batch) tuples as they complete.

prepare_for_generation(
*args: Any,
**kwargs: Any,
) → bool#

Initialize / re-enter inference mode on every worker.

First call starts the persistent inference engine, coordinator, and the OpenAI HTTP server. Subsequent calls re-enter inference mode after a refit.

finish_generation(*args: Any, **kwargs: Any) → bool#

Clean up after generation.

preinit_nvshmem_collective() → list[ray.ObjectRef]#

Pre-initialize NVShmem collectively after CUDA graph capture.

Must be called simultaneously on both training and inference workers.

suspend_for_refit() → None#

Suspend the inference engine for safe weight updates.

resume_after_refit() → None#

Resume the inference engine after weight updates.

prepare_refit_info(
state_dict_info: Optional[dict[str, Any]],
) → None#

Accept the cross-backend refit-prep contract; Megatron needs none of it.

start_gpu_profiling() → None#

Start GPU profiling on the dedicated inference workers.

No-op when colocated: the shared workers are already profiled through the training policy.

stop_gpu_profiling() → None#

Stop GPU profiling on the dedicated inference workers.

shutdown() → bool#

Shut down all inference workers and clean up resources.

__del__() → None#

Safety net to ensure workers are shut down.