Amazon SageMaker AI でベンチマークとレコメンデーションの結果を MLflow にストリーミングする

生成AIモデルのデプロイ最適化において、チームはしばしば数十のGPUインスタンスタイプ、サーバーコンテナ、並列化戦略、投機的デコードなどの最適化手法を評価する必要があります。従来の手作業によるデータ収集と比較は時間がかかり、データサイロが発生しやすいという課題がありました。Amazon SageMaker AIの新しいMLflow統合により、ベンチマークとレコメンデーションの結果を自動的に統一されたSageMaker MLflowアプリにストリーミングできるようになり、データサイロを削減し、反復サイクルを加速し、実験の再現性を確保します。

この統合の主な利点は以下の通りです：

• 手動データ統合の排除：複数のジョブの結果が自動的に同じ実験名に統合され、MLflow実験ビューで異なる実行を並べて比較できます。例えば、qwen2-0.5bモデルをml.g4dn.12xlargeとml.p4d.24xlargeで比較できます。
• 長時間ジョブのリアルタイム監視：ベンチマークやレコメンデーションジョブは数時間かかることがありますが、レイテンシーやスループットのメトリクスがテスト中にライブで更新されるため、期待したスループットが得られない場合は早期にジョブを停止できます。
• 完全な監査証跡：各実行はジョブパラメータ、タイムスタンプ、チェックポイントメトリクス、生成されたアーティファクトをキャプチャし、数か月にわたってクエリ可能で再現可能です。構成変更による改善点を追跡し、パフォーマンスの停滞を特定できます。
• コラボレーションとガバナンスの向上：共有MLflow実験は最適化作業の単一の真実の情報源となり、重複作業を削減し、チームメンバー間のスムーズな引き継ぎをサポートします。

技術的には、まずAmazon SageMaker StudioでSageMaker MLflowアプリを作成し、実行ロールに必要な権限（sagemaker-mlflow:*やエンドポイント呼び出し権限を含む）を付与します。ベンチマークまたはレコメンデーションジョブを作成する際に、OutputConfigのMlflowConfigでMLflowアプリARN、実験名、実行名を指定します。ベンチマークジョブは既存のエンドポイントを評価し、レコメンデーションジョブは内部でエンドポイントを自動的にプロビジョニングして評価し、ランク付けされた構成を返します。例えば、Qwen/Qwen2-0.5B-Instructモデルを使用するチームは、同じMLflow実験で両方のジョブを追跡でき、エンドツーエンドのウォークスルーは通常45〜120分で完了します。

重要な注意点として、この統合はSageMaker MLflowアプリのみをサポートし、セルフホストのMLflowトラッキングサーバーにはストリーミングしません。tooling.versionを0.8.0以降に設定してMLflowネストランのサポートを有効にしてください。レコメンデーションジョブは内部でエンドポイントをプロビジョニングするため、特別な理由がない限りComputeSpec.InstanceTypesを指定しないでください。実行ロールには必要なSageMaker、MLflow、S3、エンドポイント呼び出しの権限が含まれている必要があります。S3出力バケットはジョブと同じリージョンにある必要があります。

この新しいMLflow統合により、チームは生成AIモデルのデプロイ最適化プロセスを大幅に簡素化し、データ整理ではなくモデル性能の向上に集中できます。