ML グッドプット測定ライブラリでグッドプットをモニタリングする

ML グッドプット測定ライブラリ（ml-goodput-measurement）は、Cloud TPU VM で実行されている ML トレーニングワークロードの効率性を測定するのに役立つ Python パッケージです。このライブラリは、ワークロードグッドプットを測定する指標を提供します。これは、生産的なトレーニングの進捗状況を維持するために費やされた TPU 使用時間の割合です。一方、バッドプットは、起動オーバーヘッド、I/O 停止、中断復旧などの非生産的なアクティビティに費やされた合計時間の割合です。

Cloud Monitoring ダッシュボードと TensorBoard を使用して、スループット指標をリアルタイムで可視化できます。これにより、ボトルネックを特定し、リソース使用率を最適化して、最終的にトレーニング費用を削減できます。

詳細については、ML Goodput 測定 GitHub リポジトリをご覧ください。

スループット指標

ML Goodput Measurement ライブラリは、次の指標を提供します。これらの指標は、Cloud Monitoring と TensorBoard で表示することもできます。この表の指標は compute.googleapis.com/workload/ で始まりますが、たとえば、goodput_time の完全な指標名は compute.googleapis.com/workload/goodput_time です。

goodput_time: 生産的なトレーニングの合計時間（秒）。これは累積スループットと解釈できます。
badput_time: 生産性のないトレーニング時間の合計（秒単位）（起動、停止、復元）。これは、累積バッドプットとして解釈できます。
total_elapsed_time: ワークロードの経過時間（実時間）の合計（秒単位）。経過時間は、アプリケーションの開始時刻から現在時刻またはジョブの完了時刻までの時間です。
interval_goodput: 指定された期間（過去 24 時間など）のグッドプット率。この指標は、スループットのローリングウィンドウを提供します。
interval_badput: 指定された期間のバッドプット率。この指標は、バッドプットのローリングウィンドウを提供します。これは、I/O オペレーションの急増などの一時的な問題を特定するのに役立ちます。
disruptions: ジョブの中断の累積数。中断とは、トレーニングプロセスが予期せず停止し、再起動が必要になるイベントです。たとえば、ハードウェアの障害やメンテナンスイベントなどです。
step_time_deviation: トレーニングステップ時間のばらつきによる非生産的な時間の量。これは「ジッター」とも呼ばれます。この指標は、スパイクに敏感な安定性を測定します。これは、最近の「最悪のケース」のステップと過去のベースラインとの偏差を秒単位で表したものです。ステップは、トレーニングループの 1 回のイテレーションです。
performance: 最も速い安定したステップ時間の推定値（ベースライン）（秒単位）。理想的なステップ時間は、最適な条件下で、一時的なノイズやジッターの影響を受けずに、1 回のトレーニングステップにかかる最速の時間です。
max_productive_steps: 正常に保存された最大歩数。

インストール

次の手順に沿って、TPU ワークロードで ML Goodput Measurement ライブラリを設定します。

Cloud Logging API と Cloud Monitoring API を有効にします。
Google Kubernetes Engine（GKE）にデプロイする場合は、すべてのノードプールを cloud-platform アクセススコープで構成します。

トレーニングホストと分析マシンに ml-goodput-measurement パッケージをインストールします。

pip install ml-goodput-measurement

Goodput ライブラリでモニタリングする

ML グッドプット測定ライブラリを使用するには、GoodputRecorder インスタンスを初期化し、主要なセクションを record_event コンテキストマネージャーでラップしてトレーニングコードを計測し、GoodputMonitor で指標をリアルタイムでモニタリングします。モニターはバックグラウンドプロセスを実行して、記録されたイベントからスループット指標を定期的に計算し、リアルタイムの分析と可視化のために Cloud Monitoring と TensorBoard にアップロードします。

スループットレコーダーを初期化する

ML グッドプット測定ライブラリのコアコンポーネントである GoodputRecorder を初期化します。

import datetime
import jax
from ml_goodput_measurement import measurement

# Define a unique logger name for this specific run
logger_name = f'goodput_{config.run_name}'

# Instantiate the recorder
goodput_recorder = measurement.GoodputRecorder(
    job_name=config.run_name,
    logger_name=logger_name,
    logging_enabled=(jax.process_index() == 0)
)

イベントを記録する

record_event コンテキストマネージャーを使用して、トレーニングコードをラップします。

def train_loop(config):

  # 1. Wrap the entire Job (Start/End)
  with goodput_recorder.record_event(measurement.Event.JOB):

    # 2. Record Hardware Initialization
    with goodput_recorder.record_event(measurement.Event.ACCELERATOR_INIT):
      # ... perform device mesh setup ...
      initialize_tpu(config)

    # 3. Record Training Prep
    with goodput_recorder.record_event(measurement.Event.TRAINING_PREP):
      # ... create checkpoint managers, setup model ...
      model = training_prep(config)

    # 4. Main Training Loop
    for step in range(config.steps):

      # Record Data Loading
      with goodput_recorder.record_event(measurement.Event.DATA_LOADING):
        batch = get_next_batch()

      # Record Step Start (CRITICAL: Pass the step number!)
      with goodput_recorder.record_event(measurement.Event.STEP, step):
        output = train_step(model, batch)

      # 5. Record Custom Events (e.g., Evaluation)
      if step % eval_interval == 0:
        with goodput_recorder.record_event(measurement.Event.CUSTOM, "eval_step"):
          run_evaluation()

