AMD Primus Tuning AgentがMI355Xクラスター向けの最適LLM設定を自動発見

マルチノードGPUクラスター上での大規模言語モデルの訓練を手動で設定することは、AIインフラにおける最もコストの高いエンジニアリングタスクの一つだ。パラメータの組み合わせ空間——並列化方式・スケジューラー・マイクロバッチサイズ・MoEバックエンド——は有効な設定が数万に及び、クラスター上での各フルパスには何時間ものGPU時間が必要だ。AMD ROCmチームはその空間を自動的に探索するツールPrimus Tuning Agentで応答した。

Primus Tuning Agent：LLM訓練設定の自動化

Agentは合計64 GPUを持つ4ノードAMD Instinct MI355Xクラスター上で、FP8精度のMixtral 8×22Bモデル・グローバルバッチサイズ512・シーケンス長8,192トークンで検証された。対象メトリクスはGPUあたりのスループット——カード1枚あたり1秒間に生成されるトークン数だ。

設定空間は5次元の並列化（テンソル・パイプライン・エキスパート・コンテキスト・データ）・1から8のマイクロバッチサイズ・再計算の粒度（none/selective/full）・MoEバックエンド（All-to-AllとDeepEP）・さまざまなパイプラインスケジューラー（1F1B・Interleaved・Zero-Bubbleバリアント）を含む。トポロジー的に有効な点の数は約10の4乗——ブルートフォース探索は現実的ではない。

Agentはどのように組み合わせ空間を探索するか？

Primus Tuning Agentは確定性とLLM誘導の推論を組み合わせた二段階アプローチを使用する。

最初のフェーズでは、確定的シードプランナーがパラメータ軸の各々を順番に通過する——他のすべての値を固定して1つのパラメータの変化による効果を測定する。この単軸スイープは基準点を確立し、高コストのクロス軸実験なしに明らかに弱いオプションを排除する。

第2フェーズでは、LLM誘導探索ループがLiteLLMルーティングを持つDSPyプランナーを使用して、単軸スイープでは見えなかった組み合わせを特定する。LLMは軸間のトレードオフ——例えばMoEバックエンドの選択がパイプラインスケジューラーの最適性にどう影響するか——を検討し、測定のための候補設定を提案する。AgentはOpenAI・Anthropic・その他のLLMプロバイダーを別個のプロキシインフラなしでサポートする。

分析的なメモリ事前フィルターがGPUカーネルを実行せずに各候補設定を評価する。ケーススタディでは、フィルターは30の試験設定のうち7つをゼロのGPUコストで除外し、分析的に利用可能なメモリを超えるものを排除した。これは探索のスケーラビリティにとって重要だ。

Mixtral 8×22BとMI355Xクラスターでの結果

AgentはGPU 1枚あたり毎秒4,402トークンを達成する設定を発見した。これは64 GPU MI355Xクラスター上のAMDの公開BF16リファレンス3,475トークン/秒/GPUを**+27%上回る実測向上だ。推定最適結果は4,908トークン/秒/GPU**（+41.2%）で、測定値に対する文書化された約10%の楽観バイアスがある——Agentはこれを異常ではなく既知の特性として記録している。

発見された改善の鍵はクロス軸シナジーにある。DeepEP MoEバックエンド・特定のパイプラインスケジューラー・再計算粒度の組み合わせは、単軸スキャンにも軸が独立して最適化されるという素朴な仮定にも見えなかった。LLM誘導ループは、確定的プランナーが構造的に検討できないトレードオフを推論することでそのシナジーを特定した。

探索全体は主にシングルノードプロファイリングデータ——その後フルクラスターシナリオに外挿——によって約30分かかった。

統合と再現性

Agentは探索の各フェーズを文書化する。どの軸を探索したか・どのメモリフィルターを適用したか・推定結果と測定値がどのように比較されたか。各実行は再現可能で監査可能であり、これは組織内でインフラの決定を正当化しなければならないチームや、探索ループで異なるLLMを使って実験を繰り返すチームにとって重要だ。

LiteLLMレイヤーを持つDSPyフレームワークにより、ユーザーは探索ループのために独自のLLMエンドポイントを使用でき、探索コストを制御可能に保ち特定のクラウドプロバイダーから独立させられる。OpenAI・Anthropic・その他のプロバイダーへのサポートは別個のプロキシインフラなしで提供され、Agentの設定における運用オーバーヘッドを削減する。

AMD Instinctワークロードをスケールするチームにとって、Primus Tuning Agentは初期設定から測定可能な改善への再現可能なパスを提供する——高価なクラスター時間での何週間にもわたる手動実験なしで。**+41.2%の推定向上と実測の+27%**はツール自体が文書化した投影差と一致しており、ベンチマーク表よりも容量計画の結果を信頼できるものにする。