`nemo_rl.models.automodel.draft.checkpoint`#

Draft checkpoint save/load: optimizer-layout record and the DCP sibling entry.

Module Contents#

Functions#

`optimizer_layout_record`	Record named param-group layout for resume validation.
`draft_meta_record`	Versioned per-algo draft checkpoint metadata.
`draft_checkpoint_dir`	The draft’s DCP directory: a SIBLING of the policy weights directory.
`save_draft_checkpoint`	Save the draft’s sharded weights (DCP) plus the dspark metadata record.
`validate_dspark_checkpoint_meta`	Hard-error on inconsistent dspark checkpoint metadata.
`load_draft_checkpoint`	Load the draft’s weights from a checkpoint, validating the metadata record.
`_draft_meta_from_config`	`draft_meta_record` fed directly from a validated `policy.draft` config.
`load_checkpoint_with_draft`	Load a checkpoint with the composite-optimizer pairing rule.
`save_checkpoint_with_draft`	Save a checkpoint with the composite-optimizer pairing rule.

Data#

`DRAFT_CHECKPOINT_DIRNAME`
`DSPARK_META_FILENAME`
`DSPARK_OPTIMIZER_GROUP_NAMES`
`DRAFT_META_VERSION`

API#

nemo_rl.models.automodel.draft.checkpoint.DRAFT_CHECKPOINT_DIRNAME#: ‘draft’

nemo_rl.models.automodel.draft.checkpoint.DSPARK_META_FILENAME#: ‘dspark_meta.json’

nemo_rl.models.automodel.draft.checkpoint.DSPARK_OPTIMIZER_GROUP_NAMES#: (‘policy’, ‘draft’)

nemo_rl.models.automodel.draft.checkpoint.optimizer_layout_record( optimizer: torch.optim.Optimizer, ) → list[dict[str, Any]]#

Record named param-group layout for resume validation.

A DSpark co-training optimizer must consist of exactly the named groups [“policy”, “draft”] in that order; anything else means the optimizer was not built by the dspark setup path and its state cannot be paired safely.

nemo_rl.models.automodel.draft.checkpoint.DRAFT_META_VERSION#: 1

nemo_rl.models.automodel.draft.checkpoint.draft_meta_record( draft_model: torch.nn.Module, model_name: str, optimizer: Optional[torch.optim.Optimizer], algo: str = 'dspark', train_embed_and_head: Optional[bool] = None, ttt_steps: Optional[int] = None, ) → dict[str, Any]#

Versioned per-algo draft checkpoint metadata.

Legacy dspark checkpoints predate versioning (no meta_version/algo); the validator treats those as dspark/v0. ttt_steps is part of the eagle3 training contract (the unroll depth the drafter was trained for) and is required there; the block-drafter algos ignore it.

nemo_rl.models.automodel.draft.checkpoint.draft_checkpoint_dir(weights_path: str) → str#

The draft’s DCP directory: a SIBLING of the policy weights directory.

The draft must live outside the policy weight tree because detect_checkpoint_format(weights_path) walks it recursively and would mis-detect the safetensors policy checkpoint as DCP after seeing the draft’s .distcp files.

nemo_rl.models.automodel.draft.checkpoint.save_draft_checkpoint( draft_model: torch.nn.Module, weights_path: str, meta: dict[str, Any], ) → None#: Save the draft’s sharded weights (DCP) plus the dspark metadata record.

nemo_rl.models.automodel.draft.checkpoint.validate_dspark_checkpoint_meta( saved_meta: dict[str, Any], expected_meta: dict[str, Any], ) → None#

Hard-error on inconsistent dspark checkpoint metadata.

Architecture fields are always compared. The optimizer layout is compared only when the current run restores optimizer state (expected layout is non-None): weights-only loads (init_optimizer=False, e.g. eval or logprob-only workers) legitimately carry no layout while training checkpoints do.

nemo_rl.models.automodel.draft.checkpoint.load_draft_checkpoint( draft_model: torch.nn.Module, weights_path: str, expected_meta: dict[str, Any], ) → None#: Load the draft’s weights from a checkpoint, validating the metadata record.

nemo_rl.models.automodel.draft.checkpoint._draft_meta_from_config( draft_model: torch.nn.Module, draft_config: Any, optimizer: Optional[torch.optim.Optimizer], ) → dict[str, Any]#

draft_meta_record fed directly from a validated policy.draft config.

Both Eagle3DraftConfig and the block-drafter configs carry model_name/speculator_type/train_embed_and_head; only Eagle3DraftConfig carries ttt_steps.

nemo_rl.models.automodel.draft.checkpoint.load_checkpoint_with_draft( checkpoint_manager: Any, model: torch.nn.Module, draft_model: torch.nn.Module, composite_model: torch.nn.Module, optimizer: Optional[torch.optim.Optimizer], scheduler: Any, weights_path: str, optimizer_path: Optional[str], draft_config: Any, ) → None#

Load a checkpoint with the composite-optimizer pairing rule.

Policy weights load exactly as without a draft; draft weights load from the sibling draft entry with validated metadata; optimizer state pairs with the composite module whose param groups span policy + draft. This is the single home of that pairing invariant for both resume-at-setup and mid-run checkpoint loads. Works for eagle3 as well as dspark/dflash – all three validate the same policy.draft config shape.

nemo_rl.models.automodel.draft.checkpoint.save_checkpoint_with_draft( checkpoint_manager: Any, model: torch.nn.Module, draft_model: torch.nn.Module, composite_model: torch.nn.Module, optimizer: Optional[torch.optim.Optimizer], scheduler: Any, weights_path: str, optimizer_path: Optional[str], tokenizer: Any, tokenizer_path: Optional[str], is_final_checkpoint: bool, peft_config: Any, draft_config: Any, ) → None#

Save a checkpoint with the composite-optimizer pairing rule.

Mirrors load_checkpoint_with_draft: policy weights save without an optimizer, the optimizer (spanning policy + draft param groups) pairs with composite_model, then the draft’s own DCP entry and metadata record save last.

nemo_rl.models.automodel.draft.checkpoint#

Module Contents#

Functions#

Data#

API#

`nemo_rl.models.automodel.draft.checkpoint`#