nemo_rl.models.automodel.checkpoint#

Automodel checkpoint utilities for the policy workers.

This module provides a wrapper class around the nemo_automodel Checkpointer for saving and loading model checkpoints in Automodel policy workers.

Module Contents#

Classes#

AutomodelCheckpointManager

Manages checkpointing for Automodel models using nemo_automodel’s Checkpointer.

Functions#

_resolve_lora_adapter_dir

Resolve a lora_cfg.restore_from path to the adapter directory.

build_checkpoint_config

Build the Automodel checkpoint config for a worker.

_patch_qwen_vl_vision_key_mapping

Re-add the Qwen2.5-VL ^visual -> model.visual checkpoint key rename.

Data#

logger

API#

nemo_rl.models.automodel.checkpoint.logger#

‘getLogger(…)’

nemo_rl.models.automodel.checkpoint._resolve_lora_adapter_dir(restore_from: str) → str#

Resolve a lora_cfg.restore_from path to the adapter directory.

Accepts a NeMo RL checkpoint weights directory (step_*/policy/weights), its model subdirectory, or any directory directly containing adapter_model.safetensors + adapter_config.json (HF PEFT layout).

nemo_rl.models.automodel.checkpoint.build_checkpoint_config(
automodel_cfg: collections.abc.Mapping[str, Any],
*,
model_repo_id: str,
dequantize_base_checkpoint: bool,
is_peft: bool,
is_async: bool,
skip_task_head_prefixes_for_base_model: Optional[list[str]] = None,
) → dict[str, Any]#

Build the Automodel checkpoint config for a worker.

Read NeMo-RL settings from the YAML-backed checkpoint block and forward them to Automodel. This integration boundary is shared by the two Automodel v2 workers; defaults belong in the exemplar configs.

Parameters:

  • automodel_cfg – The worker’s policy.automodel_cfg / value.automodel_cfg mapping. Automodel checkpoint settings are read from its nested checkpoint block; all other keys are ignored.

  • model_repo_id – Forwarded to Automodel’s CheckpointingConfig.model_repo_id.

  • dequantize_base_checkpoint – Forwarded to CheckpointingConfig.dequantize_base_checkpoint; only takes effect on a base-checkpoint init load, not on a resume load.

  • is_peft – Forwarded to CheckpointingConfig.is_peft.

  • is_async – Forwarded to CheckpointingConfig.is_async.

  • skip_task_head_prefixes_for_base_model – Optional parameter-prefix list forwarded to CheckpointingConfig.skip_task_head_prefixes_for_base_model (only set on the returned dict when not None); the value worker passes ["score."] to skip the reward head on base-model init loads.

Returns:

A dict of only the fields Automodel’s CheckpointingConfig dataclass accepts, meant to be splatted into AutomodelCheckpointingConfig(enabled=True, checkpoint_dir="", **result).

nemo_rl.models.automodel.checkpoint._patch_qwen_vl_vision_key_mapping() → None#

Re-add the Qwen2.5-VL ^visual -> model.visual checkpoint key rename.

Workaround for a transformers v5.5.0 regression. transformers #44627 moved VLM checkpoint conversions into the main mapping, but copied the Qwen-VL visual key mapping incorrectly: visual.* checkpoint keys no longer map to model.visual.*. transformers #45358 fixed those VLM mappings in v5.6, but the Automodel commit NeMo-RL can currently pin to still depends on transformers v5.5.0. Automodel’s get_combined_key_mapping mirrors the transformers WeightRenaming entries, so the bad v5.5.0 mapping leaves vision-tower checkpoint keys unmapped and FSDP2 set_model_state_dict(strict=False) drops them in load_base_model. The vision tower is then left randomly initialized, making the training forward diverge from vLLM (token_mult_prob_error).

This wraps get_combined_key_mapping to inject the missing rule for qwen2_5_vl/qwen2_vl. It is idempotent: the rule is only added when no existing rule already targets model.visual. Remove this after Automodel upgrades its transformers dependency to a version that includes #45358.

class nemo_rl.models.automodel.checkpoint.AutomodelCheckpointManager(
dp_mesh: torch.distributed.device_mesh.DeviceMesh,
tp_mesh: torch.distributed.device_mesh.DeviceMesh,
moe_mesh: Optional[torch.distributed.device_mesh.DeviceMesh] = None,
)#

Manages checkpointing for Automodel models using nemo_automodel’s Checkpointer.

This class provides a clean interface for saving and loading model checkpoints, wrapping the nemo_automodel Checkpointer with configuration management.

.. attribute:: checkpointer

The underlying nemo_automodel Checkpointer instance.

Initialization

Initialize the AutomodelCheckpointManager.

Parameters:

  • dp_mesh – The data parallel device mesh.

  • tp_mesh – The tensor parallel device mesh.

  • moe_mesh – Optional MoE device mesh.