MaxText を使用してトレーニングコードを Goodput ライブラリと統合する方法の例については、goodput.py をご覧ください。

イベントをモニタリングする

GoodputMonitor を使用して指標をモニタリングします。これにより、ジョブの実行中に指標を計算してアップロードするバックグラウンドプロセスが開始されます。アップローダプロセスが意図したとおりに開始および停止するように、ロジックをコンテキストマネージャーでラップします。

構成とライフサイクル管理を処理するヘルパーコンテキストマネージャーを定義します。

import contextlib
from ml_goodput_measurement import monitoring

@contextlib.contextmanager
def maybe_monitor_goodput(config):
  """Monitor goodput if enabled and on the main process."""
  if not config.monitor_goodput or jax.process_index() != 0:
    yield
    return

  goodput_monitor = None
  try:
    # Configure GCPOptions for Cloud Monitoring
    gcp_options = monitoring.GCPOptions(
      enable_gcp_goodput_metrics=config.enable_gcp_goodput_metrics
    )

    # Instantiate the monitor
    goodput_monitor = monitoring.GoodputMonitor(
      job_name=config.run_name,
      logger_name=f"goodput_{config.run_name}",
      tensorboard_dir=config.tensorboard_dir,
      upload_interval=config.goodput_upload_interval_seconds,
      monitoring_enabled=True,
      pathway_enabled=config.enable_pathways_goodput,
      include_badput_breakdown=True,
      gcp_options=gcp_options,
    )

    # Start the background upload process
    goodput_monitor.start_goodput_uploader()
    print("Started Goodput upload to Tensorboard & GCM in the background!")
    yield

  finally:
    # Ensure clean shutdown of the background process
    if goodput_monitor:
      goodput_monitor.stop_goodput_uploader()
      print("Flushed final metrics and safe exited from Goodput monitoring.")

ローリングウィンドウのパフォーマンスを測定するには、start_rolling_window_goodput_uploader と stop_rolling_window_goodput_uploader を使用するようにモニターを構成します。

try:
    self._rolling_window_monitor.start_rolling_window_goodput_uploader(
        self.config.rolling_window_size
    )
finally:
    if self._rolling_window_monitor:
        self._rolling_window_monitor.stop_rolling_window_goodput_uploader()

メインのトレーニングエントリポイントをコンテキストマネージャーでラップします。これにより、モニタリングはトレーニングの開始前に開始され、トレーニングの終了時に最終的な指標がクリアされます。

def main():
  # ... Load configuration ...

  # Wrap the entire execution
  with maybe_monitor_goodput(config):
    # Run the training loop (which contains the GoodputRecorder events)
    train_loop(config)

後処理と分析

標準の CPU VM やノートパソコンなど、任意のマシンから完了したジョブのグッドプット指標を計算できます。後処理と分析に TPU を使用する必要はありません。

次のコードは、ジョブの総グッドプットを出力します。

from ml_goodput_measurement import goodput

calculator = goodput.GoodputCalculator(
    job_name="my-run-name",
    logger_name="goodput_my-run-name"
)

goodput, badput, last_step = calculator.get_job_goodput(
    include_badput_breakdown=True
)

print(f"Goodput: {goodput}%")
print(f"Badput (Infra Recovery): {badput[goodput.BadputType.INFRASTRUCTURE_RECOVERY_FROM_DISRUPTION]}%")

バッドプットタイプ（BadputType）については、バッドプットの内訳の詳細をご覧ください。

特定の期間を指定するには、get_job_goodput_interval を使用します。

goodput_pct, badput, _, _, _ = calculator.get_job_goodput_interval(
    start_time_utc,
    end_time_utc
)

開始時間（start_time_utc）と終了時間（end_time_utc）は datetime オブジェクトです。詳細については、goodput.py をご覧ください。

Goodput ダッシュボードでモニタリングする

ML トレーニングワークロードのモニタリングと可視化に役立つように、Google Cloud には GKE JobSet Goodput ダッシュボードと Cloud ML Goodput ダッシュボードの 2 つの Goodput ダッシュボードが用意されています。GKE JobSet ダッシュボードを使用してインフラストラクチャまたはスケジューリングの問題を診断し、Cloud ML Goodput ダッシュボードを使用してトレーニングコード内のボトルネックを特定します。

Google Cloud ML グッドプットダッシュボード

ML グッドプットダッシュボードは、トレーニングスクリプトのアプリケーションレベルの効率を測定します。生産的なトレーニングに費やされた時間や、データ読み込み、初期化、中断からの復元などのバッドプットのソースに関する分析情報を提供します。

