nemo_rl.models.automodel.setup

Setup utilities for automodel-based training in NeMo RL.

Module Contents

Functions

_has_optimizer_fp32_master	Whether the resolved optimizer and kwargs provide FP32 master weights.
_maybe_set_force_hf	Validate and maybe auto-set force_hf based on adapter compatibility.
get_tokenizer	Get tokenizer using NeMoAutoTokenizer for automodel workers.
validate_and_prepare_config	Validate configuration and prepare runtime settings.
setup_reference_model_state	Set up reference model state dict by creating a CPU copy of the model’s state dict.
setup_distributed	Set up distributed training environment and create device meshes.
_validate_lora_adapter_config	Fail closed if the donor adapter’s metadata is incompatible with this run.
_validate_lora_adapter_keys	Fail closed if the donor adapter tensors don’t cover this model’s LoRA params.
_load_initial_lora_adapter	Warm-start the model’s LoRA adapters from a donor PEFT adapter checkpoint.
setup_model_and_optimizer	Set up model, parallelization, and optimizer.

Data

STRING_TO_DTYPE

API

nemo_rl.models.automodel.setup.STRING_TO_DTYPE

None

nemo_rl.models.automodel.setup._has_optimizer_fp32_master(

optimizer_cls: type[torch.optim.Optimizer],

optimizer_kwargs: dict[str, Any],

) → bool

Whether the resolved optimizer and kwargs provide FP32 master weights.

nemo_rl.models.automodel.setup._maybe_set_force_hf(automodel_kwargs: dict, model_config) → None

Validate and maybe auto-set force_hf based on adapter compatibility.

Custom model implementations (e.g. Qwen2, Llama) use state_dict_adapters to convert between native and HF weight formats. NeMo RL’s weight syncing requires the adapter to implement convert_single_tensor_to_hf. Some adapters (like CombinedProjectionStateDictAdapter) don’t implement this yet.

This function checks the adapter BEFORE model loading to avoid wasting time:

• force_hf=True: no check needed, HF model won’t have an adapter.

• force_hf not set + adapter incompatible: auto-set force_hf=True with a warning.

• force_hf=False + adapter incompatible: raise an error telling the user to set force_hf=True or file an issue to NeMo-Automodel.

See: https://github.com/NVIDIA-NeMo/RL/issues/2072

nemo_rl.models.automodel.setup.get_tokenizer(

tokenizer_config: nemo_rl.models.policy.TokenizerConfig,

get_processor: bool = False,

) → Union[transformers.PreTrainedTokenizerBase, transformers.AutoProcessor]

Get tokenizer using NeMoAutoTokenizer for automodel workers.

Uses NeMoAutoTokenizer which provides custom tokenizer dispatch per model type and falls back to NeMoAutoTokenizerWithBosEosEnforced for default handling.

Parameters:

• tokenizer_config – A dictionary containing tokenizer configuration. Required keys: - name: The name or path of the pretrained tokenizer Optional keys: - chat_template: The chat template to use. Can be: - None: Uses a passthrough template that just returns message content - “default”: Uses the tokenizer’s default template - A file path ending in “.jinja”: Loads template from file - A custom jinja2 template string If not specified, the tokenizer’s default template will be used. - chat_template_kwargs: Arguments passed to tokenizer.apply_chat_template() - tokenizer_kwargs: Extra keyword arguments forwarded to NeMoAutoTokenizer.from_pretrained(), e.g. {“fix_mistral_regex”: False}. When get_processor=True, these are passed through AutoProcessor.from_pretrained().

• get_processor – Whether to return a processor (via AutoProcessor) instead of a tokenizer.

Returns:

The configured tokenizer or processor instance.

nemo_rl.models.automodel.setup.validate_and_prepare_config(

config: nemo_rl.models.policy.PolicyConfig,

processor: Optional[transformers.AutoProcessor],

rank: int,

) → nemo_rl.models.automodel.config.RuntimeConfig

Validate configuration and prepare runtime settings.

This function validates the policy configuration, sets environment variables, determines model configuration, and returns runtime settings as a named tuple.

Parameters:

• config – Policy configuration dictionary

• processor – Optional processor for multimodal models

• rank – Current process rank

Returns:

RuntimeConfig named tuple containing validated configuration values

Raises:

• ValueError – If configuration is invalid

