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

Module Contents#

Classes#

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

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.

async_engine: bool#: None

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

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, 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]#: None

logging_step_interval: NotRequired[int]#: None

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

Bases: nemo_rl.models.generation.interfaces.GenerationConfig

Generation config for Megatron Inference.