`nemo_rl.models.generation.megatron.megatron_worker`#

Module Contents#

Classes#

`_MegatronRefitTask`	A Bridge import task and its mutable inference destination.
`_MegatronBulkRefitPiece`	One HF-local shard and the Megatron destination region it populates.
`MegatronGenerationMixin`	Engine lifecycle, coordinator, HTTP server, and finish-generation machinery.
`MegatronGenerationRefitMixin`	Refit collective, weight transfer, and engine suspend/resume around refits.

Functions#

`_inference_optimized_transformer_layer_spec`	Build the generic GPT layer spec backed by MCore inference linears.
`_configure_inference_optimized_layer_spec`	Select MCore inference linears for a Bridge generic-GPT provider.
`_resolve_mxfp8_refit_backend`	Resolve the MXFP8 storage required by the grouped-GEMM backend.
`_http_sampling_defaults`	Sampling defaults for HTTP requests that omit temperature, top_p or top_k.
`_resolve_coordinator_policy`	Effective DP-coordinator routing policy for this generation config.
`_apply_optional_inference_config_kwargs`	Forward configured optional MCore inference settings with validated types.
`_apply_inference_cuda_graph_scope`	Apply the optional CUDA graph scope to the engine model.

API#

nemo_rl.models.generation.megatron.megatron_worker._inference_optimized_transformer_layer_spec( config: Any, ) → Any#: Build the generic GPT layer spec backed by MCore inference linears.

nemo_rl.models.generation.megatron.megatron_worker._configure_inference_optimized_layer_spec( model_provider: Any, ) → None#

Select MCore inference linears for a Bridge generic-GPT provider.

Only the generic path needs this. Model-specific providers (DeepSeek, MiniMax, GLM, …) install get_gpt_decoder_block_spec, which already dispatches transformer_impl='inference_optimized' itself and — unlike the uniform spec built here — gives dense and MoE layers different specs. Overwriting those would rebuild a moe_layer_freq model’s dense layers as MoE layers, so leave any non-default spec alone.

Non-GPT providers are left alone too. Hybrid/Mamba stacks (Nemotron-H via HybridModelProvider) are not GPTModelProvider subclasses and resolve inference_optimized through their own stack spec, so this is a no-op for them rather than an error.

nemo_rl.models.generation.megatron.megatron_worker._resolve_mxfp8_refit_backend(model_config: Any) → str#: Resolve the MXFP8 storage required by the grouped-GEMM backend.

class nemo_rl.models.generation.megatron.megatron_worker._MegatronRefitTask#

A Bridge import task and its mutable inference destination.

conversion_task: Any#: None

destination: Any#: None

target_id: int#: None

property param_name: str#

property cache_key: tuple[int | None, str]#: Stable identity across Bridge task rebuilds.

property dependencies: tuple[str, ...]#

property expected_shape: torch.Size#

property is_mxfp8: bool#

property is_quantized: bool#

class nemo_rl.models.generation.megatron.megatron_worker._MegatronBulkRefitPiece#

One HF-local shard and the Megatron destination region it populates.

task: nemo_rl.models.generation.megatron.megatron_worker._MegatronRefitTask#: None

spec: Any#: None

shape: torch.Size#: None

device: torch.device#: None

destination: torch.Tensor | None#: None

nemo_rl.models.generation.megatron.megatron_worker._http_sampling_defaults( generation_cfg: Mapping[str, Any], ) → dict[str, Any]#

Sampling defaults for HTTP requests that omit temperature, top_p or top_k.

Training recomputes log probabilities with the configured generation sampling parameters. Without explicit server defaults, an omitted request field falls back to the server’s hardcoded value or to the model’s generation_config.json, so the request would sample off-policy.

nemo_rl.models.generation.megatron.megatron_worker._resolve_coordinator_policy( mcore_generation_config, ) → megatron.core.inference.config.PrefixCachingCoordinatorPolicy#

Effective DP-coordinator routing policy for this generation config.

Resolved in one place so the engine and the HTTP frontends cannot disagree: if the engine routes on prefix affinity but the frontends skip hashing, the coordinator silently sees no hashes and degrades to load balancing.

nemo_rl.models.generation.megatron.megatron_worker._apply_optional_inference_config_kwargs( inference_config_kwargs: dict[str, Any], mcore_generation_config: dict[str, Any], ) → None#: Forward configured optional MCore inference settings with validated types.

nemo_rl.models.generation.megatron.megatron_worker._apply_inference_cuda_graph_scope( engine_model: megatron.core.transformer.MegatronModule, mcore_generation_config: dict[str, Any], ) → None#: Apply the optional CUDA graph scope to the engine model.

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.
processor: optional multimodal processor.
is_generation_colocated: Whether colocated or distributed.
_reserved_http_server_port: driver-reserved server port, or None.

processor: Optional[Any]#: None

inference_model#: None

_colocated_reshard_plan#: None

_image_preprocessing_config: Optional[Any]#: 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.

_get_megatron_inference_wrapper_cls() → Optional[type]#

Resolve the configured Megatron inference wrapper, if any.

Returns:: The wrapper class, or None when no wrapper is configured.

static _wrapper_supports_modality( inference_wrapper_cls: Optional[type], modality: str, ) → bool#: Whether the configured inference wrapper advertises modality support.

_inference_model_and_media_parts(inference_wrapper_cls=None)#: Return the language model and its optional multimodal parent.

_vision_pixel_dtype() → torch.dtype#

The dtype the vision encoder casts pixels to before encoding them.

Read from the encoder’s own parameters rather than the policy dtype: a config may keep the vision tower in fp32 under a bf16 language model.