_get_dp_rank() → int#

Get the data parallel rank.

_get_tp_rank() → int#

Get the tensor parallel rank.

init_checkpointer(
config_updates: Optional[dict[str, Any]] = None,
) → None#

Initialize the Automodel Checkpointer if not already created.

This method creates a new Checkpointer instance with the provided configuration. If a checkpointer already exists, this method does nothing.

Parameters:

config_updates – Automodel checkpoint fields to set during initialization.

finalize_async_save() → None#

Block until in-flight async checkpoint writes have landed on disk.

With is_async=True the Automodel Checkpointer hands both the model and optimizer state to dcp.async_save, which stages them and uploads from a separate process. Those writes address files by path, so the caller must not rename tmp_step_N to step_N until they finish – otherwise the writer re-creates tmp_step_N and the promoted checkpoint is missing its optimizer shards and .metadata.

Safe to call when async saving is off or no save is in flight; both underlying calls are no-ops in that case.

static _save_tokenizer_on_rank0(
tokenizer: Any,
tokenizer_path: str,
) → None#

Save a tokenizer (or processor) to tokenizer_path from rank 0 only.

Unlike model/optimizer state, the tokenizer is replicated rather than sharded: every rank holds an identical copy, and save_pretrained writes to the same rank-independent filenames (tokenizer_config.json, …). Letting all ranks write means N-1 redundant writes racing on the same file. save_pretrained is not atomic (it opens with O_TRUNC, writes, and may read files back), so concurrent writers can deadlock on the inode lock. On a hard-mounted NFS share this manifests as ranks stuck indefinitely in uninterruptible disk sleep, hanging the whole job at checkpoint time.

This mirrors the rank-0 guard nemo_automodel applies to the same artifacts in nemo_automodel.components.checkpoint.addons.ConsolidatedHFAddon.pre_save.

Note this must NOT be applied to dcp.save-based model/optimizer saves: those are collective calls whose output is already rank-disjoint, so guarding them would deadlock and drop every non-zero rank’s shard.

Parameters:

  • tokenizer – The tokenizer or processor to save.

  • tokenizer_path – Directory to save the tokenizer into.

save_checkpoint(
model: torch.nn.Module,
weights_path: str,
optimizer: Optional[torch.optim.Optimizer] = None,
optimizer_path: Optional[str] = None,
scheduler: Optional[torch.optim.lr_scheduler.LRScheduler] = None,
tokenizer: Optional[transformers.AutoTokenizer] = None,
tokenizer_path: Optional[str] = None,
*,
is_final_checkpoint: bool,
peft_config: Optional[nemo_automodel.components._peft.lora.PeftConfig] = None,
) → None#

Save a checkpoint of the model.

The optimizer states are saved only if optimizer and optimizer_path are provided. Any previous async save is completed before a new one starts. When async saving is enabled, this method returns after model and optimizer staging is complete; upload and consolidation may continue in the background.

Parameters:

  • model – The model to save.

  • weights_path – Path to save model weights.

  • optimizer – Optional optimizer to save.

  • optimizer_path – Optional path to save optimizer state.

  • scheduler – Optional learning rate scheduler.

  • tokenizer – Optional tokenizer to save with the checkpoint.

  • tokenizer_path – Optional path to save tokenizer separately.

  • is_final_checkpoint – Whether this checkpoint completes the training run, either at the configured final step or after a deliberate early stop. Automodel’s save_consolidated="final" mode consolidates these checkpoints. Timeout recovery checkpoints are resumable and are not considered final.

  • peft_config – Optional PEFT configuration.

load_lora_adapter(model: torch.nn.Module, adapter_dir: str) → None#

Load validated donor adapter weights through the PEFT checkpointer path.

Parameters:

  • model – Model whose adapters will be initialized.

  • adapter_dir – Adapter directory or checkpoint weights directory.

load_checkpoint(
model: torch.nn.Module,
weights_path: str,
optimizer: Optional[torch.optim.Optimizer] = None,
optimizer_path: Optional[str] = None,
scheduler: Optional[torch.optim.lr_scheduler.LRScheduler] = None,
) → None#

Load a checkpoint into the model using Automodel Checkpointer.

Parameters:

  • model – The model to load weights into.

  • weights_path – Path to the checkpoint weights.

  • optimizer – Optional optimizer to load state into.

  • optimizer_path – Optional path to optimizer checkpoint.

  • scheduler – Optional learning rate scheduler.

check_master_weight_dtypes(
model: torch.nn.Module,
optimizer: torch.optim.Optimizer,
optimizer_path: str,
) → None#

Hard-error if a resumed optimizer’s on-disk master dtype has drifted.

Must run before checkpointer.load_optimizer: legacy FP32 masters must not be silently cast into BF16 int16 remainders (or vice versa) just because store_param_remainders/master_weight_dtype changed between save and resume.