背景信息

大模型訓練面臨的困難在于確保訓練任務能夠持續進行而不中斷。在訓練過程中，可能會遇到硬件故障、系統問題、連接錯誤、以及其他未知的問題。這種頻繁中斷導致的訓練進度的損失對于耗時又耗資源的大模型訓練來說是難以承受的。盡管可以通過做Checkpoint來保存和恢復進度，但Checkpoint本身的耗時與模型的大小成正比。對于當前典型的擁有百億到千億參數的大模型而言，Checkpoint的時間開銷通常在幾分鐘到十幾分鐘之間。此時，訓練任務需要暫停，使得用戶難以頻繁進行Checkpoint操作。然而，一旦發生中斷，之前損失的迭代次數在恢復時需要重新計算，通常會花費數小時的時間。考慮到大模型使用的GPU規模，以1千卡為例，損失將會是數千個卡的時間。

因此，在發生故障時，迫切需要一種以低成本的方法來保存最新的Checkpoint。這樣在重新啟動訓練時就無需重復計算，從而減少時間和成本的浪費。

功能介紹

針對頻繁故障的情況，PAI通過之前的故障失敗場景總結出以下GPU和深度學習場景故障的功能特點：

• 特點1：任務的故障是部分的。

  通常，故障的根因是一到兩臺機器的故障，這只會影響部分Worker。對于大規模分布式訓練任務而言，不會所有的Worker都出現故障。

  

• 特點2：機器的故障是部分的。

  通過對眾多故障失敗案例的分析，對于訓練集群而言：

  • GPU易損壞，但機器的CPU和內存通常仍可正常使用。

  • 以節點為單位，內存的閑置空間很大（通常遠大于模型狀態）。

  • 在節點上，通常很難出現所有網卡都損壞的情況，即使節點出現故障，它通常仍然具備通信能力。

    

• 特點3：模型的失敗是部分的。

  在大規模模型訓練中通常采用3D并行或ZeRO系列優化，大多數任務的數據并行副本數大于1，這使得模型訓練參數在多個副本上都有備份。因此，當某個GPU設備發生故障時，可以通過其他機器的GPU上保留的副本來實現恢復。

針對上述內容提到的大模型場景下Checkpoint的特點， PAI開發了EasyCkpt框架，該框架提供了高性能的Checkpoint功能。通過采用異步化層次化保存、重疊模型拷貝和計算、網絡感知異步存儲的策略， EasyCkpt實現了近0開銷的模型保存機制和大模型訓練全過程精度無損的模型保存與恢復的能力。EasyCkpt已經支持當前流行的兩種大模型訓練框架Megatron和DeepSpeed，您只需修改幾行代碼即可使用。

操作入口

安裝AIMaster SDK

本功能依賴AIMaster SDK，其安裝方式如下：

# py36
pip install -U http://odps-release.cn-hangzhou.oss.aliyun-inc.com/aimaster/pai_aimaster-1.2.1-cp36-cp36m-linux_x86_64.whl

# py38
pip install -U http://odps-release.cn-hangzhou.oss.aliyun-inc.com/aimaster/pai_aimaster-1.2.1-cp38-cp38-linux_x86_64.whl

# py310
pip install -U http://odps-release.cn-hangzhou.oss.aliyun-inc.com/aimaster/pai_aimaster-1.2.1-cp310-cp310-linux_x86_64.whl

Megatron

改動示例

在Megatron框架下的training.py文件中，只需修改4行代碼：

并在您的實際訓練代碼（以pretrain_gpt.py為例）里，import一行代碼即可：

training.py文件修改后的代碼如下：

from megatron.core.utils import get_model_config
from megatron import print_rank_0
from megatron import print_rank_last
# from megatron.checkpointing import load_checkpoint
from megatron.checkpointing import save_checkpoint
from megatron.model import Float16Module
from megatron.model import GPTModel
from megatron.utils import report_memory
from megatron.model.vision.knn_monitor import compute_feature_bank
from aimaster.python.torch.easyckpt.megatron import (load_checkpoint,
                                                    initialize_easyckpt,
                                                    save_checkpoint_if_needed)

def print_datetime(string):
    """Note that this call will sync across all ranks."""
    timers('interval-time', log_level=0).start(barrier=True)
    print_datetime('before the start of training step')
    report_memory_flag = True

    initialize_easyckpt(save_mem_interval=1, save_storage_interval=5, max_ckpt_num=5, log_file_path='./test.log')

    while iteration < args.train_iters:
        if args.profile and \
           iteration == args.profile_step_start and \
                                       args.micro_batch_size * \
                                       get_num_microbatches()

        save_checkpoint_if_needed(iteration, model, optimizer, opt_param_scheduler)

        # Logging.
        loss_scale = optimizer.get_loss_scale().item()
        params_norm = None

實際訓練代碼（以pretrain_gpt.py為例）修改后的代碼如下：

