Google、TPU開発を加速するAIエージェントツール集「Accelerator Agents」公開

Accelerator Agentsは、Google Cloud TPU 上での機械学習開発を加速するために設計された AI 活用ツールの集合体です。このリポジトリには、コードの移行やカーネル最適化、パフォーマンスチューニングを支援するエージェントが格納されており、開発者が TPUs の性能を最大限に引き出しつつ、より高速な開発を実現できるようになっています。

免責事項: これは公式にサポートされている Google 製品ではありません。このプロジェクトは Google Open Source Software Vulnerability Rewards Program の対象外です。

概要

機械学習モデルの複雑さが増すにつれ、TPU などの特定のハードウェアアクセラレータ向けに最適化することはますます困難になっています。このプロジェクトは、Gemini を搭載した専門的な AI ツールである「エージェント」を提供し、これらの複雑な作業を自動化・支援することを目指しています。

本プロジェクトには、主に以下の 2 つのエージェントが含まれています。

1. MaxCode

MaxCode エージェントは、既存の PyTorch モデルやコードベースを JAX へ変換するのを支援します。高性能なフレームワークである MaxText を活用し、TPU で効率的に実行できるようワークロードを移行したいユーザーをサポートするために設計されています。

*注記：MaxCode は現在も開発中で、移行品質の向上と対応モデル範囲の拡大に取り組んでいます。

主な機能:

• 自動変換: PyTorch の機能的なコードブロックやモデル層を JAX へ変換します。

MaxText 連携機能では、TPU 上で即座にトレーニングや推論を実行できるよう、MaxText フレームワークと互換性のある JAX コードを生成します。

人間による監視（Human-in-the-Loop）に対応しており、開発者がレビューして改良を加えられる初期実装案のドラフト作成をサポートします。

2. MaxKernel

MaxKernel エージェントは、TPU 向けの高性能カーネル開発に特化した専用ツールです。エンジニアがカスタムカーネルの記述、最適化、デバッグを支援し、特に JAX のカーネル言語である Pallas に焦点を当てています。

主な機能:

• カーネル作成: Pallas カーネルをゼロから作成するか、JAX のリファレンス実装に基づいてドラフトを作成します。
• CUDA から Pallas への移植: カスタム CUDA/GPU カーネルを TPU で最適に実行できるよう移行するのを支援します。
• 最適化とプロファイリング: カーネルパフォーマンス（MFU）向上のためのプロファイリング洞察と最適化の提案を提供します。
• テストハッチ生成: 正しさの確認やコンパイルチェック用のボイラープレートコードを自動的に生成します。

はじめに

事前準備

• Google Cloud VM。MaxCode では CPU のみの VM で十分ですが、MaxKernel では TPU VM の利用をお勧めします。
• Python 3.11 以上
• Gemini API へのアクセス（エージェントの推論機能のため）

インストール

リポジトリをクローンしてください：bash git clone https://github.com/AI-Hypercomputer/accelerator-agents.git cd accelerator-agents

*(注：各エージェント固有のインストール手順は、それぞれのサブディレクトリに記載されています。)*

コントリビューション

本プロジェクトへの貢献を歓迎します。プルリクエストの提出方法、課題の報告方法、およびプロジェクトへの参加方法の詳細については、CONTRIBUTING.md をご覧ください。

ライセンス

本プロジェクトは Apache License, Version 2.0 の下でライセンスされています。完全なライセンステキストについては LICENSE をご確認ください。