Configuring Recovery Acceleration

Configuring Periodic Checkpoint Saving

This section describes key steps for periodic checkpoint saving. For details on the features of periodic checkpoint saving, see Periodic Checkpoint Saving.

Configuring Storage Checkpoint Loading

Loading checkpnoints from storage can be performed using the loading interface provided by the AI framework. You need to pass the file path to be loaded into the AI framework. Taking the MindSpeed-LLM framework as an example, you can refer to the following example if you need to configure the storage checkpoint loading function.

In the job YAML, add the --load /data/ckpt/XXX \ parameter to enable storage checkpoint loading. --load is the unified switch for training process recovery; training process recovery takes effect only after this switch is turned on.

...
spec:
  replicaSpecs:
    Master:
      template:
        spec:
          containers:
          - name: ascend # Do not modify
            args:
              - |
                bash scripts/train_start.sh /job/code /job/output pretrain_gpt.py \
                  ...
                  --load /data/ckpt/XXX \  # Checkpoint storage path
                  ...
    Worker:
      template:
        spec:
          containers:
          - name: ascend # Do not modify
            ...
            args:
              - |
                ...
                bash scripts/train_start.sh /job/code /job/output pretrain_gpt.py \
                  ...
                --load /data/ckpt/XXX \    # Checkpoint storage path
                  ...
...

Configuring Dying Gasp Checkpoint Saving

This section provides key steps for dying gasp checkpoint saving. For details, see Dying Gasp Checkpoint Saving.

Building an Image

Use a Dockerfile to build a container image and add the startup command.

...
# Adaptation Script to MindCluster lossless resumable training
RUN pip3 install $TASKD_WHL
RUN pip3 install $MINDIO_TTP_PKG

# Optional. The following commands must be configured when using graceful fault tolerance, Pod-level rescheduling, or process-level rescheduling.
RUN sed -i '/import os/i import taskd.python.adaptor.patch' $(pip3 show torch | grep Location | awk -F ' ' '{print $2}')/torch/distributed/run.py

Preparing the Job YAML

In the training job YAML, add the following fields to enable process-level recovery. recover-strategy is the strategy used for training process recovery, where dump indicates dying gasp checkpoint saving. Under ports, add ttp-port 8000 and port 9601 for TaskD communication.

Saving dying gasp checkpoint can be used as a policy named dump of recover-strategy for process-level recovery. An example is shown below.

...
metadata:
   labels:
     ...
 ...
...
   annotations:
     ...
     recover-strategy: "dump"       # Ding gasp checkpoint saving
 ...

...
spec:
   replicaSpecs:
      Master:
         template:
            spec:
              containers:
                 env:
                   - name: TTP_PORT
                     value: "8000"
                 args: […]
                 ports:
                   - containerPort: 8000
                     name: ttp-port
                   - containerPort: 9601
                     name: taskd-port
     ...
     Worker:
        template:
          spec:
            containers:
               env:
                 - name: TTP_PORT
                   value: "8000"
               args: […]
               ports:
                 - containerPort: 8000
                   name: ttp-port
                 - containerPort: 9601
                   name: taskd-port
  ...

Adapting the Training Script

1. After the distributed environment initialization is complete and the global rank is obtained, modify the training script to launch TaskD Manager within the training script.

   1. Create a manager.py and save it to the directory where the training script is called. The content of the manager.py file is as follows.

      from taskd.api import init_taskd_manager, start_taskd_manager
      import os
      
      job_id=os.getenv("MINDX_TASK_ID")
      node_nums=XX          # Total number of nodes
      proc_per_node=XX     # Number of training processes per node
      
      init_taskd_manager({"job_id":job_id, "node_nums": node_nums, "proc_per_node": proc_per_node})
      start_taskd_manager()

      For detailed parameter descriptions in the manager.py file, see def init_taskd_manager(config:dict) -> bool:.

   2. Add the following code to the training script to start TaskD Manager.

      export TASKD_PROCESS_ENABLE="on"
      
      # Under the PyTorch framework
      if [[ "${RANK}" == 0 ]]; then
          export MASTER_ADDR=${POD_IP}
          python /job/code/manager.py 2>> /job/code/alllogs/$MINDX_TASK_ID/taskd/error.log &           # The specific execution path of manager.py is determined by the current path, and the error.log log path must be created in advance.
      fi
      
      torchrun ...

