Megatron-DeepSpeedからHFのTransformersにチェックポイントを変換する

こちらのプロジェクトでMegatron-DeepSpeedからHFのTransformersにチェックポイントを変換する必要があったので

[GENIAC 採択プロジェクト] 多様な日本語能力の向上を目指した公開の基盤モデル開発 | 東京大学松尾研究室 - Matsuo Lab[GENIAC 採択プロジェクト] 多様な日本語能力の向上を目指した公開の基盤モデル開発

Megatron-DeepSpeedは、NVIDIAのMegatron-LMをもとにした拡張版で、DeepSpeedライブラリを使用しています。

Megatron-DeepSpeedには、DeepSpeedからMegatron-LMやHugging FaceのTransformersライブラリ形式へのチェックポイントの変換にするためのPythonスクリプトが付属しています。

変換コード

#!/usr/bin/env python

import os
import torch
import json

from deepspeed_checkpoint import DeepSpeedCheckpoint
from deepspeed_to_megatron import _create_rank_checkpoint, parse_arguments

# the import was tested to work with this version
# https://github.com/huggingface/transformers/commit/0af901e83 if it diverges we may consider
# copying that version here instead
from transformers.models.megatron_gpt2.convert_megatron_gpt2_checkpoint import convert_megatron_checkpoint
from transformers import GPT2Config

def main():

    # this first part comes mainly from deepspeed_to_megatron.main
    args = parse_arguments()
    print(f'Converting DeepSpeed checkpoint in {args.input_folder} to HF Transformers checkpoint in {args.output_folder}')

    ds_checkpoint = DeepSpeedCheckpoint(args.input_folder, args.target_tp, args.target_pp)
    iteration = ds_checkpoint.get_iteration()
    input_state_dict = _create_rank_checkpoint(ds_checkpoint, 0, 0, args.for_release)

    # the 2nd part comes from transformers.models.megatron_gpt2.convert_megatron_gpt2_checkpoint.main
    # Spell out all parameters in case the defaults change.
    config = GPT2Config(
        vocab_size=50257,
        n_positions=1024,
        n_ctx=1024,
        n_embd=1024,
        n_layer=24,
        n_head=16,
        n_inner=4096,
        activation_function="gelu",  # used to be "gelu_new" in earlier versions
        resid_pdrop=0.1,
        embd_pdrop=0.1,
        attn_pdrop=0.1,
        layer_norm_epsilon=1e-5,
        initializer_range=0.02,
        summary_type="cls_index",
        summary_use_proj=True,
        summary_activation=None,
        summary_proj_to_labels=True,
        summary_first_dropout=0.1,
        scale_attn_weights=True,
        gradient_checkpointing=False,
        use_cache=True,
        bos_token_id=50256,
        eos_token_id=50256,
    )

    # Convert.
    print("Converting to HF Checkpoint")
    output_state_dict = convert_megatron_checkpoint(args, input_state_dict, config)

    basename = args.output_folder
    os.makedirs(basename, exist_ok=True)

    # Print the structure of converted state dict.
    #if args.print_checkpoint_structure:
    #    recursive_print(None, output_state_dict)

    # Store the config to file.
    output_config_file = os.path.join(basename, "config.json")
    output_config = config.to_dict()
    output_config["architectures"] = ["GPT2LMHeadModel"]
    output_config["model_type"] = "gpt2"
    print(f'Saving config to "{output_config_file}"')
    with open(output_config_file, "w") as f:
        json.dump(output_config, f)

    # Store the state_dict to file.
    output_checkpoint_file = os.path.join(basename, "pytorch_model.bin")
    print(f'Saving checkpoint to "{output_checkpoint_file}"')
    torch.save(output_state_dict, output_checkpoint_file)

    print("Now add tokenizer files and upload to the hub")


if __name__ == "__main__":
    main()

Megatron-DeepSpeedからHugging FaceのTransformersライブラリへの変換を行うには、以下のように直接行うことができます。

python tools/convert_checkpoint/deepspeed_to_transformers.py  \
--input_folder /path/to/Megatron-Deepspeed/checkpoint/global_step97500 \
--output_folder /path/to/transformers/checkpoint

transformersは現在PP=1/TP=1でのみ動作するため、デフォルトの--target_tp 1/--target_pp 1を使用します。

ステップ1: 引数の解析

parse_arguments()関数を使用してコマンドライン引数を解析します。これには入力と出力のフォルダパス、ターゲットのテンソル並列度(TP)とパイプライン並列度(PP)、などが含まれます。

ステップ2: DeepSpeedチェックポイントの読み込み

引数をもとに、DeepSpeedCheckpointクラスのインスタンスを作成し、指定されたディレクトリからチェックポイントを読み込みます。これにより、モデルの状態や訓練の進行状況などの情報が取得されます。

ステップ3: チェックポイントの変換準備

_create_rank_checkpoint関数を使用して、DeepSpeedチェックポイントからMegatron-LMチェックポイントに変換します。現在はPP=1/TP=1でのみ動作するらしい。

ステップ4: GPT-2モデル構成の設定

TransformersライブラリのGPT2Configクラスを使用して、変換後のモデルの構成を定義します。この構成は、モデルの各層のサイズやドロップアウト率、アクティベーション関数などを指定します。

ステップ5: チェックポイントの変換実行

Transformersライブラリのconvert_megatron_checkpoint関数を呼び出し、Megatron-LMチェックポイントとGPT-2のモデル構成を渡して、Transformersチェックポイントに変換します。

ステップ6: 変換されたチェックポイントの保存

出力ディレクトリを作成し、変換されたモデルの構成をJSONファイルとして、モデルの状態をPyTorchのバイナリファイルとしてそれぞれ保存します。

応用編

Llama2、Mistral、Mixtralの対応は記述予定です。

Transformersライブラリのconvert_megatron_checkpoint関数に変更を加える必要があります。