なぜチェックポイント間隔が重要なのか
AIモデルの学習、特に大規模言語モデル(LLM)のトレーニングにおいて、チェックポイント(checkpoint)は生命線です。学習中にハードウェア障害やプリエンプション(preemption)が発生すると、それまでの計算結果がすべて失われる可能性があります。
しかし、チェックポイントの保存は単に「頻繁に行えば良い」というものではありません。保存間隔が短すぎるとI/O帯域を過剰に消費し、実際の学習速度が低下します。逆に長すぎると、障害発生時に失われる学習量(例：100ステップ分)が大きくなります。この従来のトレードオフ(trade-off)を解決するために、GoogleのOrbaxチームが提案したのが**連続チェックポインティング(Continuous Checkpointing)**です。
連続チェックポインティングの動作原理
従来のチェックポインティングは、一定ステップ(例：100ステップ)ごとに同期的(synchronous)に保存していました。一方、連続チェックポインティングは前回の保存が完了したら即座に次の保存を非同期的(asynchronous)に開始します。Orbaxが内部で保存キュー(queue)を管理し、ホストマシンとI/O帯域を最大限に活用しながら、性能低下を最小限に抑えます。
MaxTextでの有効化方法
MaxTextを使用している場合、たった一行の設定で連続チェックポインティングを有効にできます。
# 学習ジョブ設定ファイル (例: config.yaml)
enable_checkpointing: True # チェックポインティング有効化
async_checkpointing: True # 非同期チェックポインティング有効化
enable_continuous_checkpointing: True # 連続チェックポインティング有効化 (核心)
max_num_checkpoints_to_keep: 10 # 最新10個のみ保持 (ディスク節約)
この設定により、MaxTextは前回のチェックポイント保存リクエストがバックグラウンドで完了したら即座に新しい保存を試みます。実際のベンチマーク結果を見ると、連続チェックポインティングを有効にした場合、P50チェックポイント間隔が顕著に減少します。平均学習ステップ時間は若干増加しますが、障害時の復旧コストを考慮すると、全体のGoodput(有効スループット)は大幅に向上します。
Orbaxの高度なポリシー：細かい制御
MaxTextの基本設定に加えて、Orbaxはより柔軟なポリシー(policy)を提供します。特に軽量モデルを学習する場合、チェックポイントが頻発し過ぎて不要なI/Oオーバーヘッドが発生することがあります。その場合はminimum_interval_secsを設定してクールダウン時間を設けることができます。
# 最小30秒間隔でチェックポイント保存 (過度な保存を防止)
continuous_checkpointing_policy_with_minimum_interval = \
save_decision_policy.ContinuousCheckpointingPolicy(minimum_interval_secs=30)
# 180秒ごとに最低1つのチェックポイントを保持するポリシー
every_n_seconds_preservation_policy = \
preservation_policy.EveryNSeconds(180)
また、独自の保存ポリシーを定義することも可能です。
@dataclasses.dataclass
class CustomizedPreservationPolicy(PreservationPolicy):
"""チェックポイント保存のためのカスタムポリシー"""
def should_preserve(
self,
checkpoints: Sequence[PolicyCheckpointInfo],
*, context: PreservationContext,
) -> Sequence[bool]:
result = [is_checkpoint_preservable(cp) for cp in checkpoints]
_log_preservation_decision("カスタム保存ポリシー", checkpoints, result)
return result
注意点と制限
連続チェックポインティングは万能ではありません。特に以下の2点を必ず確認してください。
-
ストレージバケットと学習クラスタの物理的な配置: 連続チェックポインティングはネットワーク帯域に大きく依存します。ストレージが学習クラスタと同じデータセンター内(co-located)にある必要があります。クロスメトロ(cross-metro)ネットワークを使用するとチェックポインティング速度が著しく低下し、かえって信頼性を損なう可能性があります。
-
マルチスライス環境でのDCN帯域: 複数スライス(slice)にまたがって学習する場合、DCN(Data Center Network)帯域はモデル重み更新とチェックポインティングの両方で使用されます。幸い、Orbaxはチェックポインティングの主要負荷を単一スライス(主にスライス0)とストレージサーバ間の通信に制限し、スライス間通信をブロックしない設計になっています。ベンチマークの結果、マルチスライス環境でも連続チェックポインティングによる有意な速度低下は観測されていません。
実務適用のアドバイス：チェックポイント戦略の策定
連続チェックポインティングは「設定して忘れる(set-and-forget)」ソリューションに近いものです。しかし、完璧な戦略のためには以下を考慮しましょう。
- MTBF(Mean Time Between Failure)の推定: 自身のクラスタで予想される平均障害間隔を把握してください。この値に基づいてチェックポイント間隔を調整すると、最適なGoodputが得られます。
- ストレージ費用 vs. 復旧費用: チェックポイントを保持しすぎるとストレージコストが増加します。
max_num_checkpoints_to_keepを適切に設定し、コストと信頼性のバランスを取りましょう。 - モニタリング: チェックポイントの保存時間と間隔を継続的に監視し、ネットワークやストレージのボトルネックが発生していないか確認してください。
この技術は特に大規模学習(数百〜数千のチップ)で真価を発揮します。小規模クラスタでは従来の定期的チェックポインティングと大きな差はないかもしれませんが、スケールが大きくなるほど連続チェックポインティングの効率は最大化されます。
合わせて読みたい記事
次のステップとしての学習方向
連続チェックポインティングを超えて、さらに高度なトピックに興味がある場合は以下をご覧ください。
- Orbaxのカスタム保存/復元戦略: データ並列、パイプライン並列など様々な並列化戦略に応じたチェックポイント最適化。
- 異なるハードウェア(TPU v4, v5p, GPU)での性能比較: 各環境に最適な設定を見つける。
- Fault-tolerant trainingの最新研究動向: チェックポイントなしで障害を復旧する手法(例：elastic training)。
Google Cloud公式ブログの原文で、より詳細なベンチマークデータを確認できます。