Platform For AI:EasyCkpt

トレーニングループの変更

Megatron ネイティブの load_checkpoint インポートを EasyCkpt のバージョンに置き換え、トレーニングループに initialize_easyckpt と save_checkpoint_if_needed の呼び出しを追加します。

from megatron.core.utils import get_model_config
from megatron import print_rank_0
from megatron import print_rank_last
# Megatron ネイティブの load_checkpoint を EasyCkpt のものに置き換えます
# 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

# EasyCkpt の関数をインポートします
from aimaster.python.torch.easyckpt.megatron import (load_checkpoint,
                                                    initialize_easyckpt,
                                                    save_checkpoint_if_needed)

def print_datetime(string):
    """この呼び出しはすべてのランクで同期されることに注意してください。"""
    timers('interval-time', log_level=0).start(barrier=True)
    print_datetime('before the start of training step')
    report_memory_flag = True

    # トレーニングループの前に EasyCkpt を初期化します。
    # save_mem_interval=1: イテレーションごとにメモリにコピー
    # save_storage_interval=5: 5 イテレーションごとにストレージに書き込み
    # max_ckpt_num=5: ディスクに最大 5 つのチェックポイントを保持
    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)

        # ログ記録
        loss_scale = optimizer.get_loss_scale().item()
        params_norm = None

EasyCkpt フックのインポート

モデルトレーニングファイル (例: pretrain_gpt.py) で、EasyCkpt Megatron フックをインポートします：

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

# 非同期チェックポイント処理を有効にするために EasyCkpt フックをインポートします
import aimaster.python.torch.easyckpt.megatron.hook

def model_provider(pre_process=True, post_process=True):
    """モデルをビルドします。"""

関連 API

EasyCkpt は、Megatron との統合用に 3 つの関数を提供します。

initialize_easyckpt(save_mem_interval, save_storage_interval, max_ckpt_num, log_file_path=None)

load_checkpoint(model, optimizer, opt_param_scheduler, load_arg='load', strict=True, concat=False)

save_checkpoint_if_needed(iteration, model, optimizer, opt_param_scheduler)

チェックポイントパラメータの設定

EasyCkpt は、標準の Transformers のチェックポイントパラメータを再利用します。これらを起動スクリプトに追加します。

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

この例では、2 ステップごとにチェックポイントを保存し、最大 2 つのチェックポイントを保持します。

Trainer のラップ

EasyCkpt の Transformers モジュールから TrainerWrapper をインポートし、お使いの Trainer インスタンスをラップしてチェックポイントの再開を有効にします:

import datasets
import transformers

# Transformers Trainer 用の EasyCkpt ラッパーをインポートします
from aimaster.python.torch.easyckpt.transformers import TrainerWrapper

logger = logging.getLogger(__name__)

    tokenizer=tokenizer,
    data_collator=transformers.default_data_collator,
  )
  # 非同期チェックポイント処理を有効にするために Trainer を EasyCkpt でラップします
  trainer = TrainerWrapper(trainer)
  # 最新のチェックポイントから自動的に再開します
  trainer.train(resume_from_checkpoint=True)

if __name__ == "__main__":
  main()

関連 API

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

trainer = TrainerWrapper(trainer)
trainer.train(resume_from_checkpoint=True)