kubeflow/trainer

何を解決するか

Kubeflow Trainer は、大規模言語モデル（LLM）向けの分散AIトレーニングおよびファインチューニングの複雑さに対処するように設計されています。Kubernetesクラスタ上でマルチノード・マルチGPUワークロードのオーケストレーションを簡素化し、高性能な通信と効率的なリソース利用を保証します。

動作方法

Kubernetesネイティブなプラットフォームとして、TrainJob および Runtimes などのAPIを使用して分散ジョブをオーケストレーションします。既存のKubernetesのビルディングブロックである JobSet および LeaderWorkerSet を活用してオーケストレーションを実現します。パフォーマンス最適化のため、GPUノード間の高速同期を実現するMPI（Message Passing Interface）をKubernetesに導入し、大規模データをGPUにゼロコピーで直接転送する分散データキャッシュを提供します。

対象ユーザー

PyTorch、JAX、DeepSpeed、Megatron-LMなどのさまざまなフレームワークを使用して、ハイパフォーマンスコンピューティング（HPC）クラスタ上でLLMやその他のAIモデルのトレーニングやファインチューニングをスケーリングする必要があるAI実践者やMLエンジニア。

主な特徴

• マルチフレームワーク対応: PyTorch、MLX、HuggingFace、DeepSpeed、Megatron-LM、JAX、XGBoostなど、幅広いフレームワークをサポート。
• トポロジー対応スケジューリング: Kueue、Slurm Bridge、KAI Schedulerと統合し、高度なGPUおよびマルチクラスタスケジューリングを実現。
• ハイパフォーマンス通信: MPIを使用して、ノード間のプロセス間で高スループット通信を可能に。
• 分散データキャッシュ: GPUの利用効率を最大化し、メモリオーバーヘッドを削減するゼロコピーデータストリーミングメカニズムを実装。