「Amazon SageMaker Hyperband 自動モデルチューニングを使用して、分散トレーニングの収束問題を効果的に解決する」

Using Amazon SageMaker Hyperband Auto Model Tuning to effectively solve the convergence problem in distributed training.

最近の数年間で、深層学習ニューラルネットワーク（DNN）は驚異的な成長を遂げています。この成長は、より正確なモデルや生成AIによって新たな可能性を開拓することで見られます。大規模な言語モデル（LLM）による自然言語合成、テキストから画像を生成するモデルなど、これらのDNNの能力の向上は、トレーニングに膨大な計算リソースを必要とする巨大なモデルによるものです。分散トレーニングは、データ並列化とモデル並列化という2つの技術でこの問題に対処します。データ並列化は、トレーニングプロセスを複数のノードとワーカーにスケールさせるために使用されます。モデル並列化はモデルを分割し、指定されたインフラストラクチャに適合させます。Amazon SageMakerの分散トレーニングジョブは、1クリック（または1つのAPI呼び出し）で分散コンピュートクラスタを設定し、モデルをトレーニングし、結果をAmazon Simple Storage Service（Amazon S3）に保存し、完了時にクラスタをシャットダウンすることができます。さらに、SageMakerは、異種クラスタとデータ並列化およびモデル並列化のための分散トレーニングライブラリなどの機能を立ち上げることで、分散トレーニング領域での革新を継続的に行っています。

分散環境での効率的なトレーニングには、ハイパーパラメータの調整が必要です。複数のGPUでトレーニングする場合、バッチ（またはミニバッチ）サイズをGPUの数で乗算して、同じバッチサイズを各GPUで保つのが一般的な方法です。しかし、ハイパーパラメータの調整はモデルの収束に影響を与えることがよくあります。したがって、分散トレーニングでは、分散、ハイパーパラメータ、モデルの精度の3つの要素をバランスさせる必要があります。

この記事では、分散トレーニングが収束に与える影響と、Amazon SageMakerの自動モデルチューニングを使用して、データ並列化を使用した分散トレーニングのためのモデルのハイパーパラメータを微調整する方法について探究します。

この記事で言及されているソースコードは、GitHubリポジトリで見つけることができます（m5.xlargeインスタンスが推奨されています）。

• 自動化された進化が厳しい課題に取り組む
• 「抗議を取り締まっているロシアは、監視供給チェーンを構築しています」
• エロン・マスクのxAIがOpenAIのChatGPTに挑戦します

シングルから分散環境へのトレーニングのスケーリング

データ並列化は、トレーニングプロセスを複数のコンピューティングリソースにスケールさせ、より高速なトレーニング時間を実現する方法です。データ並列化では、データはコンピューティングノード間で分割され、各ノードは自分のパーティションに基づいて勾配を計算し、モデルを更新します。これらの更新は、非同期、一対多、または全対全の方法で、1つまたは複数のパラメータサーバを使用して行うことができます。別の方法はAllReduceアルゴリズムを使用することです。たとえば、リングオールリデュースアルゴリズムでは、各ノードは隣接する2つのノードとのみ通信し、データ転送を全体的に減らすことができます。パラメータサーバとリングオールリデュースについて詳しくは、「Horovodを使用したTensorFlow分散トレーニングの簡単な起動」または「Amazon SageMakerのパラメータサーバ」を参照してください。データの分割に関しては、n個のコンピューティングノードがある場合、各ノードはデータの約1/nのサイズのサブセットを取得する必要があります。

トレーニングのスケーリングの効果を示すために、次の2つのシンプルな実験を実行します：

• MXNetとGluonフレームワークを使用して、完全に接続されたDNNとReLU活性化関数を使用した画像分類モデルをトレーニングします。トレーニングデータには、手書き数字のMNISTデータセットを使用しました。SageMakerの例リポジトリで提供されるソースコードを使用しました。
• SageMakerの組み込みXGBoostアルゴリズムを使用して、バイナリ分類モデルをトレーニングします。銀行の顧客が特定のオファーに応答する可能性を予測するため、直接マーケティングデータセットを使用しました。この実験のソースコードと再現手順は、GitHubリポジトリで見つけることができます。

各モデルのトレーニングは、単一のインスタンスと複数のインスタンスで実行されました。DNNの分散トレーニングでは、分散プロセッサを最大限に活用するために、ミニバッチサイズをインスタンス数（4つ）で乗算しました。以下の表は、セットアップと結果をまとめたものです。

