「Amazon SageMakerを使用してビジョントランスフォーマーモデルのトレーニング時間を短縮するKTの取り組み」

「ビジョントランスフォーマーモデルのトレーニング時間を短縮するため、KTがAmazon SageMakerを活用」

KT Corporationは韓国で最大の通信事業者の一つであり、固定電話、モバイル通信、インターネット、AIサービスなど幅広いサービスを提供しています。KTのAI Food Tagは、コンピュータビジョンモデルを使用して写真の中の食品の種類と栄養成分を識別するAIベースの食事管理ソリューションです。KTによって開発されたこのビジョンモデルは、大量の未ラベルの画像データで事前トレーニングされたモデルに依存して、さまざまな食品の栄養成分とカロリー情報を分析します。AI Food Tagは、糖尿病などの慢性疾患を持つ患者が食事管理をサポートすることができます。KTは、このAI Food Tagモデルのトレーニングを従来よりも29倍高速化し、モデル蒸留技術を用いて本番環境への展開を最適化するためにAWSとAmazon SageMakerを使用しました。この記事では、KTのモデル開発の道のりとSageMakerの成功を紹介します。

K社プロジェクトの紹介と課題の定義

K社によって事前トレーニングされたAI Food Tagモデルは、ビジョン変換器（ViT）アーキテクチャを基にしており、従来のビジョンモデルよりもモデルパラメータが多く、精度を向上させるために設計されています。本番環境でのモデルサイズを縮小するために、KTは知識蒸留（KD）技術を使用してモデルパラメータの数を減らし、精度に重大な影響を与えることなく設定することができます。知識蒸留では、事前トレーニングされたモデルは教師モデルと呼ばれ、軽量な出力モデルは生徒モデルとしてトレーニングされます。以下の図に示すように、軽量な生徒モデルは教師モデルよりもモデルパラメータが少なく、メモリ要件を削減し、より小さな、より低コストなインスタンスに展開することができます。生徒モデルは教師モデルの出力から学習することにより、受け入れられる精度を維持します。

知識蒸留では、教師モデルはKD中に変化せず、生徒モデルは教師モデルの出力ロジットをラベルとして使用して損失を計算するためにトレーニングされます。このKDパラダイムでは、教師と生徒の両方が単一のGPUメモリ上でトレーニングする必要があります。KTは初めに、内部のオンプレミス環境で2つのGPU（A100 80 GB）を使用して生徒モデルをトレーニングしましたが、そのプロセスには約40日かかりました。トレーニングを加速し、より少ない時間で生徒モデルを生成するために、KTはAWSと提携しました。両チームはモデルのトレーニング時間を大幅に削減しました。この記事では、チームが軽量なAI Food Tagモデルを開発するためにAmazon SageMaker Training、SageMaker Data Parallelism Library、Amazon SageMaker Debugger、Amazon SageMaker Profilerをどのように使用したかについて説明します。

• テンセントAIラボは、検索補完された言語モデルの堅牢性と信頼性を高めるために、Chain-of-Noting（CoN）を導入します
• 量子コンピュータを使ってより高度な機械学習モデル
• 「Juliaプログラミング言語の探索：アプリケーションプログラミングインターフェース（API）—パート1」

SageMakerを使用して分散トレーニング環境を構築する

SageMaker Trainingは、AWS上の管理された機械学習（ML）トレーニング環境であり、トレーニングの体験を簡素化するためのさまざまな機能とツールを提供しています。また、以下の図に示すように、分散コンピューティングで役立つことがあります。

SageMakerのお客様は、モデルトレーニングのために事前にインストールされたさまざまなディープラーニングフレームワークと必要なLinux、NCCL、Pythonパッケージが組み込まれたDockerイメージにアクセスすることもできます。モデルトレーニングを実行したいデータサイエンティストやMLエンジニアは、トレーニングインフラストラクチャの設定やDockerの管理、さまざまなライブラリの互換性を気にすることなくモデルトレーニングを行うことができます。