from megatron.utils import average_losses_across_data_parallel_group
from megatron.arguments import core_transformer_config_from_args

import aimaster.python.torch.easyckpt.megatron.hook

def model_provider(pre_process=True, post_process=True):
    """Build the model."""

接口詳解

適用于Megatron的EasyCkpt框架提供了以下接口，其插入位置如上所示：

• load_checkpoint（model, optimizer, opt_param_scheduler, load_arg='load', strict=True, concat=False）：該接口在Megatron框架原生load_checkpoint()函數簽名基礎上增加了concat參數，若您使用的是Megatron2304版本，只需平替掉Megatron的load_checkpoint即可；如果您使用2305或2306版本，請參考下面注意事項的描述。

• initialize_easyckpt（save_mem_interval, save_storage_interval, max_ckpt_num, log_file_path=None）：該接口用于初始化Easyckpt框架。其中，通過save_mem_interval指定memory copy的頻率，通過save_storage_interval指定異步存儲到存儲設備的頻率，通過max_ckpt_num指定存儲設備中最多保存的Checkpoints數目。若需要保存詳盡log信息可指定log路徑到log_file_path里。

• save_checkpoint_if_needed（iteration, model, optimizer, opt_param_scheduler）：該接口用于調用EasyCkpt框架進行內存中的Checkpoint操作，其中各參數均為Megatron代碼內已有的變量名，您無需額外指定。

注意：如果您使用Megatron2305或2306的穩定版本，并且您開啟了distributed-optimizer。那么，當您在load的時候需要改變機器數目，或是希望合并分布式優化器參數時，您需要在training.py文件里load_checkpoint()函數被調用處設置concat參數為True。

DeepSpeed

DeepSpeed的大模型用戶通常使用Transformer的Trainer來啟動DeepSpeed任務，EasyCkpt兼容了這種使用方法，最小化了所需要的改動。

改動示例

啟動參數：適用于DeepSpeed的EasyCkpt框架復用了Transformer的Checkpoint參數來傳遞給DeepSpeed，這些參數的含義與Transformer中定義的一致，并且在接口詳解中均有詳細解釋。當前配置示例表示每2個mini-batch保存一次Checkpoint，最多同時保留最近的兩個持久化的Checkpoint副本。

修改后的代碼如下（對應上圖右側）：

--max_steps=10 \
--block_size=2048 \
--num_train_examples=100000 \
--gradient_checkpointing=false \
--save_strategy="steps" \
--save_steps="2" \
--save_total_limit="2"

訓練代碼改動：Transformer的Trainer需要用EasyCkpt提供的TrainerWrapper進行包裝，并打開resume_from_Checkpoint參數。

修改后的代碼如下（對應上圖右側）：

import datasets
import transformers

from aimaster.python.torch.easyckpt.transformers import TrainerWrapper

logger = logging.getLogger(__name__)

    tokenizer=tokenizer,
    data_collator=transformers.default_data_collator,
  )
  trainer = TrainerWrapper(trainer)
  trainer.train(resume_from_checkpoint=True)

if __name__ = ""__main__":
  main()

接口詳解

適用于DeepSpeed的EasyCkpt框架提供了以下接口，其插入位置如上所示：

• save_strategy：訓練過程中Checkpoint的保留方式。取值如下：

  • no： 訓練期間不進行保存。

  • epoch：在每個epoch結束時進行保存。

  • steps：根據save_steps的設置進行保存。

• save_steps：訓練過程中每隔多少步保存一次Checkpoint，這里要求save_strategy必須是"steps"。

• save_total_limit：最多保留多少個Checkpoint。

注意：按照Transformers官方文檔，使用save_total_limit時會將過時的checkpoints文件夾刪除掉，請確保過時的checkpoints文件夾是能夠安全刪除的。

數據安全須知

EasyCkpt需要在您提供的用戶存儲空間中讀寫數據，并且需要控制最大Checkpoint數量，可能引入刪除數據的動作。為了確保您的數據安全，PAI在此定義了EasyCkpt涉及的所有讀寫動作，并給出EasyCkpt提供的數據安全保證以及推薦的用戶使用方式。

EasyCkpt引入了以下讀寫動作，使用EasyCkpt即表示默認授權：

• 從load目錄中讀取Checkpoint數據，并在需要時將其拼接成新的checkpoint數據。

• 將checkpoints數據保存到save目錄中，并在必要時根據功能設定刪除save目錄中具有Megatron或Transformers格式的checkpoint文件夾。

在此基礎上，EasyCkpt提供以下保證：

• 不會對save和load目錄以外的數據做任何操作。

• EasyCkpt會以日志形式記錄其進行的所有保存或刪除的操作。

強烈建議您不要在模型的save或load文件路徑下存放任何其他數據。這種設置不符合EasyCkpt的預期用法，可能影響EasyCkpt的軟件能力。由此造成的數據風險和數據損失問題，使用者應自行承擔責任。