問題の種類	画像分類	バイナリ分類
モデル	DNN	XGBoost
*インスタンス**	ml.c4.xlarge	ml.m5.2xlarge
データセット	MNIST (ラベル付き画像)	直接マーケティング（表形式、数値およびベクトル化されたカテゴリ）
*検証メトリクス**	精度	AUC
エポック/ラウンド	20	150
インスタンス数	1	4	1	3
分布タイプ	N/A	パラメータサーバ	N/A	AllReduce
トレーニング時間（分）	8	3	3	1
最終的な検証スコア	0.97	0.11	0.78	0.63

両モデルにおいて、分散係数によってトレーニング時間がほぼ線形に短縮されました。しかし、モデルの収束性は著しく低下しました。この振る舞いは、2つの異なるモデル、異なるコンピュートインスタンス、異なる分散方法、および異なるデータ型に対して一貫しています。では、なぜトレーニングプロセスの分散はモデルの精度に影響を与えたのでしょうか？

この効果を説明しようとするいくつかの理論があります：

• テンソルの更新が大きい場合、ワーカーとパラメーターサーバー間のトラフィックが混雑する可能性があります。そのため、非同期パラメーターサーバーは重みの更新に遅延が生じるため、収束性が著しく悪化します [1]。
• バッチサイズを増やすと、過学習や一般化の不足につながる可能性があり、検証精度が低下します [2]。
• モデルパラメーターを非同期に更新する場合、一部のDNNは最新のモデルの重みを使用していないかもしれません。そのため、数回の反復遅れた重みに基づいて勾配を計算することになります。これにより、重みの遅れ [3] が生じ、さまざまな理由で引き起こされる可能性があります。
• 一部のハイパーパラメーターはモデルやオプティマイザーに固有です。たとえば、XGBoostの公式ドキュメントによれば、tree_modeハイパーパラメーターのexact値は分散トレーニングをサポートしていません。XGBoostは行分割データ分布を採用しているため、exactツリーメソッドはソートされた列形式で動作します。
• 一部の研究者は、より大きなミニバッチを設定すると、勾配の確率性が減少する可能性があると提案しています。これは、損失関数に局所的な最小値や鞍点が含まれており、ステップサイズに変更がない場合、最適化がこのような局所的な最小値や鞍点に固執することがあるためです [4]。

分散トレーニングの最適化

ハイパーパラメーターの最適化（HPO）は、学習アルゴリズムに最適なハイパーパラメーターセットを探索および選択するプロセスです。SageMakerの自動モデルチューニング（AMT）は、提供されたデータセット上で複数のトレーニングジョブを実行することで、HPOを管理されたサービスとして提供します。SageMaker AMTは、指定したハイパーパラメーターの範囲を検索し、選択したメトリックに基づいて最適な値を返します。ビルトインアルゴリズムを使用するか、独自のアルゴリズムとコンテナを使用してSageMaker AMTを利用できます。

ただし、分散トレーニングの最適化は一般的なHPOとは異なります。なぜなら、単一のインスタンスを起動する代わりに、各ジョブは実際にはインスタンスのクラスターを起動するからです。これはコストに大きな影響を与えます（特にDNN向けの高価なGPUアクセラレートインスタンスを考慮すると）。AMTの制限に加えて、トレーニングインスタンスの同時数に関するSageMakerのアカウント制限に達する可能性もあります。最後に、クラスターの起動により、起動時間が長くなるため、運用上のオーバーヘッドが発生する可能性があります。SageMaker AMTには、これらの問題に対処するための特定の機能があります。早期停止付きのHyperbandは、性能の良いハイパーパラメーター構成を微調整し、性能が低いものは自動的に停止します。これにより、トレーニング時間の効率的な利用と不要なコストの削減が可能です。また、SageMaker AMTはAmazon EC2 Spot Instancesの使用を完全にサポートしており、オンデマンドインスタンスに比べてトレーニングコストを最大90%削減できます。起動時間の長さに関しては、SageMaker AMTは各チューニングジョブ内でトレーニングインスタンスを自動的に再利用するため、各トレーニングジョブの平均起動時間が20倍短縮されます。また、関連するハイパーパラメーター、適切な範囲とスケール、最適な同時トレーニングジョブ数を選択し、結果を再現するためにランダムシードを設定するなど、AMTのベストプラクティスに従う必要があります。

次のセクションでは、先に述べたXGBoostの例を使用して、AMTジョブを構成、実行、分析する方法を説明します。

チューニングジョブの構成、実行、分析