1日間のワークショップ中に、KTのAWSアカウント内でSageMakerをベースとした分散トレーニング構成をセットアップすることができ、SageMaker Distributed Data Parallel (DDP) ライブラリを使用してKTのトレーニングスクリプトを高速化し、さらに2つのml.p4d.24xlargeインスタンスを使用してトレーニングジョブをテストすることもできました。このセクションでは、KTがAWSチームと協力し、SageMakerを使用してモデルを開発した経験について説明します。

概念実証では、分散トレーニング時にAWSインフラに最適化されたSageMaker DDPライブラリを使用してトレーニングジョブを高速化したかったのです。PyTorch DDPからSageMaker DDPに変更するには、単にtorch_smddpパッケージを宣言し、バックエンドをsmddpに変更するだけです。以下のコードを参照してください：

import smdistributed.dataparallel.torch
torch_smddpdist.init_process_group(backend='smddp',rank=args.rank,world_size=args.world_size)

SageMaker DDPライブラリの詳細については、SageMakerのデータ並列処理ライブラリを参照してください。

SageMakerデバッガーとプロファイラを使用して遅いトレーニング速度の原因を分析する

トレーニングワークロードを最適化して高速化するための最初のステップは、ボトルネックが発生する場所を理解し診断することです。KTのトレーニングジョブでは、データローダ、フォワードパス、バックワードパスごとのトレーニング時間の計測を行いました：

1回のイテレーションでの時間 – データローダー：0.00053秒、フォワード：7.77474秒、バックワード：1.58002秒

2回のイテレーションでの時間 – データローダー：0.00063秒、フォワード：0.67429秒、バックワード：24.74539秒

3回のイテレーションでの時間 – データローダー：0.00061秒、フォワード：0.90976秒、バックワード：8.31253秒

4回のイテレーションでの時間 – データローダー：0.00060秒、フォワード：0.60958秒、バックワード：30.93830秒

5回のイテレーションでの時間 – データローダー：0.00080秒、フォワード：0.83237秒、バックワード：8.41030秒

6回のイテレーションでの時間 – データローダー：0.00067秒、フォワード：0.75715秒、バックワード：29.88415秒

各イテレーションの標準出力の時間を見ると、バックワードパスの実行時間がイテレーションごとに大きく変動していることがわかりました。この変動は異常であり、トレーニング時間全体に影響を与える可能性があります。この一貫性のないトレーニング速度の原因を見つけるために、まずはリソースのボトルネックを特定するためにSystem Monitor（SageMakerデバッガーUI）を利用しました。これにより、SageMakerトレーニング上のトレーニングジョブをデバッグし、指定された秒数以内で管理されたトレーニングプラットフォームのCPU、GPU、ネットワーク、I/Oなどのリソースのステータスを表示することができます。

SageMakerデバッガーUIは、トレーニングジョブのボトルネックを特定し診断するのに役立つ詳細なデータを提供します。具体的には、CPU利用率の折れ線グラフとインスタンスごとのCPU/GPU利用率のヒートマップが注目されました。

CPU利用率の折れ線グラフでは、一部のCPUが100%使用されていることがわかりました。

ヒートマップ（濃い色ほど利用率が高い）では、いくつかのCPUコアがトレーニング中ずっと高い利用率を示している一方、GPUの利用率は一定ではありませんでした。

ここから、トレーニングスピードの遅さの理由の一つがCPUボトルネックである可能性があると疑い始めました。CPUボトルネックを引き起こす要素があるかどうか、トレーニングスクリプトのコードを見直しました。最も疑わしい部分は、データローダー内のnum_workersの大きな値でしたので、CPU利用率を削減するためにこの値を0または1に変更しました。そして、再度トレーニングジョブを実行し、結果を確認しました。

以下のスクリーンショットは、CPU利用率のラインチャート、GPU利用率、そしてCPUボトルネックを軽減した後のヒートマップを示しています。

num_workersを単純に変更することで、CPU利用率が大幅に減少し、GPU利用率が全体的に向上しました。これはトレーニングスピードを大幅に改善した重要な変更でした。それでも、GPU利用率を最適化できる箇所を見つけたいと思いました。そのため、SageMaker Profilerを使用しました。

