torch_npu

torch_npu の機能と使用ガイド

このスキルを使用するタイミング

• ユーザーが @torch_npu を指定した場合、または昇腾 NPU、CANN、Ascend、PyTorch の NPU 上での実行について言及した場合。
• ユーザーが 環境チェック／デプロイ を必要としている場合：PyTorch のチェックまたはデプロイ、環境が NPU をサポートしているかの確認。
• ユーザーが @torch_npu_doc を使用する場合：このスキルの reference.md に基づいて、プロジェクト内の中国語ドキュメントによる機能説明と操作手順を提供します。
• ユーザーが NPU 形式変換のコードヒント を必要としている場合：torch_npu.npu_format_cast、npu_format_cast_、torch_npu.Format、get_npu_format の記述または補完時に、§2.1 のサブ機能に従ってパラメーターと列挙値のヒントを提供します。この機能が分からない場合は、このスキルの MCP を通じてバージョン別ドキュメントを取得し、補足説明することを提案できます。

---

1. 環境チェックとデプロイ

1.1 PyTorch と NPU 環境の自動チェック

回答またはスクリプト生成時に、必要に応じて以下のチェックを実行するか、ユーザーに実行を提案します。

PyTorch と Python のバージョンが対応範囲内か確認します（README.zh.md の「PyTorch与Python版本配套表」を参照）：

• PyTorch 1.11.0～2.9.0 などの複数のバージョンをサポートし、対応する Python は 3.7～3.11 です（具体的な PyTorch のバージョンによって異なります）。

環境が NPU をサポートしているか確認します：

import torch
import torch_npu  # 2.5.1 及以后可不显式 import，仍建议写以便兼容

# 是否可用 NPU、设备数量
if torch.npu.is_available():
    count = torch.npu.device_count()
    # 使用 device='npu' 或 .npu()
else:
    # 未安装 CANN / 未 source set_env.sh / 无 NPU 设备
    pass

CANN 環境変数を確認します（インストール後の検証）：

• 使用前に source /usr/local/Ascend/ascend-toolkit/set_env.sh を実行する必要があります（パスは実際の CANN インストール先に合わせてください）。
• ASCEND_HOME_PATH、ASCEND_OPP_PATH が設定されていない、またはパスが存在しない場合、torch_npu はエラーを報告し、source set_env.sh の実行を促します。

1.2 デプロイ手順の概要

1. CANN をインストールする：CANN インストールガイドに従ってインストールし、README.zh.md の「昇腾辅助软件」表で CANN と PyTorch／Extension のバージョンを照合します。
2. PyTorch をインストールする：x86 では pip3 install torch==2.7.1+cpu --index-url https://download.pytorch.org/whl/cpu、aarch64 では pip3 install torch==2.7.1 を使用します（バージョンは例です。必要に応じて置き換えてください）。
3. 依存関係をインストールする：pip3 install pyyaml setuptools。
4. torch_npu をインストールする：pip3 install torch-npu==2.7.1（PyTorch および CANN と対応するバージョンが必要です）。
5. 検証する：source set_env.sh の実行後、簡易検証コードを実行します（README.zh.md の「インストール後の検証」を参照）。

---

2. torch_npu の機能一覧（概要）

カテゴリ	機能の説明
デバイスとメモリ	torch.npu：デバイス管理、device_count、current_device、set_device、synchronize、Stream／Event、メモリの統計と割り当て（memory_allocated、empty_cache、MemPool など）。
テンソル／ストレージ	tensor.npu()、tensor.is_npu、NPU Storage、シリアライズ torch.save／load の NPU サポート、DDP／マルチプロセスの reductions。
トレーニング／最適化	torch.npu.amp 混合精度、torch_npu.optim、FSDP パッチ（ShardedGradScaler）、勾配チェックポイントのデフォルト NPU 化。
分散処理	torch_npu.distributed：HCCL／LCCL バックエンド、is_hccl_available、reinit_process_group、RPC、対称メモリ、DTensor ルール。
拡張 API	torch_npu.contrib：NMS、IoU 系列、ROIAlign、DCN、FusedAttention、カスタムモジュール（DropoutWithByteMask など）。
グラフとコンパイル	NPU Graph（npugraphify）、Dynamo、Inductor、torch.compile の NPU サポート。
推論／ONNX	ONNX エクスポートと NPU カスタム演算子のラッパー（OneHot、RoiAlign、NMS、FastGelu、MultiHeadAttention など）。
プロファイリング	torch_npu.profiler、MSTX パッチ、パフォーマンスダンプ。
その他	HiFloat8Tensor、erase_stream、matmul_checksum、transfer_to_npu（オプション）、op_plugin。

詳細な API については、昇腾 Ascend Extension for PyTorch カスタム API リファレンス およびプロジェクトの README.zh.md を参照してください。

---

2.1 サブ機能：torch_npu.npu_format_cast のコードヒント

ユーザーがコードで torch_npu.npu_format_cast、npu_format_cast_、NPU テンソル形式の変換を記述または質問した場合、IDE で各種ヒントを利用できるよう、以下のコードヒントと補完ガイドを提供します。

API シグネチャとパラメーター

• torch_npu.npu_format_cast(tensor, acl_format, customize_dtype=None)

  • tensor：NPU 上の torch.Tensor（先に .npu() を呼び出す必要があります）。
  • acl_format：変換先のストレージ形式。int または torch_npu.Format の列挙メンバーを指定できます。
  • customize_dtype：オプション。ONNX などの用途でカスタム dtype を指定します。
  • 戻り値：新しいテンソル（元のテンソルは変更されません）。