先にも述べたように、ソースコードはGitHubリポジトリで入手できます。ステップ1から5では、データのダウンロードと準備、xgb3エスティメーターの作成（分散XGBoostエスティメーターは3つのインスタンスを使用するように設定されています）、トレーニングジョブの実行、結果の観察を行います。このセクションでは、既にステップ1から5を実行したものとして、そのエスティメーターのためのチューニングジョブを設定する方法について説明します。

チューニングジョブは、起動するトレーニングジョブのための最適なハイパーパラメーターを計算するため、パフォーマンスを評価するメトリックを使用します。独自のメトリックを設定することもできますが、SageMakerは、設定した正規表現に基づいてパースし、stdoutに出力するメトリックを使用します。また、SageMakerのビルトインアルゴリズムのメトリックを使用することもできます。この例では、ビルトインのXGBoostオブジェクティブメトリックを使用するため、正規表現の設定は必要ありません。モデルの収束を最適化するために、検証AUCメトリックに基づいて最適化を行います：

objective_metric_name="validation:auc"

私たちは以下の7つのハイパーパラメータをチューニングします：

• num_round – トレーニング中のブースティングのラウンド数。
• eta – 過学習を防ぐために更新時に使用されるステップサイズの縮小。
• alpha – 重みのL1正則化項。
• min_child_weight – 子ノードに必要なインスタンスの重み（ヘシアン）の最小合計。ツリーパーティションのステップによってインスタンスの重みの合計がmin_child_weightよりも小さい葉ノードが生成された場合、それ以上の分割は行われません。
• max_depth – ツリーの最大深度。
• colsample_bylevel – 各スプリットの各レベルでの列のサブサンプルの比率。このサブサンプリングは、ツリーの新しい深度レベルに到達するたびに1回行われます。
• colsample_bytree – 各ツリーの構築時の列のサブサンプルの比率。構築される各ツリーごとに、サブサンプリングが1回行われます。

XGBoostのハイパーパラメータについて詳しくは、「XGBoostハイパーパラメータ」を参照してください。以下のコードは、7つのハイパーパラメータとその範囲を示しています：

hyperparameter_ranges = {
    "num_round": IntegerParameter(100, 200),
    "eta": ContinuousParameter(0, 1),
    "min_child_weight": ContinuousParameter(1, 10),
    "alpha": ContinuousParameter(0, 2),
    "max_depth": IntegerParameter(1, 10),
    "colsample_bylevel": ContinuousParameter(0, 1),
    "colsample_bytree": ContinuousParameter(0, 1),
}

次に、Hyperband戦略とSageMaker SDKを使用したチューナーオブジェクトの設定を提供します。HyperbandStrategyConfigは2つのパラメータを使用できます：max_resource（オプション）は目的を達成するためにトレーニングジョブに使用される最大イテレーション数で、min_resourceはトレーニングを停止する前にトレーニングジョブによって使用される最小イテレーション数です。HyperbandStrategyConfigを使用してStrategyConfigを設定し、それは後でチューニングジョブの定義で使用されます。以下のコードを参照してください：

hsc = HyperbandStrategyConfig(max_resource=30, min_resource=1)
sc = StrategyConfig(hyperband_strategy_config=hsc)

これで、HyperparameterTunerオブジェクトを作成します。このオブジェクトには、以下の情報が渡されます：

• XGBoost estimator（3つのインスタンスで実行するように設定されています）
• 目的メトリック名と定義
• ハイパーパラメータの範囲
• トレーニングジョブの実行総数や並行して実行できるトレーニングジョブの数などのチューニングリソースの設定
• Hyperbandの設定（前のステップで設定した戦略と構成）
• Early stopping（early_stopping_type）をOffに設定

なぜEarly stoppingをOffに設定するのですか？トレーニングジョブは、ハイパーパラメータチューニングジョブの目的メトリックの改善の可能性が低い場合、早期に停止することがあります。これにより、計算時間を短縮し、モデルの過学習を防ぐことができます。ただし、Hyperbandは早期停止を適用するための高度な組み込みメカニズムを使用します。したがって、Hyperbandの内部早期停止機能を使用する場合、パラメータearly_stopping_typeはOffに設定する必要があります。以下のコードを参照してください：

tuner = HyperparameterTuner(
    xgb3,
    objective_metric_name,
    hyperparameter_ranges,
    max_jobs=30,
    max_parallel_jobs=4,
    strategy="Hyperband",
    early_stopping_type="Off",
    strategy_config=sc
)

最後に、fitメソッドを呼び出すことで、自動モデルチューニングジョブを開始します。非同期でジョブを開始したい場合は、waitをFalseに設定します。以下のコードを参照してください：

