> This page is for version 0.7.0 · 26.10. > For other versions, use one of these documentation indexes: > - Nightly (default): https://docs.nvidia.com/nemo/automodel/nightly/llms.txt > - Latest: https://docs.nvidia.com/nemo/automodel/latest/llms.txt > - 0.7.0 · 26.10: https://docs.nvidia.com/nemo/automodel/v0.7/llms.txt > - 0.5.0 · 26.06: https://docs.nvidia.com/nemo/automodel/v0.5/llms.txt > - 0.4.0 · 26.04: https://docs.nvidia.com/nemo/automodel/v0.4/llms.txt > For clean Markdown of any page, append .md to the page URL. > For a complete documentation index, see https://docs.nvidia.com/nemo/automodel/llms.txt. > For AI client integration (Claude Code, Cursor, etc.), connect to the MCP server at https://docs.nvidia.com/nemo/automodel/_mcp/server. # Checkpointing in NeMo AutoModel > Save and manage NeMo AutoModel checkpoints in Safetensors and PyTorch DCP formats, including consolidated and sharded layouts, PEFT adapters, asynchronous saves, and preemption handling. ## Introduction During machine learning experiments, the model training routine regularly saves checkpoints. A checkpoint is a complete snapshot of a run that includes model weights, optimizer states, and other metadata required to resume training exactly where it left off. Writing these snapshots at regular intervals lets you recover quickly from crashes or pauses without losing progress. NeMo AutoModel checkpoints capture the complete state of a distributed training run across multiple GPUs or nodes. This approach reduces memory overhead, improves GPU utilization, and allows you to resume training with a different parallelism strategy. NeMo AutoModel writes checkpoints in two formats: [Hugging Face Safetensors](https://github.com/safetensors/safetensors) and [PyTorch Distributed Checkpointing (DCP)](https://docs.pytorch.org/docs/stable/distributed.checkpoint.html). It also supports two layouts: * **Consolidated Checkpoints**: The complete model state is saved as a Hugging Face-compatible bundle, typically in a single file or a compact set of files with an index. Because tensors are not split across GPUs (unsharded), tools such as Hugging Face, vLLM, and SGLang can load these checkpoints directly. * **Sharded Checkpoints**: During distributed training with parameter sharing, each GPU holds a subset (or "shard") of the full state, such as model weights and optimizer states. When checkpointing, each GPU writes its own shard independently without reconstructing the full model state. The following table provides an overview of the available checkpoint formats. | Task | Model Domain | DCP (Sharded) | Safetensors (Sharded) | Safetensors (Consolidated) | | ---- | ------------ | :-----------: | :-------------------: | :------------------------: | | SFT | LLM | ✅ | ✅ | ✅ | | SFT | VLM | ✅ | ✅ | ✅ | | PEFT | LLM / VLM | 🚧 | 🚧 | ✅ | Change output formats through the recipe's YAML configuration file: ```yaml checkpoint: ... model_save_format: safetensors # Format for saving (torch_save or safetensors) save_consolidated: final # Recommended: export consolidated HF weights only for the final checkpoint. # Other modes: false (sharded only) or every/true (export every checkpoint). consolidation_timeout_minutes: 30 # Timeout for inline consolidated-export synchronization. ... ``` `save_consolidated` accepts the following values: * `final` (recommended): Keep intermediate checkpoints sharded and export consolidated Hugging Face weights only for the final checkpoint. * `false`: Save sharded checkpoints only. Run the generated `model/consolidate.sh` helper later if you need Hugging Face weights. * `every` (or legacy `true`): Export consolidated Hugging Face weights during every checkpoint save. Use this only when every checkpoint must be immediately loadable by Hugging Face tools. NeMo AutoModel writes a `model/consolidate.sh` helper next to Safetensors model shards. Use this helper to create a Hugging Face-compatible `model/consolidated/` directory after training for `save_consolidated: false` checkpoints, or for earlier checkpoints when using `save_consolidated: final`. Creating consolidated Hugging Face weights requires `model_save_format: safetensors`. Inline consolidated export synchronizes ranks through a dedicated Gloo process group. `consolidation_timeout_minutes` controls that synchronization independently of `dist_env.timeout_minutes`, which continues to control training collectives. The default is 30 minutes. > **Note** > > The optimizer states are always saved in DCP format with the `.distcp` extension. ## Checkpoint Symbolic Links NeMo AutoModel automatically creates symbolic links in the checkpoint directory to provide convenient access to important checkpoints: * **LATEST**: Points to the most recently saved checkpoint. This is useful for resuming training from the last saved state. * **LOWEST_VAL**: Points to the checkpoint with the lowest validation score or loss. This provides easy access to the best-performing checkpoint based on validation metrics, making it ideal for model evaluation or deployment. These symbolic links eliminate the need to manually track checkpoint names or search through directories to find the best model. When validation is enabled in your training run, both links are automatically maintained and updated as training progresses. ## Interrupted Saves A save can be cut short by a wall-clock limit, a preemption, or an out-of-memory error, which leaves a partially written `epoch__step_` directory behind. NeMo AutoModel marks each checkpoint directory as in progress while it is being written and clears the mark only after every component is saved and the checkpoint is published. Directories left by an interrupted save stay marked, which determines how they are treated: * **Resume skips them.** A checkpoint becomes resumable only after its marker is cleared. Both the `LATEST` pointer lookup and the step-order fallback ignore marked directories. Training resumes from the most recent complete checkpoint. Setting `checkpoint.restore_from` directly to a marked directory, or to a different pointer such as `LOWEST_VAL` whose target is marked, raises an error instead of loading an incomplete checkpoint. * **The next save of that step replaces them.** Checkpoint directory names are derived from the step, so the resumed run targets the same directory when it reaches that step again. Because the marked directory holds a checkpoint that was never published, NeMo AutoModel removes it, logs a warning, and writes the new checkpoint in its place. * **Retention does not preserve them.** When `max_recent_checkpoints` is set, a marked directory never occupies a slot in the recent window, so it is pruned as an orphan. > **Warning** > > Only marked directories are replaced. A published checkpoint is never overwritten: if you resume from an older checkpoint with `restore_from` and training runs past a step that already has a published checkpoint, the save fails with `FileExistsError` rather than discarding it. Remove that directory, or point `checkpoint.checkpoint_dir` at a fresh directory, before rerunning the step. > **Note** > > Checkpoint saves are synchronous by default and are published before training continues. With opt-in asynchronous checkpointing (`checkpoint.is_async: true`), the next save call or a clean training shutdown waits for the background write and publishes the checkpoint. If the job terminates before either point, the pending checkpoint stays marked as in progress (even if its file writes finished), and resume uses the previous published checkpoint. > **Warning** > > `msc://` paths use NVIDIA MultiStorageClient (MSC) object storage. NeMo AutoModel supports them for sharded Distributed Checkpoint (DCP) model and optimizer data when `checkpoint.save_consolidated` is `false`. However, a full recipe checkpoint is not self-contained in object storage. RNG, dataloader, and other recipe state written with `torch.save`, along with the `LATEST` and `LOWEST_VAL` pointers, remain local to the training environment. The interrupted-save marker and automatic stale-directory replacement described above also require a local filesystem and are unavailable for `msc://` roots. > > Use a local or shared filesystem checkpoint directory when you need a portable, exact training resume. If you use an `msc://` checkpoint directory for DCP data, configure it as follows: > > ```yaml > checkpoint: > checkpoint_dir: msc://bucket/path > save_consolidated: false > max_recent_checkpoints: null > ``` ## Checkpoint Retention By default, NeMo AutoModel keeps all intermediate checkpoints. To limit disk usage when checkpoints are saved frequently, configure a bounded recent-checkpoint window: ```yaml checkpoint: max_recent_checkpoints: 2 ``` `max_recent_checkpoints` sets the number of checkpoint directories in the recent window, ordered by step. It is not a hard limit on the total number of checkpoint directories. Checkpoint-root pointers protect their targets in addition to the recent window. These pointers include `LATEST`, `LOWEST_VAL`, other top-level symbolic links, and `.txt` fallback pointer files. This behavior supports reliable resumption from the latest checkpoint and preserves explicitly referenced checkpoints, such as the best validation checkpoint. For example, set `max_recent_checkpoints: 2` to retain the two most recent checkpoints. If `LOWEST_VAL` points to an older checkpoint, NeMo AutoModel also retains the `LOWEST_VAL` target. Bounded retention removes only standard `epoch__step_` checkpoint directories. It does not remove recipe-specific exports, such as retrieval distillation `step_` directories. Leave the value unset, or set it to `null`, to keep all intermediate checkpoints: ```yaml checkpoint: max_recent_checkpoints: null ``` Bounded retention supports only local filesystem checkpoint directories. If `checkpoint.checkpoint_dir` uses `msc://` storage, leave `max_recent_checkpoints` unset. ## Safetensors To integrate with the Hugging Face ecosystem, NeMo AutoModel saves checkpoints in the [Safetensors](https://github.com/safetensors/safetensors) format. Safetensors is a memory-safe, zero-copy alternative to Python's pickle (PyTorch `.bin`). Hugging Face Transformers supports Safetensors natively, and Safetensors offers safety and performance advantages over Python pickle-based approaches. ### Key Benefits * **Native Hugging Face Compatibility**: Checkpoints can be loaded directly into Hugging Face-compatible tools, including vLLM, SGLang, and others. * **Memory Safety and Speed**: The Safetensors format prohibits saving serialized Python code, ensuring memory safety, and supports zero-copy loading for improved performance. * **Optional Consolidation**: Sharded checkpoints can be merged into a standard Hugging Face model format for easier downstream use. This format also supports optional consolidation of multiple shards into a complete Hugging Face format model. ### Example The following command runs the LLM fine-tuning recipe on 2 GPUs and saves the resulting checkpoint in the Safetensors format: ```bash automodel --nproc-per-node=2 examples/llm_finetune/llama3_2/llama3_2_1b_squad.yaml \ --step_scheduler.ckpt_every_steps 20 \ --checkpoint.model_save_format safetensors \ --checkpoint.save_consolidated final ``` > **Note** > > The preceding command uses the [`llama3_2_1b_squad.yaml`](https://github.com/NVIDIA-NeMo/Automodel/blob/492add84a2b9d495946fe211c28973cd00051f3e/examples/llm_finetune/llama3_2/llama3_2_1b_squad.yaml) configuration file as a running example. Adjust it for your use case. > More configuration examples are available in the [`examples/`](https://github.com/NVIDIA-NeMo/Automodel/tree/r0.7.0/examples) directory. If you are running on a single GPU, run the following command: ```bash automodel examples/llm_finetune/llama3_2/llama3_2_1b_squad.yaml \ --step_scheduler.ckpt_every_steps 20 \ --checkpoint.model_save_format safetensors \ --checkpoint.save_consolidated final ``` After running for a few seconds, the standard output should be: ``` ... > Saving checkpoint to checkpoints/epoch_0_step_20 ... ``` The `checkpoints/` directory should have the following contents: ``` checkpoints/ ├── LATEST -> epoch_0_step_20 ├── LOWEST_VAL -> epoch_0_step_20 └── epoch_0_step_20 ├── model │ ├── consolidate.sh │ ├── shard-00001-model-00001-of-00001.safetensors │ └── shard-00002-model-00001-of-00001.safetensors └── optim ├── __0_0.distcp └── __1_0.distcp ... ``` The `epoch_0_step_20/` directory stores the full training state from step `20` of the first epoch, including both the model and optimizer states. Because this example uses `save_consolidated: final`, intermediate checkpoints such as `epoch_0_step_20/` do not include `model/consolidated/` before the run reaches the final checkpoint. To export this intermediate checkpoint for Hugging Face-compatible tools, run the generated helper: ```bash bash checkpoints/epoch_0_step_20/model/consolidate.sh ``` Run the helper from the AutoModel repo root so it can find `tools/offline_hf_consolidation.py`, or set `CONSOLIDATION_TOOL=/path/to/tools/offline_hf_consolidation.py`. The helper defaults to one CPU worker process with five writer threads so it is safe on small machines. For large checkpoints, run it on a CPU compute node and increase parallelism: ```bash NPROC_PER_NODE=16 NUM_THREADS=5 bash checkpoints/epoch_0_step_20/model/consolidate.sh ``` `NPROC_PER_NODE` controls worker processes, and `NUM_THREADS` controls writer threads per process. Keep `NPROC_PER_NODE * NUM_THREADS` within your CPU allocation. You can also submit the helper to a CPU Slurm partition, as shown in the following example: ```bash sbatch --cpus-per-task=80 --wrap='NPROC_PER_NODE=16 NUM_THREADS=5 bash /path/to/checkpoints/epoch_0_step_20/model/consolidate.sh' ``` By default, consolidated export uses the original Hugging Face Safetensors headers when they are available. Ordinary floating-point tensors are restored to their original per-tensor Hugging Face dtype, such as BF16, FP16, or FP32, even if the saved sharded checkpoint uses a different floating dtype. If the run started from configuration-only weights or the original Hugging Face metadata is unavailable, export keeps the saved checkpoint dtype. If an original quantized or packed tensor was saved as a floating-point tensor, export leaves it as float and emits a warning. You can request an explicit floating-point dtype cast during offline export: ```bash CAST_DTYPE=bf16 bash checkpoints/epoch_0_step_20/model/consolidate.sh ``` Use `CAST_DTYPE` when the consolidated Hugging Face bundle should override the default per-tensor dtype behavior, such as `CAST_DTYPE=bf16` to export ordinary floating-point tensors as BF16 for serving. Supported values include `bf16`, `fp16`, `fp32`, and `fp64`. Only ordinary floating-point tensors with a different source dtype are cast. Tensors already in the cast dtype, FP8 tensors, and non-floating tensors are left unchanged. The helper writes `checkpoints/epoch_0_step_20/model/consolidated/`. Load and run that consolidated checkpoint directly with the Hugging Face Transformers API: ```python import torch from transformers import pipeline model_id = "checkpoints/epoch_0_step_20/model/consolidated/" pipe = pipeline( "text-generation", model=model_id, torch_dtype=torch.bfloat16, device_map="auto", ) print(pipe("The key to life is")) >>> [{'generated_text': 'The key to life is to be happy. The key to happiness is to be kind. The key to kindness is to be'}] ``` Although this example uses the Hugging Face Transformers API, the `consolidated/` checkpoint is compatible with any Hugging Face-compatible tool, such as vLLM, SGLang, and others. ## PEFT For Parameter-Efficient Fine-Tuning (PEFT), NeMo AutoModel saves the trained adapters instead of the full model, often reducing checkpoints to just a few megabytes. Reloading combines those adapters with the original base model, so adapter-only saving assumes the base model's weights and buffers stay unchanged. PEFT checkpoints save adapter files directly under `model/` and do not generate or need `model/consolidate.sh`. ### Router Correction Biases in MoE Models Some Mixture-of-Experts (MoE) routers update a correction-bias buffer during training to balance which experts process tokens. LoRA freezes parameters but does not stop these buffer updates. Adapter-only checkpoints omit the buffer, so reloading can change expert selection and resumed training. For models using the shared gate with a pretrained `e_score_correction_bias`, such as DeepSeek V3/V3.2 and GLM MoE DSA (including GLM-5.1), add these settings to the existing `model` section of the PEFT recipe: ```yaml model: moe_overrides: gate_bias_update_factor: 0.0 force_e_score_correction_bias: true ``` The zero update factor freezes adaptation; `force_e_score_correction_bias` retains the pretrained buffer and its effect on routing. Do not remove or zero the buffer. These settings apply to this correction-bias mechanism, not ordinary router bias parameters. Intentional router-bias adaptation requires additional save, restore, and export support, tracked in [issue #4075](https://github.com/NVIDIA-NeMo/Automodel/issues/4075). ### Benefits of Consolidated Adapter Checkpoints Because the PEFT state is so lightweight, sharded checkpointing adds unnecessary overhead. Instead, NeMo AutoModel automatically saves a compact Hugging Face-compatible adapter checkpoint when using PEFT. This makes it: * Easier to manage and share (just the adapters). * Compatible with Hugging Face Transformers out of the box. * Ideal for deployment and downstream evaluation. ### Run PEFT Fine-Tuning on 2 GPUs To fine-tune a model using PEFT and save a Hugging Face-compatible checkpoint: ```bash automodel --nproc-per-node=2 examples/llm_finetune/llama3_2/llama3_2_1b_hellaswag_peft.yaml --step_scheduler.ckpt_every_steps 20 --checkpoint.model_save_format safetensors ``` After training, the run produces a compact Safetensors adapter checkpoint that can be loaded directly with Hugging Face tools: ``` checkpoints/ ├── LATEST -> epoch_0_step_20 ├── LOWEST_VAL -> epoch_0_step_20 ├── epoch_0_step_20 │ ├── config.yaml │ ├── dataloader │ │ ├── dataloader_dp_rank_0.pt │ │ └── dataloader_dp_rank_1.pt │ ├── losses.json │ ├── model │ │ ├── adapter_config.json │ │ ├── adapter_model.safetensors │ │ ├── automodel_peft_config.json │ │ ├── special_tokens_map.json │ │ ├── tokenizer.json │ │ └── tokenizer_config.json │ ├── optim │ │ ├── __0_0.distcp │ │ └── __1_0.distcp │ ├── rng │ │ ├── rng_dp_rank_0.pt │ │ └── rng_dp_rank_1.pt │ └── step_scheduler.pt ├── training.jsonl └── validation.jsonl ``` The following example shows the direct compatibility of NeMo AutoModel with Hugging Face and PEFT: ```python from peft import AutoPeftModelForCausalLM from transformers import AutoTokenizer checkpoint_path = "checkpoints/epoch_0_step_20/model/" model = AutoPeftModelForCausalLM.from_pretrained(checkpoint_path) tokenizer = AutoTokenizer.from_pretrained(checkpoint_path) model = model.to("cuda") model.eval() inputs = tokenizer("Preheat the oven to 350 degrees and place the cookie dough", return_tensors="pt") outputs = model.generate(input_ids=inputs["input_ids"].to("cuda"), max_new_tokens=50) print(tokenizer.batch_decode(outputs.detach().cpu().numpy(), skip_special_tokens=True)[0]) >>> Preheat the oven to 350 degrees and place the cookie dough in a large bowl. Roll the dough into 1-inch balls and place them on a cookie sheet. Bake the cookies for 10 minutes. While the cookies are baking, melt the chocolate chips in the microwave for 30 seconds. ``` ## PyTorch DCP NeMo AutoModel also offers native PyTorch DCP checkpointing support (`.distcp` extension). Similar to Safetensors, it provides load-time resharding and parallel saving. As a simple example, run the following command to launch the training recipe on 2 GPUs. ```bash automodel --nproc-per-node=2 examples/llm_finetune/llama3_2/llama3_2_1b_squad.yaml \ --step_scheduler.ckpt_every_steps 20 \ --checkpoint.model_save_format torch_save ... > Saving checkpoint to checkpoints/epoch_0_step_20 ... ``` After 20 steps, AutoModel saves the following checkpoint: ``` checkpoints/ ├── LATEST -> epoch_0_step_20 ├── LOWEST_VAL -> epoch_0_step_20 └── epoch_0_step_20 ├── config.yaml ├── dataloader │ ├── dataloader_dp_rank_0.pt │ └── dataloader_dp_rank_1.pt ├── losses.json ├── model │ ├── __0_0.distcp │ └── __1_0.distcp └── optim ├── __0_0.distcp └── __1_0.distcp ... ``` If you rerun the script, NeMo AutoModel automatically detects and restores the most recent checkpoint. ```bash automodel --nproc-per-node=2 examples/llm_finetune/llama3_2/llama3_2_1b_squad.yaml \ --step_scheduler.ckpt_every_steps 20 \ --checkpoint.model_save_format torch_save ... > Loading checkpoint from checkpoints/epoch_0_step_20 ... ``` ## Save Checkpoints When Using Docker When training inside a Docker container (refer to the [Installation Guide](/nemo/automodel/v0.7/get-started/installation)), any files written to the container's filesystem are lost when the container exits, especially with `--rm`. To keep your checkpoints, bind-mount a host directory to the checkpoint path before starting the container: ```bash docker run --gpus all -it --rm \ --shm-size=8g \ -v "$(pwd)"/checkpoints:/opt/Automodel/checkpoints \ nvcr.io/nvidia/nemo-automodel:26.10.00 ``` You can also set a custom checkpoint directory through the YAML config or CLI override: ```yaml checkpoint: checkpoint_dir: /mnt/shared/my_checkpoints ``` ```bash # Or with a CLI override: automodel examples/llm_finetune/llama3_2/llama3_2_1b_squad.yaml \ --checkpoint.checkpoint_dir /mnt/shared/my_checkpoints ``` When using a custom path, make sure the corresponding host directory is mounted into the container with `-v`. > **Tip** > > Mount additional host directories for datasets and the Hugging Face model cache to avoid downloading large models again across container restarts. Refer to the [Installation Guide](/nemo/automodel/v0.7/get-started/installation) for a complete `docker run` example with all recommended mounts. ## Enable Asynchronous Checkpointing Checkpointing is synchronous by default. NeMo AutoModel can instead write checkpoints asynchronously to reduce training stalls caused by I/O. When enabled, checkpoint writes are scheduled in the background using PyTorch Distributed Checkpointing's async API while training continues. * **Enable** (YAML): ```yaml checkpoint: is_async: true ``` * **Enable** (CLI): Add `--checkpoint.is_async True` to your run command. * **Requirements**: PyTorch ≥ 2.9.0. If an older version is detected, async mode is automatically disabled. * **Behavior**: At most one checkpoint is written at a time. The next save waits for the previous write to finish. * **Publication**: The next save call waits for and publishes the previous asynchronous checkpoint. On clean shutdown, each recipe waits for and publishes the final checkpoint before closing the checkpointer. A forced termination before either point can leave a finished background write unpublished, as described in [Interrupted Saves](#interrupted-saves). * **PEFT**: Adapter model files are written synchronously on rank 0. Optimizer states can still be written asynchronously. ## Preemption Checkpointing On shared clusters, jobs can be preempted or reach their wall-time limits before training finishes. NeMo AutoModel can catch an operating system signal sent ahead of the shutdown, save a checkpoint at the next step boundary, and exit gracefully. This lets training resume from the moment of preemption instead of the last periodic checkpoint. ```yaml step_scheduler: ... preemption_signal: SIGUSR1 # default: SIGTERM ``` `preemption_signal` accepts the following values: * A signal name as a string, case-insensitive, with or without the `SIG` prefix (`"SIGUSR1"`, `"usr1"`, and `"SIGTERM"` are all valid). * A signal number, such as `15`. * A list of preceding options to watch multiple signals at once, such as `["SIGUSR1", "SIGTERM"]`. * `null` to disable preemption handling entirely. By default, `SIGTERM` is watched, so runs already checkpoint on a graceful termination request without any configuration. ### How It Works When any rank receives the configured signal, the signal handler records it locally. The flags are then gathered across ranks at the next step boundary, so every rank observes the preemption at the same step. All ranks write a checkpoint together and shut down cleanly. This means it is sufficient for the signal to reach a single rank. For example, the signal can reach only the workers on one node. Since the checkpoint is written at a step boundary, the delay between signal delivery and the checkpoint is at most one training step. Set the signal lead time to cover one step and the time to write a checkpoint for your model size and storage backend. If the signal arrives while the final steps of a run are completing, the run finishes normally and writes its regular final checkpoint, so no separate preemption checkpoint is made. ### Slurm Usage The reference `slurm.sub` submission script demonstrates the full setup. Slurm delivers the warning signal to the batch script through the `--signal` directive: ```bash #SBATCH --signal=B:USR1@300 # send USR1 to the batch script 300s before the time limit ``` The `B:` prefix is required: it delivers the signal to the batch script only. The batch script then forwards `SIGUSR1` to the local training worker processes. The workers must be signaled directly rather than `srun` or `torchrun`, because `torchrun` terminates on an unhandled `USR1`, taking the job down before a checkpoint can be written. Forwarding to the workers on the batch node alone is sufficient because the all-gather propagates the flag to every rank on every node. To test preemption behavior without waiting for a time limit, deliver the signal manually to a running job's batch script: ```bash scancel --batch --signal=USR1 ``` > **Note** > > `--signal=B:USR1@` guarantees the warning before a *wall-time* termination. Whether the same warning is delivered on a *scheduler preemption*, such as a higher-priority job taking the nodes, depends on the cluster's Slurm configuration (`GraceTime` on the partition and `PreemptParameters=send_user_signal`). Ask your cluster administrator if preemption-triggered signals are required. ### Resume Training After Preemption A preemption checkpoint is a regular checkpoint: the `LATEST` symbolic link is updated to point to it. Resubmitting the same job resumes from the preempted step with model, optimizer, and scheduler state intact. > **Note** > > Signals that arrive before training starts (during model loading or dataloader > construction) cannot trigger a checkpoint, since the signal handler is > installed when the step scheduler is created and there is no training state to > save at that point. ## Save Additional States You can also save additional states in NeMo AutoModel. By default, AutoModel automatically checkpoints the `dataloader`, `rng`, and `step_scheduler` states that are necessary to resume training accurately. A complete Safetensors consolidated checkpoint looks like this: ``` checkpoints/ ├── LATEST -> epoch_0_step_20 ├── LOWEST_VAL -> epoch_0_step_20 ├── epoch_0_step_20 │ ├── config.yaml │ ├── dataloader │ │ ├── dataloader_dp_rank_0.pt │ │ └── dataloader_dp_rank_1.pt │ ├── losses.json │ ├── model │ │ ├── consolidated │ │ │ ├── config.json │ │ │ ├── generation_config.json │ │ │ ├── model-00001-of-00001.safetensors │ │ │ ├── model.safetensors.index.json │ │ │ ├── special_tokens_map.json │ │ │ ├── tokenizer.json │ │ │ └── tokenizer_config.json │ │ ├── shard-00001-model-00001-of-00001.safetensors │ │ └── shard-00002-model-00001-of-00001.safetensors │ ├── optim │ │ ├── __0_0.distcp │ │ └── __1_0.distcp │ ├── rng │ │ ├── rng_dp_rank_0.pt │ │ └── rng_dp_rank_1.pt │ └── step_scheduler.pt ├── training.jsonl └── validation.jsonl ``` To define a new state to be checkpointed in the recipe, create a new attribute in the recipe class (defined using `self.` inside the recipe). Ensure that the new attribute uses both the `load_state_dict` and `state_dict` methods. The following example shows how to define the state: ```python class NewState: def __init__(self, ...): self.state_value = ... self.another_value = ... ... def state_dict(self) -> dict[str, Any]: return { "": self.state_value, "": self.another_value, } def load_state_dict(self, state_dict: dict[str, Any]) -> None: self.state_value = state_dict[""] self.another_value = state_dict[""] ``` Inside your recipe class, define the new state as an instance attribute using `self.new_state = NewState(...)`.