翻訳は機械翻訳により提供されています。提供された翻訳内容と英語版の間で齟齬、不一致または矛盾がある場合、英語版が優先します。

モデル並列処理によるモデルのチェックポイントと微調整

SageMaker モデル並列処理ライブラリは、さまざまなモデル並列処理戦略によって分割されたモデルの状態とオプティマイザの状態を保存し、トレーニングを再開して微調整する場所から継続的なトレーニングのチェックポイントをロードするためのチェックポイント APIs を提供します。この API は、モデルとオプティマイザの状態の一部または全体を保存するオプションもサポートしています。

分散モデルのチェックポイント機能

PyTorch と の間のフレームワーク TensorFlow と、使用する SageMaker モデル並列処理ライブラリのバージョンに応じて、次のいずれかのトピックを選択します。

分散 PyTorch モデルのチェックポイント ( SageMaker モデル並列処理ライブラリ v1.10.0 以降用）

SageMaker モデル並列処理ライブラリは、分散モデルの状態とそのオプティマイザの状態のチェックポイント全体または一部を保存およびロードするためのチェックポイント APIs を提供します。

部分チェックポイント

モデル並列処理を使用してトレーニングされたモデルのチェックポイントを保存するには、部分チェックポイントオプションを true (partial=True) に設定して、smdistributed.modelparallel.torch.save_checkpoint API を使用します。これにより、各モデルのパーティションが個別に保存されます。モデルとオプティマイザの状態に加えて、user_content 引数を介して追加のカスタムデータを保存することもできます。チェックポイントモデル、オプティマイザ、およびユーザーコンテンツは個別のファイルとして保存されます。save_checkpoint API コールにより、次の構造でチェックポイントフォルダが作成されます。

- path
  - ${tag}_partial (folder for partial checkpoints)
    - model_rankinfo.pt
    - optimizer_rankinfo.pt
    - fp16_states_rankinfo.pt
    - user_content.pt
  - $tag (checkpoint file for full checkpoints)
  - user_content_$tag (user_content file for full checkpoints)
  - newest (a file that indicates the newest checkpoint)

部分チェックポイントからトレーニングを再開するには、partial=True として smdistributed.modelparallel.torch.resume_from_checkpoint API を使用し、チェックポイントディレクトリと部分チェックポイントの保存時に使用したタグを指定します。モデルの重みが実際に読み込まれるのは、モデルパーティショニングの後、smdistributed.modelparallel.torch.step - 修飾トレーニングステップ関数の初回実行時に行われることにご注意ください。

部分チェックポイントを保存するとき、ライブラリは .pt ファイル拡張子のファイルとして、モデルパーティションの決定も保存します。逆に、部分チェックポイントから再開するとき、ライブラリはパーティション決定ファイルを一緒に読み込みます。パーティション決定が読み込まれたら、パーティションを変更することはできません。

次のコードスニペットは、トレーニングスクリプトで PyTorchチェックポイント APIsを設定する方法を示しています。

import smdistributed.modelparallel.torch as smp

model = ...
model = smp.DistributedModel(model)
optimizer = ...
optimizer = smp.DistributedOptimizer(optimizer)
user_content = ...     # additional custom data
checkpoint_path = "/opt/ml/checkpoint/model_parallel"

# Save a checkpoint.
smp.save_checkpoint(
    path=checkpoint_path,
    tag=f"total_steps{total_steps}",
    partial=True,
    model=model,
    optimizer=optimizer,
    user_content=user_content
    num_kept_partial_checkpoints=5
)

# Load a checkpoint.
# This automatically loads the most recently saved checkpoint.
smp_checkpoint = smp.resume_from_checkpoint(
    path=checkpoint_path, 
    partial=True
)

完全チェックポイント

推論の目的で最終モデルアーティファクトを保存するには、partial=False として smdistributed.modelparallel.torch.save_checkpoint API を使用します。これにより、モデルパーティションが組み合わされて単一のモデルアーティファクトが作成されます。これはオプティマイザの状態は組み合わせないことに注意してください。

完全なモデルチェックポイントを指定して、特定の重みでトレーニングを初期化するには、partial=False として smdistributed.modelparallel.torch.resume_from_checkpoint API を使用できます。これはオプティマイザの状態は読み込まないことに注意してください。

次のコードは、チェックポイント APIs を使用して、 PyTorch モデル並列処理でトレーニングされたモデルを完全にチェックポイントする方法の例を示しています。

import smdistributed.modelparallel.torch as smp

model = ...
model = smp.DistributedModel(model)
optimizer = ...
optimizer = smp.DistributedOptimizer(optimizer)
user_content = ...     # additional custom data
checkpoint_path = "/opt/ml/checkpoint/model_parallel"

# Save a checkpoint.
smp.save_checkpoint(
    path=checkpoint_path,
    tag=f"total_steps{total_steps}",
    partial=False,
    model=model,
    optimizer=optimizer,
    user_content=user_content
    num_kept_partial_checkpoints=5
)

# Load a checkpoint.
# This automatically loads the most recently saved checkpoint.
smp_checkpoint = smp.resume_from_checkpoint(
    path=checkpoint_path, 
    partial=False
)

分散 PyTorch モデルのチェックポイント (v1.6.0 から v1.9.0 までの SageMaker モデル並列処理ライブラリ用）

SageMaker モデル並列処理ライブラリは、テンソル並列処理によるトレーニングジョブのチェックポイントの一部または全体を保存するための Python 関数を提供します。次の手順は、テンソル並列処理を使用するときにチェックポイントを保存および読み込むように smp.save() と smp.load() を使用する方法を示しています。

分散 TensorFlow モデルのチェックポイント

TensorFlow モデル並列処理によるトレーニング中にモデルを保存するには、 SageMaker モデル並列処理ライブラリが提供する次の関数を使用します。

分散モデルの微調整

微調整はトレーニングスクリプトで設定する必要があります。次のコードスニペットは、モジュールsmdistributed.model.parallel.torchと微調整の設定を登録するための変更を加えた Hugging Face Transformer の AutoModelForCausalLM クラスを使用するトレーニングスクリプトの構造の例を示しています。

import argparse
from transformers import AutoModelForCausalLM
import smdistributed.modelparallel
import smdistributed.modelparallel.torch as smp

def parse_args():

    parser = argparse.ArgumentParser()

    # set an arg group for model
    model_grp = parser.add_argument_group(
        title="model", description="arguments to describe model configuration"
    )

    ... # set up numerous args to parse from the configuration dictionary to the script for training

    # add arg for activating fine-tuning
    model_grp.add_argument(
        "--fine_tune",
        type=int,
        default=0,
        help="Fine-tune model from checkpoint or pretrained model",
    )

def main():
    """Main function to train GPT."""
    args = parse_args()

    ... # parse numerous args

    if args.fine_tune > 0 and args.delayed_param > 0 and smp.rank() == 0:
        pretrained_model = AutoModelForCausalLM.from_pretrained(
            args.model_name or args.model_dir
        )
        model_state_dict = pretrained_model.state_dict()
        path = os.path.join(args.model_dir, "fullmodel.pt")
        torch.save(model_state_dict, path)

    # create a Transformer model and wrap by smp.model_creation() 
    # with options to configure model parallelism parameters offered by SageMaker
    with smp.model_creation(
        tensor_parallelism=smp.tp_size() > 1 or args.use_distributed_transformer > 0,
        zero_init=args.use_distributed_transformer == 0,
        dtype=dtype,
        distribute_embedding=args.sharded_data_parallel_degree > 1 and smp.tp_size() > 1,
        use_alibi=args.alibi > 0,
        attention_in_fp32=args.attention_in_fp32 > 0,
        fp32_residual_addition=args.residual_addition_in_fp32 > 0,
        query_key_layer_scaling=args.query_key_layer_scaling > 0 and args.bf16 < 1,
        fused_softmax=args.fused_softmax > 0,
        fused_dropout=args.fused_dropout > 0,
        fused_bias_gelu=args.fused_bias_gelu > 0,
        flash_attention=args.flash_attention > 0,
    ):
        if args.fine_tune > 0 and args.delayed_param == 0:
            model = AutoModelForCausalLM.from_pretrained(
                args.model_name or args.model_dir
            )
        else:
            model = AutoModelForCausalLM.from_config(model_config)

    # wrap the model by smp.DistributedModel() to apply SageMaker model parallelism
    model = smp.DistributedModel(
        model, trace_device="gpu", backward_passes_per_step=args.gradient_accumulation
    )

    # wrap the optimizer by smp.DistributedOptimizer() to apply SageMaker model parallelism
    optimizer= ... # define an optimizer
    optimizer = smp.DistributedOptimizer(
        optimizer,
        static_loss_scale=None,
        dynamic_loss_scale=True,
        dynamic_loss_args={"scale_window": 1000, "min_scale": 1, "delayed_shift": 2},
    )

    # for fine-tuning, use smp.resume_from_checkpoint() to load a pre-trained model
    if args.fine_tune > 0 and args.delayed_param > 0:
        smp.resume_from_checkpoint(args.model_dir, tag="fullmodel.pt", partial=False)

トレーニングスクリプトと Jupyter Notebook の完全な例については、「例」リポジトリの「 の GPT-2 の例 PyTorch」を参照してください。 SageMaker GitHub