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,
skip_weight_load: bool = False,
reserved_http_server_port: Optional[int] = 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).

  • skip_weight_load – Do not load the weights from the checkpoint; refit will do it.

  • reserved_http_server_port – Driver-reserved OpenAI server port for non-colocated.

static effective_megatron_cfg(
config: nemo_rl.models.policy.PolicyConfig,
) → dict[str, Any]#

The megatron_cfg the generation workers actually run with.

Colocated generation shares the training model, so the training values apply; non-colocated builds a dedicated policy with mcore_generation_config merged on top. Always returns a fresh dict.

Largest GPU group requiring full NVLink connectivity.

Colocated reshard hosts a second, inference-layout model on the same ranks.

classmethod 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).

classmethod reserve_http_server_address(
cluster: nemo_rl.distributed.virtual_cluster.RayVirtualCluster,
config: nemo_rl.models.policy.PolicyConfig,
) → tuple[str, int, ray.actor.ActorHandle]#

Reserve the OpenAI server address before any generation worker exists.

This is megatron’s substitute for vLLM’s defer_model_load overlap. See https://github.com/NVIDIA-NeMo/RL/issues/3752

Parameters:

  • cluster – The cluster the generation workers will run on.

  • config – The full PolicyConfig.

Returns:

Tuple of (server base URL, reserved port, port-holder actor handle). The caller must keep the handle referenced until rank 0 has adopted the socket (worker init complete), then ray.kill it.

classmethod validate_settings(
master_config: nemo_rl.algorithms.single_controller_utils.config.MasterConfig,
) → None#

Reject config the Megatron generation backend cannot honor.

classmethod verify_served_address(
served_urls: list[Optional[str]],
reserved_url: str,
) → None#

Fail loud if the engine serves anywhere but the pre-published address.

property worker_group: nemo_rl.distributed.worker_groups.RayWorkerGroup#

The underlying policy’s worker group (fleet-health probes read dp_size).

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 “nccl”; “nvshmem” is currently broken and warns at setup).

Returns:

List of Ray ObjectRefs for the collective init futures.

update_weights_from_collective(
refit_timeout_s: Optional[float] = None,
) → 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(*, release_gpu: bool = True) → bool#

Clean up after generation.

When release_gpu is False, a colocated engine keeps serving instead of standing down.

blocks_training() → bool#

Whether the engine must stand down before a training step.

Colocated generation shares the training GPUs, so the training loop must wind the engine down before it can train.

wake_carries_weight_updates() → bool#

The colocated wake reshards (or shares tensors); see the ABC.

invalidate_kv_cache() → bool#

Report whether weight updates invalidate the KV cache.

Under “recompute” mode the engine drops and rebuilds its KV cache across the suspend/resume that brackets every weight update, so invalidation is genuinely handled; report it truthfully instead of inheriting the interface’s False (which makes the trajectory collector warn every step).

preinit_nvshmem_collective() → list[ray.ObjectRef]#

Pre-initialize NVShmem collectively outside 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.