• torch_npu.npu_format_cast_(tensor, acl_format)

  • 上記と同じですが、in-place 版であり、tensor の形式を直接変更します。

• torch_npu.get_npu_format(tensor)

  • テンソルの現在の NPU ストレージ形式を返します（torch_npu.Format または整数）。

よく使う Format 列挙値（torch_npu.Format）

コードヒントでは、以下のよく使われる値を優先的に提示できます（torch_npu.npu._format.Format から引用）：

列挙名	値	よくある用途
Format.NCHW	0	デフォルトの 4D 畳み込みレイアウト
Format.NHWC	1	チャネルが後ろにある 4D レイアウト
Format.ND	2	汎用 ND レイアウト
Format.NC1HWC0	3	Conv／BatchNorm などの演算子でよく使用
Format.FRACTAL_Z	4	3D 畳み込みなど
Format.FRACTAL_NZ	29	線形演算／行列乗算、Attention の重みなど
Format.NDC1HWC0	32	5D
Format.FRACTAL_Z_3D	33	3D 畳み込み
Format.UNDEFINED	-1	未定義

その他の選択肢：NC1HWC0_C04(12)、HWCN(16)、NDHWC(27)、NCDHW(30)、NC(35)、NCL(47)、FRACTAL_NZ_C0_*(50–54) など。

コードヒントと補完のルール

1. 第 2 引数の補完：ユーザーが torch_npu.npu_format_cast(x, と入力したとき、acl_format の候補として int または torch_npu.Format.xxx を提示し、よく使われる列挙値（Format.NCHW、Format.NHWC、Format.FRACTAL_NZ、Format.NC1HWC0 など）を列挙します。
2. Format 列挙値の補完：ユーザーが torch_npu.Format. と入力したとき、上記の列挙値のメンバー一覧を提示します。
3. 組み合わせて使用：コード内に get_npu_format(t) がある場合、同じ形式に変換する必要があるときは、torch_npu.npu_format_cast(other, torch_npu.get_npu_format(t)) を提示できます。
4. よくあるケース：
   • 線形層の重みの量子化／NPU への移行：torch_npu.npu_format_cast(weight.npu(), 29)（FRACTAL_NZ）；
   • パラメーターと同じ形式の勾配：torch_npu.npu_format_cast(p.grad, torch_npu.get_npu_format(p))；
   • モジュール移行時の BN／Conv の NC1HWC0：torch_npu.npu_format_cast(tensor, 3) または Format.NC1HWC0。

ドキュメントの参照元について

より信頼性の高い形式の説明や CANN との対応関係が必要な場合は、本 skill の MCP（mcp/README.md を参照）を使って fetch_torch_npu_doc を実行し、GitCode Ascend/pytorch の該当ブランチにあるバージョン別ドキュメント（docs/zh 内の format または演算子の対応に関するドキュメントなど）を取得するか、reference から「フレームワーク機能ガイド」に関する項目を探してください。

---

3. ドキュメントの機能を自動的に提供する方法（@torch_npu_doc）

• ユーザーが @torch_npu_doc を付けるか、torch_npu の中国語ドキュメントを明確に参照するよう求めた場合：
  • まず本 skill の reference ドキュメント：reference.md を使用してください。プロジェクト内のすべての中国語ドキュメントの索引と重要な内容の概要が含まれています。
  • ユーザーの質問に応じて、reference.md の「ドキュメント索引」から該当するドキュメントのパス（torch_npu/README.zh.md、torch_npu/docs/zh/quick_start/quick_start.md、torch_npu/docs/zh/framework_feature_guide_pytorch/、torch_npu/docs/zh/troubleshooting/ など）を特定し、必要に応じて作業領域内のファイルを読み、関連する段落を引用してください。
  • インストール、検証、アンインストール、バージョンの互換性については README.zh.md または reference の概要を使用してください。移行とクイックスタートについては quick_start を使用してください。演算子の対応、メモリ、パフォーマンス、グラフモードについては framework_feature_guide_pytorch 内の該当ドキュメントを使用してください。エラーやトラブルについては troubleshooting 内の該当ドキュメントを使用してください。
• 回答は中国語で簡潔に説明し、ユーザーが自分で参照できるよう、引用したドキュメントのパスを明記してください。

---

4. 参考リソース

• 本 skill のドキュメント参照：reference.md。torch_npu プロジェクト内のすべての中国語ドキュメントの索引と要点をまとめています。ユーザーが @torch_npu_doc を付けた場合、これをもとに該当ドキュメントを特定して引用します。
• プロジェクトの中国語説明：作業領域内の torch_npu/README.zh.md（バージョン表、デプロイ、検証、参考ドキュメントへのリンク）。
• 公式ドキュメント：昇腾社区 Ascend Extension for PyTorch、ソフトウェアインストールガイド、モデル移行とトレーニング、演算子の対応、PyTorch ネイティブインターフェース一覧、カスタム API リファレンス（README.zh.md 下部の表を参照）。
• GitCode からのバージョン別ドキュメントの取得：本 skill は MCP サービスを提供しており、GitCode Ascend/pytorch からブランチに応じたドキュメントを取得し、fetched_docs/ に保存できます。インストールと設定については mcp/README.md を参照してください。プロジェクトには .cursor/mcp.json が設定済みです。依存関係をインストールして Cursor を再起動すれば、fetch_torch_npu_doc、fetch_torch_npu_docs_batch、list_torch_npu_doc_paths ツールを使用できます。