SageMaker Profilerは、GPUとCPUの利用率メトリクスやトレーニングスクリプト内のGPU/CPUのカーネル消費を追跡することで、最適化の手がかりを提供します。これにより、どの操作がリソースを消費しているのかをユーザーが把握することができます。まず、SageMaker SDKを使用してトレーニングジョブを起動する関数にProfilerConfigを追加する必要があります。以下のコードに示すように設定します:

from sagemaker import ProfilerConfig, Profilerfrom sagemaker.debugger import (ProfilerRule, rule_configs)rules=[ProfilerRule.sagemaker(rule_configs.ProfilerReport())]profiler_config = ProfilerConfig(profile_params = Profiler(cpu_profiling_duration=3600))from sagemaker.pytorch import PyTorchregion_name = 'us-west-2'image_uri=f'763104351884.dkr.ecr.{region_name}.amazonaws.com/pytorch-training:2.0.0-gpu-py310-cu118-ubuntu20.04-sagemaker'estimator = PyTorch(entry_point='train.py',source_dir='src',role=role,image_uri=image_uri,instance_count=4,instance_type='ml.p4d.24xlarge',distribution={'smdistributed': {'dataparallel': {'enabled': True}}},profiler_config=profiler_config,hyperparameters=hyperparameters,sagemaker_session=sagemaker_session,)

SageMaker Python SDKでは、SageMaker Profilerのためのannotate関数を追加することで、プロファイリングが必要なコードやステップをトレーニングスクリプトで選択する柔軟性があります。以下は、トレーニングスクリプトでSageMaker Profilerのために宣言する必要があるコードの例です:

import smppySMProf = smppy.SMProfiler.instance()config = smppy.Config()config.profiler = {"EnableCuda": "1",}SMProf.configure(config)SMProf.start_profiling()…with smppy.annotate("Forward"):student_out = student_model(inp)with smppy.annotate("Backward"):loss.backward()…SMProf.stop_profiling()

上記のコードを追加した後、トレーニングスクリプトを使用してトレーニングジョブを実行すると、一定時間経過後にGPUカーネルによって消費される操作の情報を取得できます（以下の図のように）。KTのトレーニングスクリプトの場合、1エポックで実行し、以下の結果を得ました。

SageMakerプロファイラの結果の中で、GPUカーネルのトップ5の操作消費時間を確認したところ、KTトレーニングスクリプトでは、最も時間がかかるのは一般的な行列乗算（GEMM）操作であることがわかりました。SageMakerプロファイラのこの重要な洞察を得た後、これらの操作を高速化し、GPUの利用率を向上させる方法について調査を開始しました。

トレーニング時間の短縮

行列の乗算の計算時間を短縮するさまざまな方法を検討し、2つのPyTorch関数を適用しました。

ZeroRedundancyOptimizerを使用してオプティマイザの状態を分割

Zero Redundancy Optimizer (ZeRO)を見ると、DeepSpeed/ZeROテクニックは、モデルが使用するメモリの冗長性を排除することにより、大規模なモデルの訓練を効率的に行い、より高速な訓練速度を実現します。PyTorchのZeroRedundancyOptimizerは、Distributed Data Parallel (DDP)内の各プロセスごとにメモリ使用量を削減するために、オプティマイザの状態をシャードする技術を使用します。DDPは後ろ向きパスで同期された勾配を使用し、すべてのオプティマイザのレプリカが同じパラメータと勾配値で反復するようにしますが、すべてのモデルパラメータを持つ代わりに、各オプティマイザの状態は異なるDDPプロセスごとにメモリ消費量を減らすためにシャーディングによって保持されます。

使用するには、既存のオプティマイザをoptimizer_classに残して、モデルの残りのパラメータと学習率をパラメータとしてZeroRedundancyOptimizerを宣言します。

student_optimizer = ZeroRedundancyOptimizer(student_model.parameters(),optimizer_class=torch.optim.AdamW,lr=initial_lr)

自動混合精度

