NVIDIA/JAX-Toolbox

解決する課題

JAX Toolbox は、開発者が NVIDIA GPU 上で JAX ベースの機械学習ワークロードを実行、最適化、およびスケーリングするための集中リソースを提供します。事前構築済みのコンテナイメージ、最適化されたライブラリ、トレーニングと推論のための技術ガイダンスを提供することで、高性能環境の構成の複雑さに対処します。

仕組み

このプロジェクトは、NVIDIA の JAX-on-GPU エコシステムのハブとして機能します。最適化されたバージョンの JAX、XLA、および TransformerEngine などの関連ライブラリを含む、ナイトリービルドと安定版のコンテナイメージ（NGC および GitHub Container Registry 経由）を提供します。これらのコンテナには、NVIDIA ハードウェアで高性能を保証するために必要な CUDA および NCCL ドライバが事前構成されています。リポジトリには、大規模モデルのスケーリングのためのドキュメント、技術ブログ、パフォーマンスチューニングガイドもホストされています。

対象読者

このプロジェクトは、JAX を使用して大規模モデルを構築またはトレーニングし、NVIDIA GPU クラスタでピークパフォーマンスを確保する必要がある機械学習エンジニア、研究者、インフラストラクチャ開発者を対象としています。

ハイライト

• JAX および MaxText 向けの本番環境対応コンテナイメージを提供。ナイトリービルドと安定した月次リリースを含みます。
• JAX ワークロードのプロファイリング、最適化、およびスケーリングに関する詳細な技術ドキュメントを提供。
• 大規模言語モデル（LLM）および長文脈モデルのトレーニングに焦点を当てたパフォーマンス重視のガイドを含みます。
• FP4/FP8 精度、CuTe を使用したカスタムカーネル開発、分散トレーニング戦略などの高度な GPU 機能のリソースを特集。