Cloud Monitoring で ML Goodput ダッシュボードを表示するには:

Google Cloud コンソールで、[Cloud Monitoring] ページに移動します。
Monitoring コンソールに移動します
ナビゲーションパネルで [ダッシュボード] をクリックします。
[フィルタ検索] フィールドに「Cloud ML Goodput Dashboard」と入力します。

ML グッドプットダッシュボードには、次の指標（workload/ で始まる）が表示されます。

goodput_time: 生産的なトレーニングステップに費やされた累積時間。
badput_time: 生産性の低いアクティビティに費やされた累積時間と、バッドプットのソース:
- ACCELERATOR_INITIALIZATION: TPU の設定時間。
- TRAINING_PREP: チェックポイントの読み込み、モデルとオプティマイザの作成。
- PROGRAM_STARTUP: JIT コンパイル、グラフトレース。
- DATA_LOADING_SYNC: データ入力でブロックされた時間。
- CHECKPOINT_SAVE: モデルの状態を保存する時間。
- CHECKPOINT_RESTORE: モデルの状態を復元する時間。
- WASTED_PROGRESS: チェックポイントの前に発生した中断により失われた生産的な時間。
- INFRASTRUCTURE_RECOVERY: ジョブの再起動中のダウンタイム。
- CUSTOM_BADPUT_EVENTS: 評価などのユーザー定義の同期イベント。
step_time_deviation: トレーニングステップ時間のジッターと不安定さを測定します。
interval_goodput: ローリングウィンドウ（過去 1 時間など）で計算されたスループット。

GKE JobSet ダッシュボード

GKE JobSet Goodput ダッシュボードは、オーケストレーションレイヤの効率に焦点を当て、GKE が Pod を迅速にスケジュールして実行し続けているかどうかを把握するのに役立ちます。

JobSet の健全性とパフォーマンスに関する包括的な情報を表示するには、 Google Cloud コンソールの JobSet モニタリング ダッシュボードに移動します。

JobSet モニタリングダッシュボードに移動

JobSet モニタリングダッシュボードには、次の指標が表示されます。

kubernetes.io/jobset/scheduling_goodput: JobSet の作成時間に対する、JobSet の必要なすべてのリソース（Pod）が準備完了状態であり、作業に使用できる時間の割合。値が小さい場合は、GKE 内の Pod スケジューリング、イメージの pull、リソース割り当ての遅延を示します。
kubernetes.io/jobset/proxy_runtime_goodput: アクセラレータのデューティサイクルなどのシステムレベルのシグナルに基づいて、TPU がアクティブに使用されている時間の推定割合。これにより、アプリケーションのインストルメンテーションなしで、ランタイムの生産性の概要を把握できます。
ノードプールの指標: JobSet をホストするノードプールに影響する健全性、可用性、中断に関する情報。これにより、スループットの低下と基盤となるノードの問題を関連付けることができます。

GKE JobSet ダッシュボードの詳細については、JobSet のグッドプットをモニタリングするをご覧ください。

トラブルシューティングガイド

このセクションでは、ML Goodput 測定ライブラリの使用中に発生する可能性のある一般的な問題の特定と解決に役立つトラブルシューティング情報について説明します。

Cloud Monitoring の指標が見つからない

Cloud Monitoring の指標が見つからない場合:

IAM 権限を確認する: ワークロードに接続されているサービスアカウントに monitoring.timeSeries.create 権限と monitoring.metricDescriptors.create 権限が必要です。
構成を確認する: スループット指標は、構成で明示的に有効にする必要があります。たとえば、GCPOptions 内で enable_gcp_goodput_metrics=True が渡されていることを確認します。

一般的な指標が欠落している

一般的な指標が見つからない場合:

モニターのステータスを確認する: モニターのインスタンス化でモニタリングが有効になっている必要があります。monitoring_enabled=True をモニターに渡します。

破損した指標

指標が破損している場合や、奇妙な値が表示されている場合:

実行名が一意であることを確認する: すべての run_name 値は、すべてのテストで一意である必要があります。実行名を再利用すると、古い実行と新しい実行のログが混在します。

ログの欠落

Cloud Logging からログが欠落している場合:

Cloud Logging API を有効にする:Google Cloud プロジェクトで Cloud Logging API を有効にする必要があります。
レコーダーの設定を確認する: logging_enabled=True フラグを Goodput レコーダーに渡す必要があります。
メインプロセスを確認する: メインプロセス（jax.process_index() == 0）がアプリログをアクティブにレポートしている必要があります。

badput チェックポイントがない

badput チェックポイントが見つからない場合:

チェックポインティングライブラリを確認する: Goodput ライブラリは、Orbax のチェックポインティング時間を自動的に追跡します。ただし、他のライブラリ（PyTorch チェックポイントなど）を使用している場合は、record_custom_badput_event_start_time と record_custom_badput_event_end_time を使用して保存呼び出しと復元呼び出しを手動でラップする必要があります。