nemo_rl.models.megatron.config

Module Contents

Classes

MegatronGenerationConfig
RuntimeConfig	Runtime configuration for model training and inference.
ModelAndOptimizerState	Container for model and optimizer state.

API

class nemo_rl.models.megatron.config.MegatronGenerationConfig

Bases: typing.TypedDict

buffer_size_gb: int

None

buffer_guaranteed_fraction: float

None

num_cuda_graphs: int

None

block_size_tokens: int

None

use_cuda_graphs_for_non_decode_steps: bool

None

enable_chunked_prefill: bool

None

unified_memory_level: int

None

max_tokens: int

None

class nemo_rl.models.megatron.config.RuntimeConfig

Bases: typing.NamedTuple

Runtime configuration for model training and inference.

This contains all validated runtime settings needed for model initialization, parallelization, and training.

megatron_cfg: megatron.bridge.training.config.ConfigContainer

None

model_cfg: Any

None

dtype: torch.dtype

None

optimizer_cpu_offload: bool

None

offload_optimizer_for_logprob: bool

None

is_generation_colocated: Optional[bool]

None

final_padded_vocab_size: int

None

class nemo_rl.models.megatron.config.ModelAndOptimizerState

Bases: typing.NamedTuple

Container for model and optimizer state.

This named tuple holds all model-related state including the model itself, optimizer, scheduler, and metadata about the model type and configuration.

state: megatron.bridge.training.state.GlobalState

None

model: megatron.core.transformer.MegatronModule

None

optimizer: megatron.core.optimizer.MegatronOptimizer

None

scheduler: megatron.core.optimizer_param_scheduler.OptimizerParamScheduler

None

checkpointing_context: dict[str, Any]

None

param_sync_func: Optional[Callable]

None