2. In the startup script (e.g., train_start.sh), add the --max_restarts parameter. An example is shown below.

    ...
       logger "server id is: ""${server_id}"
       if [ "${framework}" == "PyTorch" ]; then
         get_env_for_pytorch_multi_node_job
         DISTRIBUTED_ARGS="--nproc_per_node $GPUS_PER_NODE --nnodes $NNODES --node_rank $NODE_RANK --master_addr $MASTER_ADDR --master_port $MASTER_PORT  --max_restarts 32767"
     ...

   Here, --max_restarts specifies the maximum number of fault triggers allowed within the container, expressed as an integer. If this limit is exceeded, the PyTorch training process will exit immediately. If this parameter is not configured, the default value is 32767.

NOTE If the error "the libtaskd.so has not been loaded" occurs during training, you need to import the LD_PRELOAD environment variable in the training script. This environment variable allows the system to preload specified .so files. An example is shown below.

export LD_PRELOAD=/usr/local/Ascend/cann/lib64/libmspti.so:/usr/local/lib/python3.10/dist-packages/taskd/python/cython_api/libs/libtaskd.so

• libmspti.so: This .so file is provided by MindStudio and integrated in the CANN package. The default installation path is /usr/local/Ascend/cann/lib64/libmspti.so.

• libtaskd.so: This .so file is provided by TaskD. After the whl package is installed, the path is TaskD installation path/taskd/python/cython_api/libs/libtaskd.so. You can run the following command to query the path where TaskD is located. The Location field in the command output is the target path.

  pip show taskd

Restoring Parameter Passing on the Parameter Plane

Currently, this capability is only supported in the process-level rescheduling and process-level online recovery features. It is enabled by default after adapting according to the Configuring Process-Level Rescheduling and Configuring Process-Level Online Recovery features.

(Optional) Disabling Parameter Passing Recovery on the Parameter Plane

For the process-level rescheduling and process-level online recovery features, if you want to disable this function and load parameters from the storage checkpoint, you need to modify the job YAML file. The following is an example of using process-level rescheduling and disabling parameter passing recovery on the parameter plane.

...
metadata:
   labels:
     ...
     fault-scheduling: "grace"
 ...
...
   annotations:
     ...
     recover-strategy: "recover"   # Recovery strategy: process-level rescheduling
...
...
spec:
  replicaSpecs:
    Master:
      template:
        spec:
          containers:
          - name: ascend # do not modify
            ...
            args:
              - |
                ...
                bash scripts/train_start.sh /job/code /job/output pretrain_gpt.py \
                  ...
                  --distributed-optimizer-no-replica \
                  ...
    Worker:
      template:
        spec:
          containers:
          - name: ascend # Do not modify
            ...
            args:
              - |
                ...
                bash scripts/train_start.sh /job/code /job/output pretrain_gpt.py \
                  ...
                  --distributed-optimizer-no-replica \
                  ...
...

distributed-optimizer-no-replica indicates whether to support periodic checkpoints for data repair, which is disabled by default. After this function is enabled, the replica optimizer does not have replicas, reducing memory usage. In process-level rescheduling and process-level online recovery scenarios, periodic checkpoints are used for repair. This function must be enabled only when process-level rescheduling or process-level online recovery is enabled.

Optimizing the Integration Time

Recovery Time Optimization (PyTorch)

This section describes the related features that you can choose to shorten the resumable training time on the PyTorch framework, including Fault Detection Time, Collective Communication Initialization Time, Training Rollback and Checkpoint Loading Time, and Operator Compilation Time.

Fault Detection Time

A parameter plane network fault in a cluster may not affect a training job. Therefore, the cluster scheduling components do not forcibly interrupt the job. When the parameter plane network fault affects a training job, the network timeout mechanism of collective communication is triggered. After a default waiting period of 30 minutes, the cluster scheduling components can detect the fault and trigger resumable training. To solve this problem, the PyTorch Adapter plugin (torch_npu) provides a watchdog fault detection function to determine if training jobs are affected and to reduce fault detection time. For details, see Table 1.

Table 1 Watchdog fault detection

...
# env for breakpoint ckpt
export RESUME_MODE_ENABLE=1


export HCCL_ASYNC_ERROR_HANDLING=0            # For details about this environment variable, see TaskD Environment Variable Description

Collective Communication Initialization Time

Parallel Store multi-thread link setup optimization: When PyTorch creates communication groups, TCP Store is used for information exchange. As the job scale increases, the information processing performance of the native TCP Store degrades, leading to prolonged times for creating communication groups. To solve this problem, torch_npu supports the optimized Parallel Store built on the native TCP Store. For details, see Table 2.

Table 2 Parallel Store

torchrun train.py --train_parameter=xxx ....

torch_npu_run train.py --train_parameter=xxx ....

