Resumable Training

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:

Verify that the load_path exists and is a directory.
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:

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.

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:
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.

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.

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.

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.

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.

# 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.

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.