`nemo_rl.models.generation.megatron.config`#

Module Contents#

Classes#

`MCoreGenerationSpecificArgs`	Megatron fields related only to inference.
`MCoreGenerationConfig`	Generation config for Megatron Inference.

Functions#

`resolve_refit_execution_batch_bytes`	Resolve native MCore staging bytes using NeMo-RL’s collective default.
`merged_inference_megatron_cfg`	The `megatron_cfg` a dedicated inference model runs with.
`dedicated_inference_megatron_cfg`	The `megatron_cfg` for a dedicated colocated inference model, or None.

API#

nemo_rl.models.generation.megatron.config.resolve_refit_execution_batch_bytes( configured_bytes: int | None, ) → int#: Resolve native MCore staging bytes using NeMo-RL’s collective default.

class nemo_rl.models.generation.megatron.config.MCoreGenerationSpecificArgs#

Bases: typing.TypedDict

Megatron fields related only to inference.

Any fields not declared here but declared in the training-side config can be overwritten. For example, Megatron inference might want transformer_impl: "inference_optimized", while Megatron training might want transformer_impl: "transformer_engine".

Initialization

Initialize self. See help(type(self)) for accurate signature.

expose_http_server: bool#: None

http_server_num_replicas: NotRequired[int]#: None

parsers: list[str]#: None

buffer_size_gb: int#: None

block_size_tokens: int#: None

max_tokens: int#: None

max_requests: NotRequired[int]#: None

max_model_len: int#: None

num_cuda_graphs: int | None#: None

use_cuda_graphs_for_non_decode_steps: bool#: None

cuda_graph_impl: str#: None

cuda_graph_sizing_distribution: NotRequired[Literal[exponential, linear, hybrid]]#: None

cuda_graph_max_tokens: NotRequired[int]#: None

inference_cuda_graph_scope: NotRequired[str]#: None

moe_pad_experts_for_cuda_graph_inference: NotRequired[bool]#: None

materialize_only_last_token_logits: bool#: None

enable_chunked_prefill: bool#: None

enable_prefix_caching: bool#: None

async_sched_mode: NotRequired[typing.Literal[legacy, async]]#: None

vision_embedding_cache_max_bytes: NotRequired[int]#: None

allow_stale_multimodal_embeddings: NotRequired[bool]#: None

refit_backend: NotRequired[Literal[gloo, nccl, nccl_m2n, nvshmem] | None]#: None

refit_execution_batch_bytes: int | None#: None

num_speculative_tokens: int#: None

mamba_inference_ssm_states_dtype: NotRequired[str]#: None

mamba_inference_conv_states_dtype: NotRequired[str]#: None

prefix_caching_mamba_gb: NotRequired[float | None]#: None

prefix_caching_eviction_policy: NotRequired[Literal[ref_zero, lru]]#: None

prefix_caching_coordinator_policy: NotRequired[Literal[load_balanced, longest_prefix, first_prefix_block]]#: None

prefix_cache_ttl_seconds: NotRequired[float]#: None

prefix_caching_routing_alpha: NotRequired[float]#: None

vision_model_type: NotRequired[str]#: None

image_dynamic_resolution: NotRequired[bool]#: None

video_num_frames: NotRequired[int]#: None

video_temporal_patch_size: NotRequired[int]#: None

video_target_num_patches: NotRequired[int]#: None

video_maintain_aspect_ratio: NotRequired[bool]#: None

megatron_inference_wrapper: NotRequired[str]#: None

kv_cache_management_mode: Literal[persist, offload, recompute]#: None

logging_step_interval: NotRequired[int]#: None

logprobs_mode: Literal[processed_logprobs, raw_logprobs]#: None

fp8_cfg: NotRequired[nemo_rl.models.policy.Fp8Config]#: None

class nemo_rl.models.generation.megatron.config.MCoreGenerationConfig#

Bases: nemo_rl.models.generation.interfaces.GenerationConfig

Generation config for Megatron Inference.

Initialization

Initialize self. See help(type(self)) for accurate signature.

refit_transport: NotRequired[Literal[mcore, nccl_reshard] | None]#: None

mcore_generation_config: nemo_rl.models.generation.megatron.config.MCoreGenerationSpecificArgs#: None

nemo_rl.models.generation.megatron.config.merged_inference_megatron_cfg( policy_config: nemo_rl.models.policy.PolicyConfig, ) → dict[str, Any]#: The megatron_cfg a dedicated inference model runs with.

nemo_rl.models.generation.megatron.config.dedicated_inference_megatron_cfg( policy_config: nemo_rl.models.policy.PolicyConfig, ) → Optional[dict[str, Any]]#

The megatron_cfg for a dedicated colocated inference model, or None.

Colocated Megatron generation shares the training model unless the resolved inference layout or transformer_impl differs from training; then the worker builds a second model and reshards into it on every wake. Inference never uses CP, and CP is already pinned to 1 (CP>1 training therefore always differs).

Returns None when the resolved config matches training (reshardless: generate directly on the shared training model).

nemo_rl.models.generation.megatron.config#

Module Contents#

Classes#

Functions#

API#

`nemo_rl.models.generation.megatron.config`#