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_ports: Optional[dict[int, 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 for a dedicated, non-colocated inference Policy.

  • policy – Existing training Policy reused for colocated 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 weights from the checkpoint; refit will do it. Inference-engine initialization is deferred until that first refit so CUDA graphs capture the final persistent weight buffers rather than placeholder checkpoint tensors.

  • reserved_http_server_ports – Driver-reserved OpenAI server ports, keyed by the distributed rank that will adopt each one (non-colocated only).

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

The megatron_cfg the generation workers actually run with.

Colocated generation uses a dedicated inference-layout model when its resolved layout differs from training, and otherwise shares the training model. Non-colocated generation always builds a dedicated policy from 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).

static frontend_ranks(
cluster: nemo_rl.distributed.virtual_cluster.RayVirtualCluster,
config: nemo_rl.models.policy.PolicyConfig,
) → list[int]#

Distributed ranks that will host an HTTP frontend.

Copy of the engine’s is_mp_coordinator = tp_rank == 0 and pp_rank == 0. Megatron orders ranks TP-fastest, PP-slowest, so PP selects a leading slice and TP a stride within it. CP and DP are not divided out because the predicate ignores them.

static _rank_placement(
cluster: nemo_rl.distributed.virtual_cluster.RayVirtualCluster,
) → list[tuple[int, int]]#

The (placement group index, bundle index) each rank will occupy, in rank order.

Mirrors how Policy hands bundles to its worker group: a unified placement group is walked in NVLink-sorted order, otherwise bundles are taken group by group. Reserving a port for a rank means placing the holder on the bundle that rank will land on, so the two must agree.

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

Reserve every 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

One address is reserved per frontend-hosting rank so NeMo Gym can be handed the full set up front. Gym distributes sessions across the URLs it is given, so reserving only one would pin every session to a single frontend no matter how many the engine goes on to start.

Ports are only ever bound per node, so two frontends on different nodes may be handed the same port number; no cross-node uniqueness is sought.

Parameters:

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

  • config – The full PolicyConfig.

Returns:

reserved port}, port-holder actor handles). The caller must keep the handles referenced until the workers have adopted the sockets (worker init complete), then ray.kill them.

Return type:

Tuple of (server base URLs, {distributed rank

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_addresses(
served_urls: list[Optional[str]],
reserved_urls: list[str],
) → None#

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

Order is not significant – NeMo Gym spreads sessions over the set it was given – but the set must match exactly: a frontend serving an address Gym never received takes no traffic, and an address Gym holds that nothing serves fails its health check.

property uses_native_refit: bool#

Whether non-colocated refit uses Megatron Core’s native mechanism.

property worker_group: nemo_rl.distributed.worker_groups.RayWorkerGroup#

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

shard_liveness_ref(shard_idx: int) → ray.ObjectRef#

Liveness of the worker leading this shard. The caller need not know the layout.

Same shape as the vLLM backend’s. Declared here because this generation object wraps a Policy, so its workers are AbstractPolicyWorkers and answer is_alive like any other – the fleet probe needs no knowledge of which of the two it holds.

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

Join the configured refit collective after the training ranks.

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

Returns:

List of Ray ObjectRefs for the collective init futures.

update_weights_from_collective(
refit_timeout_s: Optional[float] = None,
) → list[ray.ObjectRef]#

Receive weights through the configured Megatron refit mechanism.

init_nccl_reshard_comm_group(
*,
pp_ips: list[str],
pp_ports: list[int],
pp_size: int,
train_ranks_per_stage: int,
sub_world_size: int,
) → list[ray.ObjectRef]#

Join every training PP stage’s NCCL M-to-N communicator.

set_refit_membership(
membership: nemo_rl.weight_sync.membership.RefitMembership,
) → None#

Record the inference ranks participating in nccl_reshard refits.

_refit_ranked_workers(
membership: Optional[nemo_rl.weight_sync.membership.RefitMembership] = None,
) → list[tuple[Any, int, int]]#

Return (actor, rebuilt rank, original rank) for live workers.

rebuild_collective(
membership: nemo_rl.weight_sync.membership.RefitMembership,
ip: str,
port: int,
) → list[ray.ObjectRef]#

Build the misc-weight communicator over the selected Megatron ranks.

rebuild_nccl_reshard_comm_group(
membership: nemo_rl.weight_sync.membership.RefitMembership,
*,
pp_ips: list[str],
pp_ports: list[int],
pp_size: int,
train_ranks_per_stage: int,
sub_world_size: int,
) → list[ray.ObjectRef]#

Build bulk communicators over the selected Megatron ranks.

prepare_nccl_reshard_refit_info(
refit_info: dict[str, Any],
) → None#

Build each inference worker’s HF-to-Megatron M-to-N receive map.

nccl_reshard_refit(
refit_timeout_s: Optional[float] = None,
) → list[ray.ObjectRef]#

Receive one NCCL M-to-N refit on every Megatron inference worker.

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.

setup_token_capture(
dp_cfg: nemo_rl.data_plane.interfaces.DataPlaneConfig,
staging_partition: str,
*,
capture_media: bool = False,
) → MediaColumnSpec | None#

Install MInf’s canonical prompt and completion capture hooks.

Returns the media column spec the coordinators pinned (None for text-only capture).

set_rollout_weight_version(version: int) → None#

Rotate the policy epoch stamped by MInf on subsequent requests.

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]],
) → Optional[list[str]]#

Prepare Bridge conversion tasks on every dedicated inference worker.

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.