Platform For AI：EasyCkpt：AI大模型高效能狀態儲存恢複

背景資訊

大模型訓練面臨的困難在於確保訓練任務能夠持續進行而不中斷。在訓練過程中，可能會遇到硬體故障、系統問題、串連錯誤、以及其他未知的問題。這種頻繁中斷導致的訓練進度的損失對於耗時又耗資源的大模型訓練來說是難以承受的。儘管可以通過做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的軟體能力。由此造成的資料風險和資料損失問題，使用者應自行承擔責任。