# Resumable Training

[![View Source on AtomGit](https://mindspore-website.obs.cn-north-4.myhuaweicloud.com/website-images/r2.10.0/resource/_static/logo_source_en.svg)](https://atomgit.com/mindspore/docs/blob/r2.10.0/docs/mindformers/docs/source_en/feature/resume_training.md)

MindSpore Transformers dynamic graph (PyNative) supports **step-level resumable training**. After training is interrupted, the framework restores the model and optimizer status from the saved weight directory, rolls back the training progress to the step when the training was interrupted, and jumps the dataset cursor to the corresponding position. This avoids the waste of computing power caused by retraining from scratch.

Resumable training does not require an additional training entry. It reuses the same `checkpoint` configuration as the initial training. The only difference lies in the settings of several switches on the **loading side**: `load_path` points to the root directory of the previous saving, and `no_load_optim` specifies whether to restore the optimizer and training progress. This page is organized in the sequence of mechanism → configuration items → scenario-specific YAML files.

> Resumable training depends on the fact that the weight and optimizer state have been flushed to disk before the interruption. Ensure that the previous training meets the following conditions:
>- `enable_save` is set to `True` (default). Otherwise, no checkpoint will be generated during training.
>- `no_save_optim` is set to `False` (default). If the optimizer state is skipped during saving, the optimizer cannot be restored during resumable training. In this case, the training can only be retrained with weights as initialization.

## 1 Resumable Training Mechanism

### 1.1 Loading Entry and Triggering Conditions

When training starts, `Trainer` reads `checkpoint.load_path` (or the explicitly passed `checkpoint_path`) as the loading directory. **Loading is triggered as long as the value is not empty.** Whether to restore the optimizer and training progress is further determined by `no_load_optim`.

| `no_load_optim` | Loading Content        | Training Progress                              | Data Cursor  | Typical Use Case             |
|-----------------|--------------|------------------------------------|--------|-------------------|
| `False`         | Model weights + optimizer state| Resume from `global_step` of `common.json`.| Go to the resumption step.| Continue the same training after interruption.       |
| `True`          | Model weights only       | Start from 0.                            | Start from scratch.  | Retrain/Fine-tune with existing weights for initialization. |

### 1.2 Locating Specific Weights for `load_path`

`load_path` directs to the **root directory for saving**, not a subdirectory of a specific step. The framework locates the directory to be loaded in `get_checkpoint_path` based on the following rules:

1. Verify that the `load_path` exists and is a directory.
2. **Self-training weight locating**: Read the tracker file `latest_checkpointed_iteration.txt` in the root directory, obtain the latest iteration number, combine it with the `iteration_XXXXXXXX` (8 digits, padded with zeros) subdirectory, and load it. This means that during the resumption, **you do not need to manually specify the number of steps. The framework automatically selects the latest saved iteration.**

By default, the training checkpoints of MindSpore Transformers are stored in the `output/checkpoint` directory. Each checkpoint is saved as a subfolder named `iteration`. The following uses the checkpoint generated by an 8-device task as an example. The storage format is as follows:

```text
output
    ├── checkpoint
        ├── latest_checkpointed_iteration.txt               # tracker: records the latest iteration number. Assume that the iteration number is 2000.
        ├── iteration_00001000/                             # Weight directory of step 1000.
            ├── metadata.json                               # Weight sharding information.
            ├── common.json                                 # Metadata for resuming training.
            ├── {prefix}-model-0000000-0000008.safetensor   # Weight shard of device 0.
            ...
            ├── {prefix}-model-0000007-0000008.safetensor   # Weight shard of device 7.
            ├── {prefix}-opt-0000000-0000008.safetensor     # Optimizer shard of device 0.
            ...
            ├── {prefix}-opt-0000007-0000008.safetensor     # Optimizer shard of device 7.
        ...
        └── iteration_00002000/                             # Weight directory of step 2000 (latest one, and all weight files have been saved).
```

Set `load_path` to `output/checkpoint` (root directory). The framework automatically selects the weight under `iteration_00002000` based on the content in `latest_checkpointed_iteration.txt` for loading.

### 1.3 `common.json` Metadata and Resumable Training Closure

Each iteration subdirectory contains a `common.json` file, which records the metadata required for resumable training. The configuration items are defined in `CommonInfo`.

| Configuration Item                | Description                       |
|---------------------|---------------------------|
| `global_step`       | Global training step (**used to resume training**).      |
| `global_batch_size` | Global batch size during saving (**used for step scaling**).|
| `step_num`          | Step number in the current epoch.            |
| `epoch_num`         | Number of trained epochs.             |
| `loss_scale`        | Gradient scaling coefficient.                   |
| `ckpt_status`       | Checkpoint health status flag.        |

During resumable training (`no_load_optim=False`), `Trainer` reads `global_step` from `common.json`, writes it to the training state `state.global_step`, and calls `train_dataset.set_init_step(global_step)` to jump the dataset cursor to the resumption step, ensuring that data is **not repeated or missing**.

### 1.4 Step Scaling When the Batch Size Changes

If `global_batch_size` configured during resumable training is different from that saved, using the original value of `global_step` will result in inconsistent consumed training data volume. Therefore, `Trainer` performs proportional scaling based on **`global_batch_size` saved in `common.json`**.

```text
Resumption step = Saving step x (global_batch_size during saving/Current global_batch_size)
```

> **Scaling Examples**
>
> During saving, `global_batch_size` is `64` and `global_step` is `1000` (that is, 64,000 samples have been consumed). If `global_batch_size=128` is used during resumable training:
>
> ```text
> Resumption step = 1000 x (64/128) = 500
> ```
>
> That is, the training resumes from step 500. The corresponding accumulated dataset consumption is 500 x 128 = 64,000, which is the same as that when the training is interrupted.

### 1.5 Master Weight Alignment When Only Weights Are Loaded

When `no_load_optim` is set to `True`, the optimizer state is not loaded, but the FP32 master weight in hybrid precision still retains the initial value before loading. To avoid misalignment between master weights and the newly loaded model parameters, the framework calls `optimizer.reload_main_params_from_model()` to update the FP32 master weight using the model parameters, ensuring that both start from the same initial point.

## 2 Configuration Items for Resumable Training

The configuration items involved in resumable training are defined in the `checkpoint` section by `CheckpointConfig`. The following table provides a quick reference, followed by a detailed explanation of each item.

| Configuration Item                     | Default Value     | Active Side| Description                         |
|--------------------------|----------|-----|-----------------------------|
| `enable_save`            | `True`   | Saving | Specifies whether to save checkpoints.            |
| `save_path`              | `""`     | Saving | Directory where weights are stored.                  |
| `save_interleaved_steps` | `1000`   | Saving | Number of steps between each saving.                 |
| `no_save_optim`          | `False`  | Saving | Specifies whether to skip saving the optimizer state.                |
| `save_max`               | `5`      | Saving | Maximum number of checkpoints that can be saved after the current task starts.|
| `load_path`              | `""`     | Loading | Root directory for saving the loaded data for resumable training.                 |
| `no_load_optim`          | `False`  | Loading | Specifies whether to skip loading the optimizer state (whether to resume training or initialize only weights).   |
| `load_balanced`          | `False`  | Loading | Specifies whether to enable balanced loading (used together with parallelism, see section 2.3).    |

### 2.1 Configuration Items on the Saving Side (Determining Whether the Training Can Be Resumed)

When resuming training for the current task, ensure that these configurations are enabled during the previous training and that the weights are completely saved to the disk. The following table describes the configuration items on the saving side.

| Parameter                    | Data Type| Required/Optional| Default Value    | Description                                                                                                 |
|--------------------------|------|------|---------|-------------------------------------------------------------------------------------------------------|
| `enable_save`            | Boolean| Optional  | `True`  | Specifies whether to enable weight saving. After this function is disabled, data is not flushed to disks during the entire training process. Disable this function only during debugging and when interruption recovery is not required.                                                               |
| `save_path`              | String | Required  | `""`    | Root directory for saving data. Each time data is saved, the `iteration_XXXXXXXX` subdirectory and a tracker file are generated in the root directory. Generally, set `load_path` to the same directory during resumable training.                     |
| `save_interleaved_steps` | int  | Optional  | `1000`  | Number of steps between each saving. Data is flushed to disks only when `global_step % save_interleaved_steps == 0` is used. A smaller number of steps indicates a higher resistance to interruption risks, but higher I/O overhead. For large-scale pre-training, this parameter is usually set to hundreds to thousands of steps.|
| `no_save_optim`          | Boolean| Optional  | `False` | Specifies whether to skip the optimizer state during saving. If the value is `True`, only model weights are saved. The default value is `False`, which is used to ensure that the optimizer can be restored and the training can be resumed completely. Set this parameter to `True` only when you are sure that no further optimization is required (for example, you only want to export the weights).             |
| `save_max`               | int  | Optional  | `5`     | Maximum number of checkpoints that can be retained for the current training job. If the number exceeds the maximum, the earliest checkpoints are automatically deleted. Set this parameter based on the disk capacity and rollback requirements.                                                   |

### 2.2 Configuration Items on the Loading Side (Determining the Resumable Training Behavior)

The following table describes the configuration items on the loading side.

| Parameter            | Data Type | Required/Optional| Default Value     | Description                                                                                                                                                                                                                                              |
|------------------|-------|------|----------|----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|
| `load_path`      | String  | Required  | `""`     | Root directory for saving the loaded data for resumable training. If the value is not empty, loading is triggered. The framework automatically selects the latest iteration through the tracker (see 1.2). The value is that of `save_path` of the previous training (self-trained weight directory, which must contain `common.json`/`metadata.json`).                                                                                                                         |
| `no_load_optim`  | Boolean | Optional  | `False`  | Specifies whether to skip the optimizer state during loading. It is the main switch for resumable training and weight-only initialization (see 1.1). For resumable training, set this parameter to `False` (resuming the optimizer, number of steps, and data cursor). For retraining or fine-tuning with existing weights, set this parameter to `True` (using only the weights for initialization). If `no_load_optim` is set to `False`, the optimizer object must exist. Otherwise, an error will be reported.                                                                                             |
| `load_balanced`  | Boolean | Optional  | `False`  | After this parameter is enabled, the sharding balancing strategy (`apply_balance_shard_strategy`) is used to compute the redundant parameter mapping between ranks during the loading phase, and then `single_parameter_broadcast` is used to broadcast the parameters between ranks. This eliminates the repeated loading of redundant parameters and reduces the graphics memory and I/O overhead during distributed loading. This parameter is valid only in distributed parallel training. There is no benefit in single-device or non-parallel scenarios. Retain the default value `False`. For details about the meaning and configuration of the parallel dimension, see [Distributed Parallel Training](./parallel_training.md).|

## 3 Differences Between Resumable Training and Weight Loading for Training

The following table compares the configuration differences between weight loading for training and resumable training, and provides the **default values** to help distinguish between the default values of configuration items and the actual configuration.

| Project              | Weight Loading for Training                           | Resumable Training                              | Default Value                             |
|------------------|-----------------------------------|------------------------------------|----------------------------------|
| `load_path`      | Empty (random initialization) or the pre-trained weight directory.              | Root directory of the previous saving.                       | `""` (If the value is not empty, loading is performed.)                |
| `no_load_optim`  | Explicitly set to `True` (not requiring the optimizer information when weights are loaded for pre-training).| Must be `False` (restoring the optimizer).                | `False`                |
| Starting from `global_step`| Starts from 0.                           | Resumes from `global_step` of `common.json`.| `0`                        |
| Dataset cursor           | Starts from scratch.                             | Goes to the resumption step.                            | `set_init_step`                  |
| FP32 master weight        | Initialized with the network.                           | Obtained from model parameters when `no_load_optim` is set to `True`.    | `reload_main_params_from_model`  |

> **Do not be misled by the default value of no_load_optim.**
>
> The default value of `no_load_optim` in the source code is `False`. This is correct for resumable training (the optimizer is restored by default), but not suitable for initial pre-training (there is usually no optimizer state to load). The common practice is to **not set `load_path`** (no loading) or to **explicitly set `no_load_optim` to `True`** when initializing with pre-trained weights.
>
> That is, the default value of `no_load_optim` will enter the resumable training process. If you are loading weights for training, you need to adjust this configuration based on the scenario.

## 4 Scenario-Based Configuration Examples

The following provides complete `checkpoint` section in YAML files that can be directly used in three typical scenarios. For details about other sections (such as model, dataset, and parallelism), see [Configuration File Description](./configuration.md).

### Scenario A: Standard Resumable Training (Restoring the Optimizer + Number of Steps + Data Cursor)

Most common scenario: The training is interrupted unexpectedly and continues from the last saved point. If `no_load_optim` is set to `False`, the framework automatically selects the latest iteration, restores the `global_step`, and jumps to the data cursor.

```yaml
checkpoint:
  # Save the configuration (use the previous training and ensure that the data has been flushed to the disk before the interruption).
  enable_save: True
  save_path: "./output/ckpt"
  save_interleaved_steps: 1000
  no_save_optim: False
  save_max: 5

  # Load the configuration for resumable training.
  load_path: "./output/ckpt"   # Use the root directory for saving. The framework automatically obtains the latest iteration.
  no_load_optim: False         # Restore the optimizer state, training step, and data cursor.
  load_balanced: False         # Retain False in single-node or unbalanced scenarios.
```

### Scenario B: Weight-Only Initialization Retraining/Fine-Tuning

If you want to reuse the existing weights as the starting point (for example, using pre-trained weights for fine-tuning) without retaining the original optimizer state and training progress, you can set `no_load_optim` to `True`. The training will restart from step 0 and trigger FP32 master weight alignment.

```yaml
checkpoint:
  enable_save: True
  save_path: "./output/finetune_ckpt"
  save_interleaved_steps: 500
  no_save_optim: False
  save_max: 5

  # Load only weights for initialization.
  load_path: "./pretrained/ckpt"  # Root directory of the self-trained weight (the files in the specific iteration include common.json and metadata.json).
  no_load_optim: True             # Do not restore the optimizer, training steps, or data cursor. Start from step 0.
  load_balanced: False
```

> Currently, the dynamic graph does not support direct loading of the Hugging Face weight directory.

### Scenario C: Resuming Training with an Adjusted `global_batch_size`

The global batch size is adjusted during resuming training (for example, by increasing or decreasing the number of nodes). The configuration remains unchanged. The framework automatically scales the restoration step based on the formula in section 1.4. You do not need to manually modify `global_step`.

```yaml
# In training section: Change global_batch_size from 64 to 128 during resumable training.
training:
  global_batch_size: 128

checkpoint:
  enable_save: True
  save_path: "./output/ckpt"
  save_interleaved_steps: 1000
  no_save_optim: False
  save_max: 5

  load_path: "./output/ckpt"
  no_load_optim: False    # The value must be False. The scaling logic takes effect only when the optimizer is restored.
  load_balanced: False
```

According to the test case, if `global_batch_size=64` and `global_step=1000` are saved, after the value is changed to `128` for resumable training, the training continues from step `1000 × 64 / 128 = 500`, and the accumulated dataset consumption remains 64,000.

> Scaling takes effect only in the resumable training path.
> Step scaling occurs only when `no_load_optim` is set to `False`. If `no_load_optim` is set to `True` (that is, the optimizer weight is not loaded), the training starts from 0, and the dataset is not scaled.

### Scenario D: Balanced Loading in Distributed Parallel Training

When multiple devices are used for distributed parallel resumable training, enabling `load_balanced` can eliminate the repeated loading of redundant parameters, reducing the graphics memory usage and I/O overhead during loading.

```yaml
checkpoint:
  enable_save: True
  save_path: "./output/ckpt"
  save_interleaved_steps: 1000
  no_save_optim: False
  save_max: 5

  load_path: "./output/ckpt"
  no_load_optim: False
  load_balanced: True      # This parameter is valid only in parallel scenarios: shard balancing + parameter broadcast.
```

For details about the configuration of the parallel dimension, see [Distributed Parallel Training](./parallel_training.md).

## Related Documents

- Weight format and saving mechanism: [Weight Saving and Loading](./save_load_checkpoint.md)
- Configuration item overview: [Configuration File Description](./configuration.md)
- `load_balanced` and parallel dimension: [Distributed Parallel Training](./parallel_training.md)
- Data cursor restoration and `set_init_step`: [Datasets](./dataset.md)
- End-to-end training process: [Training Guide](../guide/training.md)
- Quick start: [Quick Start](../quick_start/quick_start.md)