easyckpt的技術原理和使用方法 - Platform For AI

EasyCkpt是PAI為PyTorch大模型訓練開發的高效能Checkpoint架構，通過近零開銷的模型儲存機制，實現訓練全過程進度無損的儲存與恢複。任務中斷後可從最新Checkpoint繼續訓練，避免重複計算。目前支援Megatron和DeepSpeed架構。

背景資訊

大模型訓練過程中可能遇到硬體故障、系統異常等問題，頻繁中斷會造成大量訓練進度損失。雖然Checkpoint可以儲存和恢複進度，但其耗時與模型大小成正比——百億到千億參數的模型，Checkpoint通常需要幾分鐘到十幾分鐘，期間訓練必須暫停。中斷後需重新計算丟失的迭代，通常耗時數小時。以1千卡規模為例，損失將達數千卡時。

因此，需要一種低成本的方法儲存最新Checkpoint，以便故障後快速恢複訓練、避免重複計算。

功能介紹

PAI從大量故障案例中總結出GPU和深度學習情境的以下特點：

特點1：任務的故障是部分的。
通常，故障的根因是一到兩台機器的故障，這隻會影響部分Worker。對於大規模分布式訓練任務而言，不會所有的Worker都出現故障。
特點2：機器的故障是部分的。
通過對眾多故障失敗案例的分析，對於訓練叢集而言：
- GPU易損壞，但機器的CPU和記憶體通常仍可正常使用。
- 以節點為單位，記憶體的閑置空間很大（通常遠大於模型狀態）。
- 節點上所有網卡同時損壞的機率極低，即使節點故障，通常仍具備通訊能力。
特點3：模型的失敗是部分的。
大規模模型訓練通常採用3D並行或ZeRO最佳化，資料並行副本數大於1，模型參數在多個副本上有備份。某個GPU故障時，可從其他機器的副本恢複。

基於上述特點，PAI開發了EasyCkpt架構。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啟動任務，EasyCkpt相容該方式，改動量極小。

改動樣本

啟動參數：EasyCkpt複用Transformer的Checkpoint參數，含義與Transformer中一致。以下配置樣本表示每2個mini-batch儲存一次Checkpoint，最多保留最近2個副本。

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

--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數量時可能刪除舊資料。以下是EasyCkpt的讀寫行為說明和資料安全保證。

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

從load目錄中讀取Checkpoint資料，並在需要時將其拼接成新的checkpoint資料。
將checkpoints資料儲存到save目錄中，並在必要時根據功能設定刪除save目錄中具有Megatron或Transformers格式的checkpoint檔案夾。

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

不會對save和load目錄以外的資料做任何操作。
EasyCkpt會以日誌形式記錄其進行的所有儲存或刪除的操作。

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