_build_image_preprocessing_config( generation_config: dict[str, Any], )#: Build raw-image preprocessing settings.

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

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

Install canonical TQ capture on each MInf model-parallel leader.

capture_media builds the sink/source against the media-enabled staging schema so the stager can hand the engine’s media tensors to TQ beside each call’s tokens. Returns the media column spec the coordinator pinned the staging column to (None for text-only capture and on followers, which host no hooks).

set_rollout_weight_version(version: int) → None#: Stamp subsequent MInf requests with the trainer weight version.

_build_sampling_params( greedy: bool, stop_words: Optional[list[str]], *, return_prompt_tokens: bool = False, ) → 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.

_sample_vision_tensors(data, index: int)#: Return one sample’s vision tensors from RL PackedTensors.

_build_prompt_and_multimodal_data(data, index: int)#: Build one pre-expanded token prompt and optional MCore media payload.

_prepare_data_for_generation( data: nemo_rl.distributed.batched_data_dict.BatchedDataDict[nemo_rl.models.generation.interfaces.GenerationDatumSpec], greedy: bool = False, ) → tuple[list[list[int]], list[Optional[Any]], list[megatron.core.inference.sampling_params.SamplingParams]]#: Build prompts, optional media payloads, and sampling parameters.

_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( prompts: list[list[int]], multi_modal_data_list: list[Optional[Any]], 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_generation_refit_state() → None#

Reset every refit attribute to its unprepared state.

Declaring these up front keeps prepare_refit_info readiness a value check rather than an hasattr probe, and makes the full set visible to anyone adding a __getstate__.

init_collective_mcore_generation( ip: str, port: int, world_size: int, rank_offset: int, refit_execution_batch_bytes: int | None, refit_backend: str, ) → 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_execution_batch_bytes – Configured native MCore staging limit. None uses NeMo-RL’s dynamic collective-refit default.
refit_backend – Copy-service backend (“gloo”, “nccl”, or “nccl_m2n”; “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.

init_nccl_reshard_comm_groups_generation( *, pp_ips: list[str], pp_ports: list[int], pp_size: int, train_ranks_per_stage: int, sub_world_size: int, rank_prefix: Optional[int] = None, ) → None#: Join every training PP stage’s M-to-N communicator as a gen rank.

_commit_megatron_bulk_refit_piece( piece: nemo_rl.models.generation.megatron.megatron_worker._MegatronBulkRefitPiece, tensor: torch.Tensor, ) → None#: Commit one received HF-local weight to its inference destination.

_build_destination_hf_to_local_param_map( refit_info: dict[str, Any], tasks: list[nemo_rl.models.generation.megatron.megatron_worker._MegatronRefitTask], ) → nemo_rl.weight_sync.nccl_reshard_utils.HFToLocalParamMap#: Map canonical HF FFN shards onto local Megatron destinations.

_prepare_mxfp8_refit( tasks: list[nemo_rl.models.generation.megatron.megatron_worker._MegatronRefitTask], ) → None#: Install persistent MCore MXFP8 destinations before engine initialization.

_refresh_flashinfer_mxfp8_weights() → None#: Refresh derived FlashInfer expert storage after canonical weights change.

_build_generation_refit_tasks() → tuple[list[torch.nn.Module], list[nemo_rl.models.generation.megatron.megatron_worker._MegatronRefitTask]]#: Resolve Bridge import tasks against the persistent inference model.

_install_generation_refit_plan( *, model_chunks: list[torch.nn.Module], tasks: list[nemo_rl.models.generation.megatron.megatron_worker._MegatronRefitTask], state_dict_info: dict[str, Any], ) → None#: Install the packed-broadcast subset of a Bridge receive plan.

_prepare_destination_refit_info( state_dict_info: dict[str, Any], ) → None#: Build the local HF-to-Megatron receive plan before the first refit.

_prepare_destination_nccl_reshard_refit_info( refit_info: dict[str, Any], ) → None#: Prepare Megatron as the destination of NCCL M-to-N reshard refit.

_destination_nccl_reshard_refit( refit_timeout_s: Optional[float] = None, ) → bool#: Receive bulk FFN shards via M-to-N, then import packed misc weights.

_write_generation_refit_weight( task: nemo_rl.models.generation.megatron.megatron_worker._MegatronRefitTask, converted_weight: torch.Tensor, ) → None#

Copy one converted weight into its persistent inference destination.

copy_ rather than rebinding: an MXFP8 destination quantizes in place and keeps its storage address, which is what already-captured CUDA graphs point at.

Parameters:

task – The Bridge import task owning the destination tensor.
converted_weight – The assembled weight in the destination’s layout.

Raises:

ValueError – If the converted weight does not match the destination shape, which means the refit plan and the model disagree.

_load_generation_refit_batch( weights: list[tuple[str, torch.Tensor]], ) → None#

Import every Bridge task whose HF inputs have now all arrived.

Called once per packed_broadcast_consumer batch. Weights arrive in HF order, but a task may fuse several of them (gate/up, stacked experts), so each tensor is buffered until its task’s dependency count reaches zero, then the task is converted and written in task order. Buffered tensors record their producing stream so the conversion can wait on it before reading across streams.

Parameters:: weights – (hf_name, tensor) pairs unpacked from this batch.

_update_destination_weights_from_collective( refit_timeout_s: Optional[float] = None, *, refresh_mxfp8: bool = True, ) → bool#: Receive packed HF tensors and import them into the Megatron model.

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

nemo_rl.models.generation.megatron.megatron_worker#

Module Contents#

Classes#

Functions#

API#

`nemo_rl.models.generation.megatron.megatron_worker`#