tuner.fit(
{"train": s3_input_train, "validation": s3_input_validation},
include_cls_metadata=False,
wait=True,
)

SageMakerコンソールでジョブの進捗状況と概要をフォローすることができます。ナビゲーションパネルでTrainingを選択し、Hyperparameter tuning jobsを選択し、関連するチューニングジョブを選択します。以下のスクリーンショットは、トレーニングジョブのステータスとパフォーマンスの詳細を示しています。

チューニングジョブが完了したら、結果を確認することができます。ノートブックの例では、SageMaker SDKを使用して結果を抽出する方法を示しています。まず、チューニングジョブがモデルの収束をどのように向上させたかを調べます。ジョブ名を使用してHyperparameterTunerオブジェクトをアタッチし、describeメソッドを呼び出すことができます。このメソッドは、チューニングジョブのメタデータと結果を含む辞書を返します。

以下のコードでは、目的のメトリック（検証AUC）によって測定された最も優れたトレーニングジョブの値を取得しています：

tuner = HyperparameterTuner.attach(tuning_job_name=tuning_job_name)
tuner.describe()["BestTrainingJob"]["FinalHyperParameterTuningJobObjectiveMetric"]["Value"]

結果は、検証セットでのAUCが0.78です。これは初期の0.63に比べて大幅な改善です！

次に、トレーニングジョブの実行時間を見てみましょう。そのために、SDKのHyperparameterTuningJobAnalyticsメソッドを使用してチューニングジョブの結果を取得し、分析および可視化のためにPandasデータフレームに読み込みます：

tuner_analytics = sagemaker.HyperparameterTuningJobAnalytics(tuning_job_name)
full_df = tuner_analytics.dataframe()
full_df.sort_values(by=["FinalObjectiveValue"], ascending=False).head()

Hyperband戦略によるトレーニングジョブの平均時間を見てみましょう：

full_df["TrainingElapsedTimeSeconds"].mean()

平均時間は約1分でした。これは、パフォーマンスが低いトレーニングジョブを早期に停止するHyperband戦略のメカニズムと一致しています。コストの面では、トレーニングジョブは合計30分のトレーニング時間に対して請求されました。Hyperbandの早期停止を行わない場合、請求可能なトレーニング時間は90分になる予定でした（30ジョブ * 1ジョブあたり1分 * 3インスタンス）。これはコスト削減において3倍の改善です！最後に、トレーニングジョブが30個のトレーニングジョブを実行し、合計12分かかったことがわかります。これは、予想される時間（30ジョブ / 4ジョブ並列 * 1ジョブあたり3分）のほぼ50％少ないです。

結論

この記事では、分散環境でモデルをトレーニングする際に観察された収束の問題について説明しました。SageMaker AMTを使用したHyperbandは、最適化されていないシーケンシャルなジョブよりも収束（10％以上の改善）、操作効率（チューニングジョブは半分の時間で実行されました）、およびコスト効率（トレーニングジョブ時間の90分に対して30分の請求）に関して主な懸念点を解決しました。以下の表は、結果をまとめたものです：

改善メトリック	チューニングなし/ナイーブモデルチューニング実装	SageMaker Hyperband自動モデルチューニング	測定された改善
モデルの品質（検証AUCによる評価）	0.63	0.78	15%
コスト（請求可能なトレーニング時間による評価）	90	30	66%
操作効率（合計実行時間による評価）	24	12	50%

スケーリング（クラスタサイズ）に関する微調整を行うためには、複数のクラスタ構成で調整ジョブを繰り返し、結果を比較して、速度とモデルの精度の両方を満たす最適なハイパーパラメータを見つけることができます。

この方法をノートブックの最後のセクションに含めました。

参考文献

[1] Lian, Xiangru, et al. “Asynchronous decentralized parallel stochastic gradient descent.” International Conference on Machine Learning. PMLR, 2018.

[2] Keskar, Nitish Shirish, et al. “On large-batch training for deep learning: Generalization gap and sharp minima.” arXiv preprint arXiv:1609.04836 (2016).

[3] Dai, Wei, et al. “Toward understanding the impact of staleness in distributed machine learning.” arXiv preprint arXiv:1810.03264 (2018).

[4] Dauphin, Yann N., et al. “Identifying and attacking the saddle point problem in high-dimensional non-convex optimization.” Advances in neural information processing systems 27 (2014).

We will continue to update VoAGI; if you have any questions or suggestions, please contact us!