Ascend/pytorch

何を解決するか

TorchNPUは、PyTorchのディープラーニングフレームワークを、華為（Huawei）のAscend AIプロセッサ（NPU）上で直接実行できるようにします。PyTorchの高レベルAPIとAscend NPUの特定のハードウェアアーキテクチャの間のギャップを埋め、開発者がPyTorchエコシステムを離さずにAscendの計算パワーを活用できるようにします。

動作方法

PrivateUse1と呼ばれるメカニズムを通じてPyTorchに統合されるアダプタプラグインとして機能します。以下の主要な機能モジュールを提供します：

• 基本計算：主流のAIワークロードをサポートするため、PyTorchのネイティブAPIおよびカスタムAPIを実装します。
• 分散学習：FSDP2を使用して大規模な学習を加速し、AllReduceやAllGatherなどの集約通信プリミティブおよびDTensorをサポートします。
• グラフモード：動的グラフキャプチャ、静的グラフ最適化、効率的なコード生成の組み合わせにより、学習および推論を高速化します。
• メモリ管理：仮想メモリ管理を採用してフラグメンテーションを低減し、分散環境におけるストリーム間のメモリ再利用を最適化します。
• デバッグとチューニング：計算、メモリ、通信のプロファイリングツールと、通信異常を監視する「WatchDog」を備えています。

対象ユーザー

PyTorchを使用し、モデルを華為のAscend NPUハードウェアに移行するか、その上で学習を実行してパフォーマンスを向上させたいAI開発者および研究者です。

主な特徴

• シームレスな統合：成熟したPyTorchの機能を引き継ぎつつ、NPU固有の最適化を追加します。
• 分散スケーリング：FSDP2およびDTensor向けの専用サポートにより、高性能な並列学習を実現します。
• パフォーマンス最適化：NPUGraphを活用したグラフモードにより、実行を高速化します。
• ハードウェア固有のチューニング：詳細なプロファイリングとリアルタイム通信監視を提供します。