自動混合精度 (AMP)では、一部の演算にはtorch.float32データ型を使用し、他の演算にはtorch.bfloat16またはtorch.float16を使用することで、高速な計算とメモリ使用量の削減を実現します。特に、深層学習モデルは通常、計算において指数部よりも分数部により敏感ですので、torch.bfloat16はtorch.float32の指数部に相当し、最小限の損失で迅速に学習することができます。torch.bfloat16は、ml.p4d.24xlarge、ml.p4de.24xlarge、およびml.p5.48xlargeなどのA100 NVIDIAアーキテクチャ（Ampere）以降のインスタンスでのみ実行されます。

AMPを適用するには、上記のコードのようにトレーニングスクリプトでtorch.cuda.amp.autocastと宣言し、dtypeをtorch.bfloat16として宣言します。

with torch.cuda.amp.autocast(dtype="torch.bfloat16"):teacher = teacher_model(input_data)student = student_model(input_data)loss = loss(teacher, student, target)loss.requires_grad_(True)loss.backward()student_optimizer.step()student_optimizer.zero_grad(set_to_none=True)

SageMakerプロファイラの結果

トレーニングスクリプトに2つの関数を適用し、エポックごとにトレーニングジョブを再実行した後、SageMakerプロファイラでGPUカーネルのトップ5の操作消費時間を確認しました。以下の図は結果を示しています。

2つのTorch関数を適用する前にリストのトップにあったGEMM操作は、トップ5の操作から消え、通常は分散トレーニングで発生するReduceScatter操作に置き換えられたことがわかります。

KT蒸留モデルのトレーニング速度の結果

私たちは2つのTorch関数を適用することによるメモリの節約を考慮し、トレーニングバッチサイズをさらに128増やしました。これにより、最終的なバッチサイズは1024ではなく1152になりました。最終的な学生モデルのトレーニングは、1日に210エポックを実行することができました。KTの内部トレーニング環境とSageMakerのトレーニング時間と速度向上は、以下の表にまとめられています。

トレーニング環境	トレーニングGPU仕様	GPUの数	トレーニング時間（時間）	エポック	エポックあたりの時間	削減率
KTの内部トレーニング環境	A100（80GB）	2	960	300	3.20	29
Amazon SageMaker	A100（40GB）	32	24	210	0.11	1

AWSのスケーラビリティにより、オンプレミスでは2つのGPUではなく32個のGPUを使用することで、トレーニングジョブを29倍高速化することができました。その結果、SageMaker上でより多くのGPUを使用することにより、総合的なトレーニングコストに差はないまま、トレーニング時間を大幅に短縮することができます。

結論

KTの収束技術センターのAI2XLラボのビジョンAIサービス技術チームリーダーであるパク・サンミン氏は、AWSとの協力について次のようにコメントしています：

「最近、ビジョン領域でトランスフォーマーベースのモデルが増えてきており、モデルのパラメータと必要なGPUメモリが増加しています。私たちはこの問題を解決するために軽量技術を使用しており、学習には約1か月かかります。AWSとのこのPoCを通じて、SageMaker ProfilerとDebuggerの助けを借りてリソースのボトルネックを特定し、それらを解決し、最適化されたモデルコードを使用して4つのml.p4d.24xlargeインスタンス上で約1日でトレーニングを完了することができました。」

SageMakerは、サンミン氏のチームがモデルのトレーニングと開発にかかる数週間の時間を節約するのに役立ちました。

ビジョンモデルに関するこの協力を基に、AWSとSageMakerチームは、KTとのさまざまなAI/ML研究プロジェクトでさらなる協力を続け、SageMakerの機能を活用してモデルの開発とサービスの生産性を向上させる予定です。

SageMakerの関連機能について詳しくは、以下をご覧ください：

• 機械学習モデルをトレーニングする
• SageMakerのデータ並列処理ライブラリ
• Amazon SageMaker Debuggerを使用してモデルのデバッグとパフォーマンスの改善
• Amazon SageMaker Profilerを使用してAWSの計算リソース上のアクティビティをプロファイルする

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