• RuntimeError – If incompatible settings are detected

nemo_rl.models.automodel.setup.setup_reference_model_state(

model: torch.nn.Module,

) → dict[str, torch.Tensor]

Set up reference model state dict by creating a CPU copy of the model’s state dict.

This creates a reference copy of the model weights on CPU with pinned memory for efficient CPU-GPU transfers. The reference model is typically used to compute reference log probabilities during RL training.

Parameters:

model – The model to create a reference copy from

Returns:

Dictionary mapping parameter names to CPU tensors with pinned memory

.. rubric:: Example

model = setup_model(…) reference_model_state_dict = setup_reference_model_state(model)

nemo_rl.models.automodel.setup.setup_distributed(

config: nemo_rl.models.policy.PolicyConfig,

runtime_config: nemo_rl.models.automodel.config.RuntimeConfig,

) → nemo_rl.models.automodel.config.DistributedContext

Set up distributed training environment and create device meshes.

Initializes torch.distributed process group and creates FSDP2Config, MoEParallelizerConfig, and device meshes for distributed training.

Parameters:

• config – Policy configuration dictionary

• runtime_config – RuntimeConfig named tuple from validate_and_prepare_config

Returns:

DistributedContext containing device meshes and distributed configuration

nemo_rl.models.automodel.setup._validate_lora_adapter_config(

adapter_dir: str,

lora_cfg: nemo_rl.models.policy.LoRAConfig,

model_name: str,

) → None

Fail closed if the donor adapter’s metadata is incompatible with this run.

Checks the standard HF PEFT adapter_config.json fields: the adapter must be a LoRA adapter with the same rank (r) and scaling (lora_alpha) as this run’s lora_cfg, and must have been trained on the same base model.

nemo_rl.models.automodel.setup._validate_lora_adapter_keys(

adapter_dir: str,

model: torch.nn.Module,

*,

moe_mesh: Any = None,

) → None

Fail closed if the donor adapter tensors don’t cover this model’s LoRA params.

The underlying PEFT load is unconditionally non-strict (a key mismatch is only a warning), so this check is the only guarantee that the donor covers every LoRA parameter. Expected keys are converted through the same native to HF state-dict adapter path used when saving, so the comparison also works for custom model implementations and expert-parallel models.

nemo_rl.models.automodel.setup._load_initial_lora_adapter(

model: torch.nn.Module,

checkpoint_manager: nemo_rl.models.automodel.checkpoint.AutomodelCheckpointManager,

restore_from: str,

lora_cfg: nemo_rl.models.policy.LoRAConfig,

model_name: str,

) → None

Warm-start the model’s LoRA adapters from a donor PEFT adapter checkpoint.

The load goes through the Automodel checkpointer’s PEFT path (each rank reads adapter_model.safetensors, then set_model_state_dict(broadcast_from_rank0=True) places it with DTensor/EP awareness), the same machinery used to resume NeMo RL PEFT checkpoints. Optimizer state is not loaded: warm starts begin with a fresh optimizer.

nemo_rl.models.automodel.setup.setup_model_and_optimizer(

config: nemo_rl.models.policy.PolicyConfig,

tokenizer: transformers.AutoTokenizer,

runtime_config: nemo_rl.models.automodel.config.RuntimeConfig,

distributed_context: nemo_rl.models.automodel.config.DistributedContext,

checkpoint_manager: Any,

is_vlm: bool = False,

init_optimizer: bool = True,

weights_path: Optional[str] = None,

optimizer_path: Optional[str] = None,

) → nemo_rl.models.automodel.config.ModelAndOptimizerState

Set up model, parallelization, and optimizer.

Creates the model via from_pretrained() which handles meta device init, parallelization (FSDP2/TP/CP/EP), LoRA, and base weight loading internally.

Parameters:

• config – Policy configuration dictionary

• tokenizer – Tokenizer for the model

• runtime_config – RuntimeConfig named tuple from validate_and_prepare_config

• distributed_context – DistributedContext from setup_distributed

• checkpoint_manager – Checkpoint manager for loading/saving weights

• is_vlm – Whether this is a vision-language model

• init_optimizer – Whether to initialize optimizer

• weights_path – Optional path to checkpoint weights to load

• optimizer_path – Optional path to optimizer state to load

Returns:

ModelAndOptimizerState containing model, optimizer, scheduler, and metadata