• Performance optimization of native HCCL link setup: PyTorch sets up a link between NPUs after the collective communication information is exchanged on the NPU. As the job scale increases, the link setup time increases significantly. To solve this problem, CANN is introduced to optimize the performance of the native HCCL link setup. For details, see Table 3.

  Table 3 Native HCCL link setup performance optimization

  Function Name	Native HCCL link setup performance optimization
  Feature Description	By asynchronously completing collective communication information negotiation, multiple threads reduce both the negotiation time and the overall link setup time.
  Instructions	Only CANN 8.0.RC2 and later versions are supported.
  Key Operations	None

• Link setup optimization in RankTable mode: Ascend Operator provides the function of generating a collective communication configuration file (RankTable file, also called hccl.json) for PyTorch. Links can be set up in RankTable mode to shorten cluster communication link setup time. For details,, see Table 4.

  Table 4 Link setup for collective communication in RankTable mode

  Function Name	Link setup in RankTable mode
  Feature Description	Uses Ascend Operator is used to generate the collective communication configuration file for PyTorch tasks, reducing the cluster communication link setup time.
  Instructions	The version of torch_npu must be higher than 6.0.RC3.
  Key Operations	The parent directory of the hccl.json file is already mounted by default in the startup YAML. You can change it as required. volumes: - name: ranktable-dir hostPath: path: /user/mindx-dl/ranktable # This host directory must be under a shared directory type: DirectoryOrCreate Run the following commands to create the specific mount path for the hccl.json file in the host directory and modify the owner. mkdir -m 777 /user/mindx-dl/ranktable/Task_Running_Namespace.Task_Name chown 9000:9000 /user/mindx-dl/ranktable/default.pytorch-test For example: mkdir -m 777 /user/mindx-dl/ranktable/default.pytorch-test chown 9000:9000 /user/mindx-dl/ranktable/default.pytorch-test Modify the training script and add the following environment variable. export RANK_TABLE_FILE=/user/mindx-dl/ranktable/hccl.json Modify the training YAML and add the following settings. yaml volumeMounts: - name: ranktable mountPath: /user/mindx-dl/ranktable volumes: - name: ranktable hostPath: path: /user/mindx-dl/ranktable/namespace_of_the_running_task.task_name # Actual path of the hccl.json file in the host directory

Training Rollback and Checkpoint Loading Time

• Asynchronous checkpoint saving: A training job periodically saves checkpoint files to save parameter information. Once a fault is rectified, training is rolled back from the most recently saved checkpoint file for recovery. Each time a checkpoint file is saved, a specific training period is wasted. To ensure training efficiency, the interval for saving checkpoint files is usually large. However, a larger saving interval indicates longer time wasted for training rollback upon each fault. To solve this problem, MindIO ACP is introduced to asynchronously save checkpoints. For details, see Table 5.

  Table 5 Asynchronous checkpoint saving

  Function Name	Asynchronous checkpoint saving
  Feature Description	After checkpoints are obtained from the NPU, they are asynchronously written to storage to minimize training loss and the storage period for each checkpoint saving, thereby reducing the training rollback time.
  Instructions	Only cluster scheduling components and MindIO components of version 6.0.RC2 or later are supported.
  Key Operations	To install and use MindIO, see Optimizing Checkpoint Saving and Loading.

• Efficient checkpoint recovery: During training rollback and recovery, checkpoints must be loaded from storage. Due to the large volume of checkpoint data, directly reading and loading checkpoints from storage takes considerable time. To solve this problem, MindIO ACP is introduced for efficient checkpoint recovery. For details, see Table 6.

  Table 6 Efficient checkpoint recovery

  Function Name	Efficient checkpoint recovery
  Feature Description	MindIO stores the latest checkpoint in memory, allowing it to be read directly from memory during fault recovery, thereby reducing checkpoint read time.
  Instructions	Only cluster scheduling components and MindIO components of version 6.0.RC2 or later are supported.
  Key Operations	To install and use MindIO, see Optimizing Checkpoint Saving and Loading.

Operator Compilation Time

If an operator needs to be re-executed during resumable training, building the operator takes a long time. To solve this problem, you can select the operator binary or operator building cache to reduce the building time. For details, see Table 7 and Table 8.

NOTE The operator binary and operator compliation cache are incompatible. Please choose one of them to use.

Table 7 Operator Binary Function Description

torch.npu.set_compile_mode(jit_compile=False)

Table 8 Operator compilation cache

Recovery Time (MindSpore)

This section describes the optimization items that can be used to shorten the resumable training time on MindSpore, including Fault Detection Time, Training Rollback and Checkpoint Loading Time, and Compilation Cache Time.

