Amazon Web Services ブログ

新しい Amazon SageMaker HyperPod レシピにより基盤モデルのトレーニングと微調整を加速

12 月 4 日、Amazon SageMaker HyperPod レシピの一般提供が開始されたことをお知らせいたします。このレシピは、あらゆるスキルセットのデータサイエンティストや開発者が、基盤モデル (FM) のトレーニングと微調整を数分で開始できるようにすると同時に、最先端のパフォーマンスを実現します。Llama 3.1 405B、Llama 3.2 90B、Mixtral 8x22B など、一般公開されている人気の FM をトレーニングおよび微調整するための最適化されたレシピをご利用いただけるようになりました。

AWS re:Invent 2023 では、FM のトレーニングにかかる時間を最大 40% 短縮し、事前設定された分散トレーニングライブラリと並行して 1,000 個を超えるコンピューティングリソース全体でのスケールリングを可能にする、SageMaker HyperPod を発表しました。SageMaker HyperPod を使用すると、トレーニングに必要な高速コンピューティングリソースを見つけ、最適なトレーニングプランを作成し、コンピューティングリソースの可用性に応じて、さまざまなキャパシティブロックにまたがってトレーニングワークロードを実行できます。

SageMaker HyperPod レシピには AWS でテストされたトレーニングスタックが含まれているため、さまざまなモデル構成を試す面倒な作業を省くことができ、何週間にもわたる反復的な評価とテストが不要になります。レシピは、トレーニングデータセットの読み込み、分散トレーニング手法の適用、障害からの迅速な回復のためのチェックポイントの自動化、エンドツーエンドのトレーニングループの管理など、いくつかの重要なステップを自動化します。

レシピを変更するだけで、GPU ベースと Trainium ベースのインスタンスの切り替えがシームレスに行え、トレーニングパフォーマンスをさらに最適化し、コストを削減できます。SageMaker HyperPod または SageMaker トレーニングジョブでは、本稼働のワークロードを簡単に実行できます。

SageMaker HyperPod レシピの実践
はじめに、SageMaker HyperPod レシピの GitHub リポジトリにアクセスして、一般公開されている人気の FM 用トレーニングレシピを参照します。

簡単なレシピパラメータを編集して、クラスター構成のインスタンスタイプやデータセットの場所を指定するだけで、1 行のコマンドでレシピを実行し、最先端のパフォーマンスを実現できます。

リポジトリの複製後、レシピの config.yaml ファイルを編集して、モデルとクラスタータイプを指定する必要があります。

$ git clone --recursive https://github.com/aws/sagemaker-hyperpod-recipes.git
$ cd sagemaker-hyperpod-recipes
$ pip3 install -r requirements.txt.
$ cd ./recipes_collections
$ vim config.yaml

レシピは、Slurm を使用した SageMaker HyperPod、Amazon Elastic Kubernetes Service (Amazon EKS) を使用した SageMaker HyperPod、SageMaker トレーニングジョブをサポートしています。例えば、クラスタータイプ (Slurm オーケストレーター)、モデル名 (Meta Llama 3.1 405B 言語モデル)、インスタンスタイプ (ml.p5.48xlarge)、そしてトレーニングデータ、結果、ログなどを保存するデータロケーションを設定できます。

defaults:
- cluster: slurm # サポート: slurm / k8s / sm_jobs
- recipes: fine-tuning/llama/hf_llama3_405b_seq8k_gpu_qlora # トレーニングするモデルの名前
debug: False # ランチャーの設定をデバックする場合は True に設定
instance_type: ml.p5.48xlarge # またはその他のサポートされているクラスターインスタンス
base_results_dir: # 結果、チェックポイント、ログなどを保存する場所

必要な場合、この YAML ファイルのモデル固有のトレーニングパラメータを調整できます。このパラメータには、アクセラレータデバイスの数、インスタンスタイプ、トレーニング精度、並列化とシャーディングの手法、オプティマイザー、TensorBoard による実験をモニタリングするためのログ記録など、最適な構成の概要が記載されています。

run:
  name: llama-405b
  results_dir: ${base_results_dir}/${.name}
  time_limit: "6-00:00:00"
restore_from_path: null
trainer:
  devices: 8
  num_nodes: 2
  accelerator: gpu
  precision: bf16
  max_steps: 50
  log_every_n_steps: 10
  ...
