大規模AIトレーニングおよび推論システムにおける大規模言語モデル(LLM)の性能を最適化するためには、分散ワークロード実行をモデル化するスケーラブルで表現力豊かなメカニズムが必要です。既存のプラットフォームから得られる実行トレースは、将来の大規模システム構成を研究する際に適応しにくいという課題があります。この記事では、Symbolic Tensor grAph GEnerator(STAGE)というフレームワークを紹介しています。STAGEは高忠実度の実行トレースを合成し、LLMワークロードを正確にモデル化します。また、ユーザーがLLMアーキテクチャやシステム構成の幅広いスペクトルを体系的に探求できる並列化戦略の包括的なセットをサポートしています。STAGEは、32K以上のGPUを超える高忠実度のLLMトレースを合成し、計算、メモリ、通信のテンソルレベルの精度を維持するスケーラビリティを示しました。今後、分散機械学習システムの研究を促進するために公開される予定です。