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

Module Contents#

Classes#

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

Functions#

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

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

parsers: list[str]#: None

buffer_size_gb: int#: None

block_size_tokens: int#: None

max_tokens: 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

inference_cuda_graph_scope: NotRequired[str]#: None

materialize_only_last_token_logits: bool#: None

enable_chunked_prefill: bool#: None

enable_prefix_caching: bool#: None

refit_backend: Literal[gloo, nccl, nvshmem]#: None

num_speculative_tokens: int#: None

mamba_inference_ssm_states_dtype: NotRequired[str]#: None

mamba_inference_conv_states_dtype: 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.

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