exp_manager:
  exp_dir: # TensorBoard ログ記録の場所
  name: helloworld 
  create_tensorboard_logger: True
  create_checkpoint_callback: True
  checkpoint_callback_params:
    ...
  auto_checkpoint: True # 自動チェックポイント用
use_smp: True 
distributed_backend: smddp # 最適化された集団通信
# 事前トレーニング済みのモデルからトレーニングを開始
model:
  model_type: llama_v3
  train_batch_size: 4
  tensor_model_parallel_degree: 1
  expert_model_parallel_degree: 1
  # その他のモデル固有のパラメータ

Slurm を使用した SageMaker HyperPod でこのレシピを実行するには、クラスターのセットアップ手順に従って SageMaker HyperPod クラスターを準備する必要があります。

次に、SageMaker HyperPod ヘッドノードに接続し、Slurm コントローラーにアクセスして、編集したレシピをコピーします。続いて、ヘルパーファイルを実行して、ジョブの Slurm 送信スクリプトを生成します。このスクリプトは、トレーニングジョブの開始前に内容を検査するためのドライランに使用できます。

$ python3 main.py --config-path recipes_collection --config-name=config

トレーニングが完了すると、トレーニングされたモデルは自動的に割り当てられたデータロケーションに保存されます。

Amazon EKS を使用した SageMaker HyperPod でこのレシピを実行するには、GitHub リポジトリからレシピを複製し、要件をインストールして、ノートパソコンでレシピ (cluster: k8s) を編集します。次に、ノートパソコンと実行中のEKS クラスター間のリンクを作成し、続いて HyperPod コマンドラインインターフェイス (CLI) を使用して、レシピを実行します。

$ hyperpod start-job –recipe fine-tuning/llama/hf_llama3_405b_seq8k_gpu_qlora \
--persistent-volume-claims fsx-claim:data \
--override-parameters \
'{
  "recipes.run.name": "hf-llama3-405b-seq8k-gpu-qlora",
  "recipes.exp_manager.exp_dir": "/data/<your_exp_dir>",
  "cluster": "k8s",
  "cluster_type": "k8s",
  "container": "658645717510.dkr.ecr.<region>.amazonaws.com/smdistributed-modelparallel:2.4.1-gpu-py311-cu121",
  "recipes.model.data.train_dir": "<your_train_data_dir>",
  "recipes.model.data.val_dir": "<your_val_data_dir>",
}'

SageMaker Python SDK を使用して SageMaker トレーニングジョブでレシピを実行することもできます。次の例では、トレーニングレシピをオーバーライドして、SageMaker トレーニングジョブで PyTorch トレーニングスクリプトを実行しています。

...
recipe_overrides = {
    "run": {
        "results_dir": "/opt/ml/model",
    },
    "exp_manager": {
        "exp_dir": "",
        "explicit_log_dir": "/opt/ml/output/tensorboard",
        "checkpoint_dir": "/opt/ml/checkpoints",
    },  
    "model": {
        "data": {
            "train_dir": "/opt/ml/input/data/train",
            "val_dir": "/opt/ml/input/data/val",
        },
    },
}
pytorch_estimator = PyTorch(
           output_path=<output_path>,
           base_job_name=f"llama-recipe",
           role=<role>,
           instance_type="p5.48xlarge",
           training_recipe="fine-tuning/llama/hf_llama3_405b_seq8k_gpu_qlora",
           recipe_overrides=recipe_overrides,
           sagemaker_session=sagemaker_session,
           tensorboard_output_config=tensorboard_output_config,
)
...

トレーニングの進行に伴い、モデルチェックポイントは Amazon Simple Storage Service (Amazon S3) に保存され、完全に自動化されたチェックポイント機能により、トレーニング障害やインスタンスの再起動からの迅速な復旧が可能になります。

今すぐ利用可能
Amazon SageMaker HyperPod レシピが SageMaker HyperPod レシピの GitHub リポジトリで入手可能になりました。詳細については、SageMaker HyperPod の製品ページと「Amazon SageMaker AI デベロッパーガイド」をご覧ください。

SageMaker HyperPod レシピをお試しいただき、AWS re:Post for SageMaker 宛てに、または通常の AWS サポートの連絡先を通じて、フィードバックをお寄せください。

– Channy

原文はこちらです。