Fault Detection Time

SA parameter plane network fault in a cluster may not affect a training job. Therefore, the cluster scheduling components do not forcibly interrupt the job. When the parameter plane network fault affects a training job, the network timeout mechanism of collective communication is triggered. After a default waiting period of 30 minutes, the cluster scheduling components can detect the fault and trigger resumable training. To solve this problem, MindSpore provides a watchdog fault detection function to determine if training jobs are affected and to reduce fault detection time. For details, see Table 1.

Table 1 Watchdog fault detection

...
context:
  ascend_config:
    hccl_watchdog: False
...

Training Rollback and Checkpoint Loading Time

• Asynchronous checkpoint saving: A training job periodically saves checkpoint files to save parameter information. Once a fault is rectified, training is rolled back from the most recently saved checkpoint file for recovery. Each time a checkpoint file is saved, a specific training period is wasted. To ensure training efficiency, the interval for saving checkpoint files is usually large. However, a larger saving interval indicates longer time wasted for training rollback upon each fault. To solve this problem, MindIO ACP is introduced to asynchronously save checkpoints. For details, see Table 2.

  Table 2 Asynchronous checkpoint saving

  Function Name	Asynchronous checkpoint saving
  Feature Description	After checkpoints are obtained from the NPU, they are asynchronously written to storage to minimize training loss and the storage period for each checkpoint saving, thereby reducing the training rollback time.
  Instructions	Only cluster scheduling components and MindIO components of version 6.0.RC2 and later are suppported.
  Key Operations	To install and use MindIO, see Optimizing Checkpoint Saving and Loading.

• Efficient checkpoint recovery: During training rollback and recovery, checkpoints must be loaded from storage. Due to the large volume of checkpoint data, directly reading and loading checkpoints from storage takes considerable time. To solve this problem, MindIO ACP is introduced for efficient checkpoint recovery. For details, see Table 3.

  Table 3 Efficient checkpoint recovery

  Function Name	Efficient checkpoint recovery
  Feature Description	Store the latest checkpoint in memory, allowing direct read from memory during fault recovery to reduce checkpoint read time.
  Instructions	Only cluster scheduling components and MindIO components of version 6.0.RC2 and later are supported.
  Key Operations	To install and use MindIO, see Opyimizing Checkpoint Saving and Loading.

Compilation Cache Time

During resumable training, a computational graph needs to be built. However, this process takes a long time in foundation model scenarios. To solve this problem, MindSpore can store a building cache file during the first building. During fault recovery, the graph building cache in storage can be directly read to reduce the graph building time. For details, see Table 4.

Table 4 Graph compilation cache

export MS_COMPILER_CACHE_ENABLE=1  # Enable graph compilation cache
export MS_COMPILER_CACHE_PATH=xxx  # Set the graph compilation cache path

Configuring Proactive HCCL Link Setup

If faults occur in the HCCL link setup phase, process-level rescheduling or process-level online recovery will fail. If HCCL link setup is required in other training phases in addition to the training initialization phase, you can set up the link in advance to avoid faults during the setup process.

PyTorch Single-Operator Scenario

n the PyTorch single-operator scenario, HCCL links are set up in lazy loading mode. After a Torch communication group is set up, its first operator triggers the creation of the HCCL communicator. After the creation, the inter-rank link is set up. Therefore, to ensure all communicators are linked during training initialization, a communication operator must be dispatched to each group at that stage.

The following is an example of actively creating a communication group:

rank = 0 # Set the rank of this process
sub_ranks = [0, 1, 2]  # Assume a communication group containing 0, 1, and 2
groupX = torch.distributed.new_group(ranks=sub_ranks,...) # Create communication group X
test_tensor = torch.ones(1).to(f'npu:{rank}') * (rank + 1)  # Construct a test data tensor
torch.distributed.all_reduce(test_tensor, op=dist.ReduceOp.SUM, group=groupX)  # Execute the all reduce operator in communication group X

Configuring Proactive Checkpoint Saving for Subhealthy Policy

If you want to save the dying gasp checkpoint when a subhealth fault occurs in a job, modify the job YAML. Specifically, configure the subhealth policy to graceExit, and set the fault recovery policy to dump. For the remaining startup script and job YAML configurations, see Configuring Dying Gasp Checkpoint Saving. This feature requires TaskD and ClusterD to function properly.

...
  labels:
     ...
     subHealthyStrategy: "graceExit"  # Configuring the Subhealth Policy
...
  annotations:
    ...
    recover-strategy: "dump"  # The available recovery policy for the task is to save an on-die checkpoint.
...