nemo_rl.models.automodel.config

Configuration classes for automodel-based training in NeMo RL.

Module Contents

Classes

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

API

class nemo_rl.models.automodel.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.

model_class: type

None

model_config: Any

None

hf_config_overrides: dict[str, Any]

None

allow_flash_attn_args: bool

None

attn_impl: Optional[str]

None

dtype: torch.dtype

None

enable_seq_packing: bool

None

max_grad_norm: float

None

cpu_offload: bool

None

offload_optimizer_for_logprob: bool

None

is_generation_colocated: Optional[bool]

None

is_reward_model: bool

None

class nemo_rl.models.automodel.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.

model: torch.nn.Module

None

model_state_dict_keys: list[str]

None

optimizer: Optional[torch.optim.Optimizer]

None

scheduler: Optional[Any]

None

is_hf_model: bool

None

is_moe_model: bool

None

is_reward_model: bool

None

model_class: type

None

model_config: Any

None

peft_config: Optional[nemo_automodel.components._peft.lora.PeftConfig]

None

autocast_enabled: bool

None