pytorch-fsdp

Pytorch-Fsdp スキル

公式ドキュメントから生成された、pytorch-fsdp 開発に関する包括的なサポートです。

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

このスキルは次の場合に起動されるべきです:

• pytorch-fsdp を使用しているとき
• pytorch-fsdp の機能や API について質問するとき
• pytorch-fsdp のソリューションを実装するとき
• pytorch-fsdp のコードをデバッグするとき
• pytorch-fsdp のベストプラクティスを学ぶとき

クイックリファレンス

一般的なパターン

パターン 1: 汎用 Join コンテキストマネージャー# 作成日: Jun 06, 2025 | 最終更新日: Jun 06, 2025 汎用 Join コンテキストマネージャーは、入力数が不均一な場合の分散学習を支援します。このページでは、関連するクラスである Join、Joinable、JoinHook の API について説明します。チュートリアルについては、「Join コンテキストマネージャーを使った不均一な入力による分散学習」を参照してください。 class torch.distributed.algorithms.Join(joinables, enable=True, throw_on_early_termination=False, **kwargs)[source]# このクラスは汎用の join コンテキストマネージャーを定義します。これにより、プロセスが join した後にカスタムフックを呼び出すことができます。これらのフックは、join していないプロセスの集合通信をシャドウイングし、ハングやエラーを防ぎ、アルゴリズムの正しさを保証する必要があります。フックの定義の詳細は JoinHook を参照してください。 警告: このコンテキストマネージャーでは、正しさを保つために、参加する各 Joinable が、そのイテレーションごとの集合通信の前に notify_join_context() メソッドを呼び出す必要があります。 警告: このコンテキストマネージャーでは、すべての JoinHook オブジェクトの process_group 属性が同一である必要があります。JoinHook オブジェクトが複数ある場合は、最初のオブジェクトのデバイスが使用されます。プロセスグループとデバイスの情報は、join していないプロセスの確認と、throw_on_early_termination が有効な場合に例外を投げるようプロセスへ通知するために使用されます。どちらも all-reduce を使用します。 パラメーター: joinables (List[Joinable]) – 参加する Joinable のリスト。それらのフックは与えられた順序で反復処理されます。 enable (bool) – 不均一な入力の検出を有効にするフラグ。False に設定するとコンテキストマネージャーの機能が無効になり、入力が不均一にならないことがわかっている場合にのみ設定してください (デフォルト: True)。 throw_on_early_termination (bool) – 不均一な入力を検出したときに例外を投げるかどうかを制御するフラグ (デフォルト: False)。 例: >>> import os >>> import torch >>> import torch.distributed as dist >>> import torch.multiprocessing as mp >>> import torch.nn.parallel.DistributedDataParallel as DDP >>> import torch.distributed.optim.ZeroRedundancyOptimizer as ZeRO >>> from torch.distributed.algorithms.join import Join >>> >>> # 生成された各ワーカー上で >>> def worker(rank): >>> dist.init_process_group("nccl", rank=rank, world_size=2) >>> model = DDP(torch.nn.Linear(1, 1).to(rank), device_ids=[rank]) >>> optim = ZeRO(model.parameters(), torch.optim.Adam, lr=0.01) >>> # rank 1 は rank 0 より入力が 1 つ多い >>> inputs = [torch.tensor([1.]).to(rank) for _ in range(10 + rank)] >>> with Join([model, optim]): >>> for input in inputs: >>> loss = model(input).sum() >>> loss.backward() >>> optim.step() >>> # すべての rank がハングやエラーなしにここに到達する static notify_join_context(joinable)[source]# 呼び出し元のプロセスがまだ join していないことを join コンテキストマネージャーに通知します。その後、throw_on_early_termination=True の場合は、不均一な入力が検出されたかどうか (つまり、あるプロセスがすでに join しているかどうか) を確認し、検出された場合は例外を投げます。このメソッドは、Joinable オブジェクトがイテレーションごとの集合通信を行う前に呼び出す必要があります。例えば、DistributedDataParallel の forward パスの開始時に呼び出す必要があります。コンテキストマネージャーに渡された最初の Joinable オブジェクトのみがこのメソッドで集合通信を実行し、それ以外の場合このメソッドは何もしません。 パラメーター: joinable (Joinable) – このメソッドを呼び出す Joinable オブジェクト。 戻り値: joinable がコンテキストマネージャーに渡された最初のオブジェクトである場合は、プロセスがまだ join していないことをコンテキストマネージャーに通知するための all-reduce の非同期 work ハンドル。それ以外の場合は None。 class torch.distributed.algorithms.Joinable[source]# これは join 可能なクラスの抽象基底クラスを定義します。join 可能なクラス (Joinable を継承) は、JoinHook インスタンスを返す join_hook() に加えて、それぞれデバイスとプロセスグループの情報を返す join_device() と join_process_group() を実装する必要があります。 abstract property join_device: device# join コンテキストマネージャーに必要な集合通信を実行するデバイスを返します。 abstract join_hook(**kwargs)[source]# 与えられた Joinable に対応する JoinHook インスタンスを返します。 パラメーター: kwargs (dict) – 実行時に join フックの動作を変更するための任意のキーワード引数を含む dict。同じ join コンテキストマネージャーを共有するすべての Joinable インスタンスには、同じ kwargs の値が渡されます。 戻り値の型: JoinHook abstract property join_process_group: Any# join コンテキストマネージャー自体に必要な集合通信のためのプロセスグループを返します。 class torch.distributed.algorithms.JoinHook[source]# これは join フックを定義します。join フックは join コンテキストマネージャー内の 2 つのエントリポイントを提供します。エントリポイント: 参加していないプロセスが存在する間に繰り返し呼び出されるメインフックと、すべてのプロセスが join した後に一度だけ呼び出されるポストフックです。汎用 join コンテキストマネージャー用の join フックを実装するには、JoinHook を継承するクラスを定義し、必要に応じて main_hook() と post_hook() をオーバーライドします。 main_hook()[source]# 参加していないプロセスが存在する間にこのフックを呼び出し、学習イテレーション内の集合通信をシャドウイングします。学習イテレーションとは、1 回の forward パス、backward パス、オプティマイザーのステップを指します。 post_hook(is_last_joiner)[source]# すべてのプロセスが join した後にフックを呼び出します。追加の bool 引数 is_last_joiner を受け取ります。これは、そのランクが最後に join したプロセスの一つであるかどうかを示します。 パラメーター: is_last_joiner (bool) – ランクが最後に join したプロセスの一つである場合は True、それ以外の場合は False。

Join

パターン 2： 分散通信パッケージ - torch.distributed# 作成日：Jul 12, 2017 | 最終更新日：Sep 04, 2025

注記 分散トレーニングに関連するすべての機能の簡単な概要については、PyTorch Distributed Overview を参照してください。

バックエンド#

torch.distributed は、機能の異なる 4 つの組み込みバックエンドをサポートしています。次の表は、各バックエンドで CPU または GPU に使用できる関数を示しています。NCCL の GPU は CUDA GPU、XCCL の GPU は XPU GPU を指します。MPI が CUDA をサポートするのは、PyTorch のビルドに使用した実装が CUDA をサポートしている場合のみです。

バックエンド	gloo	mpi	nccl	xccl
デバイス	CPU	GPU	CPU	GPU
send	✓	✘	✓	?
recv	✓	✘	✓	?
broadcast	✓	✓	✓	?
all_reduce	✓	✓	✓	?
reduce	✓	✓	✓	?
all_gather	✓	✓	✓	?
gather	✓	✓	✓	?
scatter	✓	✓	✓	?
reduce_scatter	✓	✓	✘	✘
all_to_all	✓	✓	✓	?
barrier	✓	✘	✓	?

PyTorch に付属するバックエンド#

PyTorch の分散パッケージは、Linux（安定版）、MacOS（安定版）、Windows（プロトタイプ）をサポートしています。Linux ではデフォルトで Gloo および NCCL バックエンドがビルドされ、PyTorch の分散パッケージに含まれます（NCCL は CUDA を使用してビルドした場合のみ）。MPI はオプションのバックエンドであり、PyTorch をソースからビルドする場合にのみ含められます（たとえば、MPI がインストールされたホスト上で PyTorch をビルドする場合）。

注記 PyTorch v1.8 以降、Windows は NCCL を除くすべての集合通信バックエンドをサポートしています。init_process_group() の init_method 引数がファイルを指す場合、次のスキーマに従う必要があります。

ローカルファイルシステムの場合、init_method="file:///d:/tmp/some_file"

共有ファイルシステムの場合、init_method="file://////{machine_name}/{share_folder_name}/some_file"

Linux プラットフォームと同様に、環境変数 MASTER_ADDR と MASTER_PORT を設定すると TcpStore を有効にできます。

どのバックエンドを使うべきか？#

以前から「どのバックエンドを使うべきか」とよく質問されてきました。

経験則

CUDA GPU を使った分散トレーニングには、NCCL バックエンドを使用してください。

XPU GPU を使った分散トレーニングには、XCCL バックエンドを使用してください。

CPU を使った分散トレーニングには、Gloo バックエンドを使用してください。

InfiniBand インターコネクトを使う GPU ホスト

InfiniBand と GPUDirect を現在サポートしている唯一のバックエンドであるため、NCCL を使用してください。

Ethernet インターコネクトを使う GPU ホスト

NCCL は現在、特に単一ノードまたは複数ノードのマルチプロセス分散トレーニングで、最も優れた分散 GPU トレーニング性能を提供するため、NCCL を使用してください。NCCL で問題が発生した場合は、代替手段として Gloo を使用してください（現在、GPU では Gloo の方が NCCL より低速です）。

InfiniBand インターコネクトを使う CPU ホスト

InfiniBand で IP over IB が有効な場合は Gloo を使用し、そうでない場合は MPI を使用してください。今後のリリースで、Gloo に InfiniBand サポートを追加する予定です。

Ethernet インターコネクトを使う CPU ホスト

MPI を使う特別な理由がない限り、Gloo を使用してください。

共通の環境変数#

使用するネットワークインターフェイスの選択#

デフォルトでは、NCCL と Gloo の両バックエンドは適切なネットワークインターフェイスを自動的に検出しようとします。自動検出されたインターフェイスが正しくない場合は、次の環境変数を使って上書きできます（それぞれ対応するバックエンドに適用されます）。

NCCL_SOCKET_IFNAME（例：export NCCL_SOCKET_IFNAME=eth0）

GLOO_SOCKET_IFNAME（例：export GLOO_SOCKET_IFNAME=eth0）

Gloo バックエンドを使う場合は、複数のインターフェイスをカンマで区切って指定できます（例：export GLOO_SOCKET_IFNAME=eth0,eth1,eth2,eth3）。バックエンドは、これらのインターフェイスに対してラウンドロビン方式で処理を振り分けます。この変数で指定するインターフェイスの数は、すべてのプロセスで同じでなければなりません。

その他の NCCL 環境変数#

デバッグ — NCCL の障害時に、明示的な警告メッセージと基本的な NCCL 初期化情報を出力するには、NCCL_DEBUG=INFO を設定できます。NCCL_DEBUG_SUBSYS を使うと、特定の項目についてより詳細な情報を取得できます。たとえば NCCL_DEBUG_SUBSYS=COLL を設定すると、集合通信呼び出しのログが出力されます。これは、特に集合通信の種類やメッセージサイズの不一致が原因のハングをデバッグする際に役立ちます。トポロジー検出に失敗した場合は、NCCL_DEBUG_SUBSYS=GRAPH を設定して検出結果の詳細を確認し、NCCL チームに追加のサポートを求める際の参照情報として保存すると便利です。

パフォーマンスチューニング — NCCL はトポロジー検出に基づいて自動的にチューニングを行い、ユーザーの調整負担を軽減します。ソケットベースのシステムによっては、NCCL_SOCKET_NTHREADS と NCCL_NSOCKS_PERTHREAD を調整して、ソケットネットワーク帯域幅を増やせる場合があります。AWS や GCP など、一部のクラウドプロバイダー向けには、これら 2 つの環境変数が NCCL によって事前調整されています。NCCL の環境変数の一覧については、NVIDIA NCCL の公式ドキュメントを参照してください。

torch.distributed.ProcessGroupNCCL.NCCLConfig と torch.distributed.ProcessGroupNCCL.Options を使って、NCCL コミュニケーターをさらに調整できます。インタープリターで help を使うと、詳細を確認できます（例：help(torch.distributed.ProcessGroupNCCL.NCCLConfig)）。

基本#

torch.distributed パッケージは、1 台以上のマシンで動作する複数の計算ノード間のマルチプロセス並列処理に、PyTorch のサポートと通信プリミティブを提供します。

torch.nn.parallel.DistributedDataParallel() クラスは、この機能を基盤として、任意の PyTorch モデルをラップする同期分散トレーニングを提供します。これは、Multiprocessing package - torch.multiprocessing や torch.nn.DataParallel() が提供する並列処理とは異なり、ネットワーク接続された複数のマシンをサポートします。また、ユーザーはトレーニングスクリプトのコピーを各プロセスで個別に起動する必要があります。

単一マシンで同期処理を行う場合でも、torch.distributed や torch.nn.parallel.DistributedDataParallel() ラッパーには、torch.nn.DataParallel() を含む他のデータ並列化手法より優れた点があります。

各プロセスは独自のオプティマイザーを保持し、各イテレーションで最適化ステップ全体を実行します。勾配はすでにプロセス間で集約・平均されているため各プロセスで同一となり、この処理は冗長に見えるかもしれません。しかし、パラメーターのブロードキャストが不要になり、ノード間のテンソル転送にかかる時間を削減できます。

各プロセスに独立した Python インタープリターが含まれるため、単一の Python プロセスから複数の実行スレッド、モデルのレプリカ、または GPU を操作した場合に発生する、余分なインタープリターのオーバーヘッドや「GIL の競合」を避けられます。これは、再帰レイヤーや多数の小さなコンポーネントを含むモデルなど、Python ランタイムを多用するモデルで特に重要です。

初期化#

他のメソッドを呼び出す前に、torch.distributed.init_process_group() または torch.distributed.device_mesh.init_device_mesh() 関数を使ってパッケージを初期化する必要があります。どちらの関数も、すべてのプロセスが参加するまでブロックします。

警告 初期化はスレッドセーフではありません。ランク間で一貫性のない「UUID」が割り当てられるのを防ぎ、ハングにつながる初期化中の競合を防ぐため、プロセスグループの作成は単一スレッドから実行してください。

torch.distributed.is_available()[source]# 分散パッケージが利用可能な場合は True を返します。利用できない場合、torch.distributed は他の API を公開しません。現在、torch.distributed は Linux、MacOS、Windows で利用できます。ソースから PyTorch をビルドする際に有効にするには、USE_DISTRIBUTED=1 を設定します。現在、Linux と Windows では USE_DISTRIBUTED=1、MacOS では USE_DISTRIBUTED=0 がデフォルト値です。

戻り値の型 bool

torch.distributed.init_process_group(backend=None, init_method=None, timeout=None, world_size=-1, rank=-1, store=None, group_name='', pg_options=None, device_id=None)[source]# デフォルトの分散プロセスグループを初期化します。これにより、分散パッケージも初期化されます。

プロセスグループを初期化する主な方法は 2 つあります。

store、rank、world_size を明示的に指定します。

ピアの検出場所や方法を示す init_method（URL 文字列）を指定します。任意で rank と world_size を指定するか、必要なパラメーターをすべて URL に含めて省略できます。どちらも指定されない場合、init_method は “env://” とみなされます。

パラメーター

backend (str or Backend, optional) – 使用するバックエンドです。ビルド時の設定によって、有効な値は mpi、gloo、nccl、ucc、xccl、またはサードパーティープラグインで登録されたバックエンドです。2.6 以降、backend が指定されない場合、c10d は device_id キーワード引数（指定されている場合）が示すデバイスタイプに登録されたバックエンドを使用します。現在知られているデフォルトの登録は、cuda に対する nccl、cpu に対する gloo、xpu に対する xccl です。backend と device_id のどちらも指定されない場合、c10d は実行時のマシン上のアクセラレーターを検出し、検出されたアクセラレーター（または cpu）に登録されたバックエンドを使用します。このフィールドには小文字の文字列（例："gloo"）を指定できます。Backend 属性（例：Backend.GLOO）からも参照できます。nccl バックエンドでマシンあたり複数のプロセスを使う場合、各プロセスは使用するすべての GPU に排他的にアクセスできなければなりません。プロセス間で GPU を共有すると、デッドロックや NCCL の無効な使用につながる可能性があります。ucc バックエンドは実験的です。デバイスのデフォルトバックエンドは get_default_backend_for_device() で取得できます。

init_method (str, optional) – プロセスグループの初期化方法を指定する URL です。init_method と store のどちらも指定されない場合、デフォルトは “env://” です。store とは同時に指定できません。

world_size (int, optional) – ジョブに参加するプロセス数です。store を指定する場合は必須です。

rank (int, optional) – 現在のプロセスのランクです（0 から world_size-1 までの値）。store を指定する場合は必須です。

store (Store, optional) – 接続先やアドレス情報の交換に使用する、すべてのワーカーからアクセス可能なキーと値のストアです。init_method とは同時に指定できません。

timeout (timedelta, optional) – プロセスグループに対して実行される操作のタイムアウトです。デフォルト値は NCCL では 10 分、その他のバックエンドでは 30 分です。この時間が経過すると、集合通信は非同期に中止され、プロセスはクラッシュします。CUDA の実行は非同期であり、失敗した非同期 NCCL 操作によって後続の CUDA 操作が破損データで実行される可能性があるため、安全にユーザーコードの実行を続けられなくなるからです。TORCH_NCCL_BLOCKING_WAIT が設定されている場合、プロセスはこのタイムアウトまでブロックして待機します。

group_name (str, optional, deprecated) – グループ名です。この引数は無視されます。

pg_options (ProcessGroupOptions, optional) – 特定のプロセスグループの構築時に渡す追加オプションを指定します。現在サポートしているのは nccl バックエンド向けの ProcessGroupNCCL.Options のみです。is_high_priority_stream を指定すると、計算カーネルの待機中に nccl バックエンドが高優先度の cuda ストリームを使用できます。nccl を設定するその他のオプションについては、https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/api/types.html#ncclconfig-t を参照してください。

device_id (torch.device | int, optional) – このプロセスが使用する単一の特定デバイスです。バックエンド固有の最適化を可能にします。現在、NCCL でのみ次の 2 つの効果があります。コミュニケーターが即座に作成されます（通常の遅延呼び出しではなく、ncclCommInit* がただちに呼び出されます）。また、可能な場合はサブグループで ncclCommSplit を使い、グループ作成の不要なオーバーヘッドを避けます。NCCL の初期化エラーを早期に確認したい場合にも、このフィールドを使用できます。int が指定された場合、API はコンパイル時のアクセラレータータイプを使用するとみなします。

注記 backend == Backend.MPI を有効にするには、MPI をサポートするシステム上で PyTorch をソースからビルドする必要があります。

注記 複数バックエンドのサポートは実験的です。現在、backend が指定されない場合、gloo と nccl の両バックエンドが作成されます。CPU テンソルを使う集合通信には gloo バックエンド、CUDA テンソルを使う集合通信には nccl バックエンドが使用されます。カスタムバックエンドは、“<device_type>:<backend_name>,<device_type>:<backend_name>” 形式の文字列（例：“cpu:gloo,cuda:custom_backend”）を渡して指定できます。

torch.distributed.device_mesh.init_device_mesh(device_type, mesh_shape, *, mesh_dim_names=None, backend_override=None)[source]# device_type、mesh_shape、mesh_dim_names の各パラメーターに基づいて DeviceMesh を初期化します。mesh_shape の長さを n とする n 次元配列のレイアウトを持つ DeviceMesh を作成します。mesh_dim_names が指定されている場合、各次元には mesh_dim_names[i] という名前が割り当てられます。

注記 init_device_mesh は SPMD プログラミングモデルに従います。つまり、クラスター内のすべてのプロセス／ランクで、同じ PyTorch Python プログラムが実行されます。デバイスレイアウトを記述する n 次元配列の次元である mesh_shape が、すべてのランクで同一であることを確認してください。mesh_shape が一致しないと、ハングする可能性があります。

注記 プロセスグループが見つからない場合、init_device_mesh は分散通信に必要なプロセスグループを内部で初期化します。

パラメーター

device_type (str) – メッシュのデバイスタイプです。現在サポートされているのは、“cpu”、“cuda/cuda-like”、“xpu” です。“cuda:0” のように GPU インデックスを含むデバイスタイプは指定できません。

mesh_shape (Tuple[int]) – デバイスのレイアウトを記述する多次元配列の次元を定義するタプルです。

mesh_dim_names (Tuple[str], optional) – デバイスのレイアウトを記述する多次元配列の各次元に割り当てるメッシュ次元名のタプルです。長さは mesh_shape の長さと一致しなければなりません。mesh_dim_names 内の各文字列は一意でなければなりません。

backend_override (Dict[int | str, tuple[str, Options] | str | Options], optional) – 各メッシュ次元に対して作成される ProcessGroups の一部またはすべてのオーバーライドです。各キーには、次元のインデックス、または mesh_dim_names が指定されている場合はその名前を指定できます。各値には、バックエンド名とオプションを含むタプル、またはこの 2 つのうちいずれか一方を指定できます。一方だけ指定した場合、もう一方にはデフォルト値が設定されます。

戻り値

デバイスのレイアウトを表す DeviceMesh オブジェクトです。

戻り値の型 DeviceMesh

例：

from torch.distributed.device_mesh import init_device_mesh

mesh_1d = init_device_mesh("cuda", mesh_shape=(8,)) mesh_2d = init_device_mesh("cuda", mesh_shape=(2, 8), mesh_dim_names=("dp", "tp"))

torch.distributed.is_initialized()[source]# デフォルトのプロセスグループが初期化済みかどうかを確認します。

戻り値の型 bool

torch.distributed.is_mpi_available()[source]# MPI バックエンドが利用可能かどうかを確認します。

戻り値の型 bool

torch.distributed.is_nccl_available()[source]# NCCL バックエンドが利用可能かどうかを確認します。

戻り値の型 bool

torch.distributed.is_gloo_available()[source]# Gloo バックエンドが利用可能かどうかを確認します。

戻り値の型 bool

torch.distributed.distributed_c10d.is_xccl_available()[source]# XCCL バックエンドが利用可能かどうかを確認します。

戻り値の型 bool

torch.distributed.is_torchelastic_launched()[source]# このプロセスが torch.distributed.elastic（別名 torchelastic）を使って起動されたかどうかを確認します。TORCHELASTIC_RUN_ID 環境変数の存在を、現在のプロセスが torchelastic によって起動されたかどうかの判定に使います。TORCHELASTIC_RUN_ID はランデブー ID に対応し、ピア検出のためのジョブ ID を示す、常に null でない値であるため、妥当な判定方法です。

戻り値の型 bool

torch.distributed.get_default_backend_for_device(device)[source]# 指定したデバイスのデフォルトバックエンドを返します。

パラメーター

device (Union[str, torch.device]) – デフォルトバックエンドを取得するデバイスです。

戻り値

指定したデバイスのデフォルトバックエンドを小文字の文字列で返します。

戻り値の型 str

現在、3 つの初期化方法がサポートされています。

TCP による初期化#

TCP を使って初期化する方法は 2 つあります。どちらも、すべてのプロセスから到達可能なネットワークアドレスと、指定する world_size が必要です。1 つ目の方法では、rank 0 プロセスに属するアドレスを指定する必要があります。この初期化方法では、すべてのプロセスにランクを手動で指定する必要があります。最新の分散パッケージでは、マルチキャストアドレスはサポートされなくなっています。group_name も非推奨です。

import torch.distributed as dist

いずれかのマシンのアドレスを使用

dist.init_process_group(backend, init_method='tcp://10.1.1.20:23456', rank=args.rank, world_size=4)

共有ファイルシステムによる初期化#

もう 1 つの初期化方法では、グループ内のすべてのマシンから共有され、アクセスできるファイルシステムと、指定する world_size を使用します。URL は file:// で始まり、共有ファイルシステム上の既存ディレクトリ内にある、まだ存在しないファイルへのパスを含む必要があります。ファイルシステムによる初期化では、ファイルが存在しない場合に自動作成されますが、ファイルは削除されません。そのため、同じファイルパス／名前で次回 init_process_group() を呼び出す前にファイルが削除されていることを確認する責任は、ユーザーにあります。最新の分散パッケージではランクの自動割り当てはサポートされなくなっており、group_name も非推奨です。

警告 この方法では、ファイルシステムが fcntl によるロックをサポートしていることを前提とします。ほとんどのローカルシステムと NFS はこれをサポートしています。

警告 この方法では常にファイルが作成され、プログラムの終了時にファイルをクリーンアップして削除しようとします。つまり、ファイル初期化メソッドを使って初期化するたびに、新しい空のファイルが必要です。前回の初期化で使ったファイル（削除されずに残ったもの）を再利用するのは想定外の動作であり、デッドロックや失敗を引き起こすことがあります。この方法ではファイルのクリーンアップを試みますが、自動削除に失敗した場合は、次回同じファイルが再利用されないよう、トレーニング終了時にファイルを削除する責任はユーザーにあります。同じファイル名で init_process_group() を複数回呼び出す場合は、特に注意が必要です。つまり、ファイルが削除／クリーンアップされないまま、そのファイルで再度 init_process_group() を呼び出すと、失敗が予想されます。経験則として、init_process_group() を呼び出すたびに、ファイルが存在しないか空であることを確認してください。

import torch.distributed as dist

rank は必ず指定する

dist.init_process_group(backend, init_method='file:///mnt/nfs/sharedfile', world_size=4, rank=args.rank)

環境変数による初期化#

この方法では環境変数から設定を読み取り、情報の取得方法を完全にカスタマイズできます。設定が必要な変数は次のとおりです。

MASTER_PORT - 必須。rank 0 のマシンで空いているポートである必要があります。

MASTER_ADDR - 必須（rank 0 を除く）。rank 0 ノードのアドレスです。

WORLD_SIZE - 必須。ここで設定するか、初期化関数の呼び出しで設定できます。

RANK - 必須。ここで設定するか、初期化関数の呼び出しで設定できます。

rank 0 のマシンを使って、すべての接続を設定します。これはデフォルトの方法であり、init_method を指定する必要はありません（または env:// を指定できます）。

初期化時間の短縮#

TORCH_GLOO_LAZY_INIT - フルメッシュの代わりに、必要に応じて接続を確立します。これにより、all2all 以外の操作で初期化時間を大幅に短縮できます。

初期化後#

torch.distributed.init_process_group() の実行後、次の関数を使えます。プロセスグループがすでに初期化されているかどうかを確認するには、torch.distributed.is_initialized() を使います。

class torch.distributed.Backend(name)[source]# バックエンドの列挙型に似たクラスです。利用可能なバックエンドは GLOO、NCCL、UCC、MPI、XCCL、およびその他の登録済みバックエンドです。このクラスの値は小文字の文字列です（例："gloo"）。Backend.NCCL のように属性としてアクセスできます。このクラスを直接呼び出して文字列を解析することもできます（例：Backend(backend_str) は backend_str が有効かどうかを確認し、有効な場合は解析済みの小文字の文字列を返します）。大文字の文字列も受け付けます（例：Backend("GLOO") は "gloo" を返します）。

注記 Backend.UNDEFINED は存在しますが、一部のフィールドの初期値としてのみ使用されます。ユーザーが直接使用したり、その存在を前提にしたりしないでください。

classmethod register_backend(name, func, extended_api=False, devices=None)[source]# 指定した名前とインスタンス化関数を使って、新しいバックエンドを登録します。このクラスメソッドは、サードパーティーの ProcessGroup 拡張機能が新しいバックエンドを登録するために使います。

パラメーター

name (str) – ProcessGroup 拡張機能のバックエンド名です。init_process_group() で使う名前と一致する必要があります。

func (function) – バックエンドをインスタンス化する関数ハンドラーです。バックエンド拡張機能で実装し、store、rank、world_size、timeout の 4 つの引数を受け取る必要があります。

extended_api (bool, optional) – バックエンドが拡張引数構造をサポートするかどうかです。デフォルト：False。True に設定すると、バックエンドは c10d::DistributedBackendOptions のインスタンスと、バックエンドの実装で定義されたプロセスグループオプションオブジェクトを受け取ります。

device (str or list of str, optional) – このバックエンドがサポートするデバイスタイプです（例：“cpu”、“cuda” など）。None の場合は、“cpu” と “cuda” の両方を想定します。

注記 サードパーティーバックエンドのサポートは実験的であり、今後変更される可能性があります。

torch.distributed.get_backend(group=None)[source]# 指定したプロセスグループのバックエンドを返します。

パラメーター

group (ProcessGroup, optional) – 操作対象のプロセスグループです。デフォルトは全体のメインプロセスグループです。別の特定グループを指定する場合、呼び出し元プロセスはそのグループの一員でなければなりません。

戻り値

指定したプロセスグループのバックエンドを小文字の文字列で返します。

戻り値の型 Backend

torch.distributed.get_rank(group=None)[source]# 指定したグループ（指定がない場合はデフォルト）における現在のプロセスのランクを返します。ランクは、分散プロセスグループ内の各プロセスに割り当てられる一意の識別子です。常に 0 から world_size までの連続した整数です。

パラメーター

group (ProcessGroup, optional) – 操作対象のプロセスグループです。None の場合、デフォルトのプロセスグループが使われます。

戻り値

プロセスグループのランク。グループに属していない場合は -1。

戻り値の型 int

torch.distributed.get_world_size(group=None)[source]# 現在のプロセスグループ内のプロセス数を返します。

パラメーター

group (ProcessGroup, optional) – 操作対象のプロセスグループです。None の場合、デフォルトのプロセスグループが使われます。

戻り値

プロセスグループの world size。グループに属していない場合は -1。

戻り値の型 int

シャットダウン#

終了時に destroy_process_group() を呼び出してリソースをクリーンアップすることが重要です。最も簡単な方法は、通信が不要になった時点（通常は main() の終わり近く）で、group 引数にデフォルト値 None を指定して destroy_process_group() を呼び出し、すべてのプロセスグループとバックエンドを破棄することです。この呼び出しは、外側のプロセスランチャーではなく、各トレーナープロセスで 1 回ずつ行う必要があります。特に N 次元並列処理などでアプリケーション内に複数のプロセスグループがある場合、タイムアウト時間内に pg 内のすべてのランクが destroy_process_group() を呼び出さないと、終了時にハングする可能性があります。これは、ProcessGroupNCCL のデストラクターが ncclCommAbort を呼び出す必要があり、この呼び出しは全プロセスで行わなければならない一方、Python の GC による ProcessGroupNCCL のデストラクター呼び出し順序は決定的ではないためです。destroy_process_group() を呼び出すと、ランク間で ncclCommAbort が一貫した順序で呼び出され、ProcessGroupNCCL のデストラクター内で ncclCommAbort が呼び出されるのを避けられます。

再初期化#

destroy_process_group は個々のプロセスグループの破棄にも使えます。たとえばフォールトトレラントトレーニングでは、実行中にプロセスグループを破棄し、新しいものを初期化する場合があります。この場合、destroy の呼び出し後、次の初期化の前に、torch.distributed のプリミティブ以外の方法でトレーナープロセスを同期することが重要です。この動作は現在サポートもテストもされていません。同期の実現が難しいため、既知の問題とみなされています。この用途が妨げになっている場合は、github issue または RFC を提出してください。

グループ#

デフォルトでは、集合通信はデフォルトグループ（world とも呼ばれます）に対して動作し、すべてのプロセスが分散関数呼び出しに参加する必要があります。ただし、ワークロードによっては、より細かい通信が有効な場合があります。ここで分散グループが役立ちます。new_group() 関数を使うと、全プロセスの任意のサブセットを含む新しいグループを作成できます。この関数は、すべての集合通信に group 引数として渡せる不透明なグループハンドルを返します（集合通信とは、よく知られたプログラミングパターンで情報を交換する分散関数です）。

torch.distributed.new_group(ranks=None, timeout=None, backend=None, pg_options=None, use_local_synchronization=False, group_desc=None, device_id=None)[source]# 新しい分散グループを作成します。グループのメンバーにならないプロセスも含め、メイングループ内のすべてのプロセス（つまり分散ジョブに参加するすべてのプロセス）がこの関数を呼び出す必要があります。また、すべてのプロセスでグループを同じ順序で作成する必要があります。

警告 安全な同時使用：NCCL バックエンドで複数のプロセスグループを使う場合、ユーザーはランク間の集合通信の実行順序がグローバルに一貫していることを確認する必要があります。プロセス内の複数のスレッドが集合通信を発行する場合は、順序の一貫性を保つために明示的な同期が必要です。torch.distributed 通信 API の非同期版を使うと、work オブジェクトが返され、通信カーネルは別の CUDA ストリームにキューイングされ、通信と計算のオーバーラップが可能になります。あるプロセスグループで 1 つ以上の非同期操作を発行した後、別のプロセスグループを使用する前に、work.wait() を呼び出して他の cuda ストリームと同期する必要があります。詳細は Using multiple NCCL communicators concurrently https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/usage/communicators.html#using-multiple-nccl-communicators-concurrently を参照してください。

パラメーター

ranks (list[int]) – グループメンバーのランクのリストです。None の場合、すべてのランクに設定されます。デフォルトは None です。

timeout (timedelta, optional) – 詳細とデフォルト値については init_process_group を参照してください。

backend (str or Backend, optional) – 使用するバックエンドです。ビルド時の設定によって、有効な値は gloo と nccl です。デフォルトではグローバルグループと同じバックエンドを使用します。このフィールドには小文字の文字列（例："gloo"）を指定します。Backend 属性（例：Backend.GLOO）からも参照できます。None を渡した場合は、デフォルトのプロセスグループに対応するバックエンドが使われます。デフォルトは None です。

pg_options (ProcessGroupOptions, optional) – 特定のプロセスグループの構築時に渡す追加オプションを指定します。たとえば nccl バックエンドでは、プロセスグループが高優先度の cuda ストリームを使用できるように、is_high_priority_stream を指定できます。nccl を設定するその他のオプションについては、https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/api/types.html#ncclconfig-t を参照してください。

use_local_synchronization (bool, optional)：プロセスグループ作成の最後に、グループ内だけのバリアを実行します。グループに属さないランクは API を呼び出す必要がなく、バリアにも参加しない点が異なります。

group_desc (str, optional) – プロセスグループを説明する文字列です。

device_id (torch.device, optional) – このプロセスを「バインド」する単一の特定デバイスです。このフィールドを指定すると、new_group の呼び出しでは、そのデバイス用の通信バックエンドをただちに初期化しようとします。

戻り値

集合通信の呼び出しに渡せる分散グループのハンドル、またはランクが ranks に含まれない場合は GroupMember.NON_GROUP_MEMBER です。

注：use_local_synchronization は MPI では動作しません。

注：use_local_synchronization=True は、大規模クラスターや小規模プロセスグループで大幅な高速化につながる場合があります。ただし、グループに属さないランクは group barrier() に参加しないため、クラスターの動作が変化する点に注意してください。

注：use_local_synchronization=True にすると、各ランクが重複する複数のプロセスグループを作成する場合にデッドロックが起きる可能性があります。これを避けるには、すべてのランクでグローバルな作成順序を同じにしてください。

torch.distributed.get_group_rank(group, global_rank)[source]# グローバルランクをグループランクに変換します。global_rank が group に含まれていない場合、RuntimeError が発生します。

パラメーター

group (ProcessGroup) – 相対ランクを取得する ProcessGroup です。

global_rank (int) – 照会するグローバルランクです。

戻り値

group における global_rank のグループランクです。

戻り値の型 int

注：デフォルトのプロセスグループに対してこの関数を呼び出すと、恒等変換になります。

torch.distributed.get_global_rank(group, group_rank)[source]# グループランクをグローバルランクに変換します。group_rank が group に含まれていない場合、RuntimeError が発生します。

パラメーター

group (ProcessGroup) – グローバルランクを取得する ProcessGroup です。

group_rank (int) – 照会するグループランクです。

戻り値

group における group_rank のグローバルランクです。

戻り値の型 int

注：デフォルトのプロセスグループに対してこの関数を呼び出すと、恒等変換になります。

torch.distributed.get_process_group_ranks(group)[source]# group に関連付けられたすべてのランクを取得します。

パラメーター

group (Optional[ProcessGroup]) – すべてのランクを取得する ProcessGroup です。None の場合、デフォルトのプロセスグループが使われます。

戻り値

グループランク順に並んだグローバルランクのリストです。

戻り値の型 list[int]

DeviceMesh#

DeviceMesh はプロセスグループ（または NCCL コミュニケーター）を管理する高水準の抽象化です。ランクの設定方法を意識せずにノード間およびノード内のプロセスグループを簡単に作成でき、分散プロセスグループの管理にも役立ちます。init_device_mesh() 関数を使うと、デバイストポロジーを記述するメッシュ形状で新しい DeviceMesh を作成できます。

class torch.distributed.device_mesh.DeviceMesh(device_type, mesh, *, mesh_dim_names=None, backend_override=None, _init_backend=True)[source]# DeviceMesh はデバイスのメッシュを表します。デバイスレイアウトは n 次元配列で表され、n 次元配列の各値はデフォルトプロセスグループのグローバルランク ID です。DeviceMesh を使うと、クラスター全体の N 次元デバイス接続を設定し、N 次元並列処理の ProcessGroups を管理できます。DeviceMesh の各次元で個別に通信できます。

DeviceMesh は、ユーザーがすでに選択したデバイスを尊重します（つまり、DeviceMesh の初期化前に torch.cuda.set_device を呼び出した場合）。事前にデバイスを設定していない場合は、現在のプロセスのデバイスを選択／設定します。手動でデバイスを選択する場合は、必ず DeviceMesh の初期化前に行ってください。

DeviceMesh は DTensor API と組み合わせて使用する場合、コンテキストマネージャーとしても使えます。

注記 DeviceMesh は SPMD プログラミングモデルに従います。つまり、クラスター内のすべてのプロセス／ランクで同じ PyTorch Python プログラムが実行されます。したがって、デバイスのレイアウトを記述するメッシュ配列がすべてのランクで同一であることを確認してください。メッシュが一致しないと、気づかないうちにハングします。

パラメーター

device_type (str) – メッシュのデバイスタイプです。現在サポートされているのは、“cpu”、“cuda/cuda-like” です。

mesh (ndarray) – デバイスのレイアウトを記述する多次元配列または整数テンソルです。ID はデフォルトのプロセスグループのグローバル ID です。

戻り値

デバイスのレイアウトを表す DeviceMesh オブジェクトです。

戻り値の型 DeviceMesh

次のプログラムは、SPMD 方式で各プロセス／ランク上で実行されます。この例では、各ホストに 4 つの GPU があるホストが 2 台あります。メッシュの最初の次元に沿った reduction は列（0, 4）、…、（3, 7）を対象に実行され、2 番目の次元に沿った reduction は行（0, 1, 2, 3）と（4, 5, 6, 7）を対象に実行されます。

例：

from torch.distributed.device_mesh import DeviceMesh

ホスト間（次元 0）とホスト内（次元 1）のトポロジーを表すため、

デバイスメッシュを (2, 4) として初期化します。

mesh = DeviceMesh(device_type="cuda", mesh=[[0, 1, 2, 3],[4, 5, 6, 7]])

static from_group(group, device_type, mesh=None, *, mesh_dim_names=None)[source]# 既存の ProcessGroup または既存の ProcessGroup のリストから、device_type を指定して DeviceMesh を構築します。構築されるデバイスメッシュの次元数は、渡されたグループの数と同じです。たとえば、単一のプロセスグループを渡すと 1D メッシュになり、2 つのプロセスグループのリストを渡すと 2D メッシュになります。複数のグループを渡す場合は、mesh と mesh_dim_names 引数が必須です。渡されたプロセスグループの順序によって、メッシュのトポロジーが決まります。たとえば、最初のプロセスグループは DeviceMesh の 0 番目の次元になります。渡す mesh テンソルの次元数は、渡したプロセスグループの数と一致する必要があります。また、mesh テンソル内の次元の順序は、渡すプロセスグループの順序と一致しなければなりません。

パラメーター

group (ProcessGroup or list[ProcessGroup]) – 既存の ProcessGroup または既存の ProcessGroup のリストです。

device_type (str) – メッシュのデバイスタイプです。現在サポートされているのは、“cpu”、“cuda/cuda-like” です。“cuda:0” のように GPU インデックスを含むデバイスタイプは指定できません。

mesh (torch.Tensor or ArrayLike, optional) – デバイスのレイアウトを記述する多次元配列または整数テンソルです。ID はデフォルトのプロセスグループのグローバル ID です。デフォルトは None です。

mesh_dim_names (tuple[str], optional) – デバイスのレイアウトを記述する多次元配列の各次元に割り当てるメッシュ次元名のタプルです。長さは mesh_shape の長さと一致しなければなりません。mesh_dim_names 内の各文字列は一意でなければなりません。デフォルトは None です。

戻り値

デバイスのレイアウトを表す DeviceMesh オブジェクトです。

戻り値の型 DeviceMesh

get_all_groups()[source]# すべてのメッシュ次元の ProcessGroup のリストを返します。

戻り値

ProcessGroup オブジェクトのリストです。

戻り値の型 list[torch.distributed.distributed_c10d.ProcessGroup]

get_coordinate()[source]# メッシュのすべての次元に対する、このランクの相対インデックスを返します。このランクがメッシュに含まれていない場合は None を返します。

戻り値の型 Optional[list[int]]

get_group(mesh_dim=None)[source]# mesh_dim で指定された単一の ProcessGroup を返します。mesh_dim が指定されず、DeviceMesh が 1 次元の場合は、メッシュ内の唯一の ProcessGroup を返します。

パラメーター

mesh_dim (str/python:int, optional) – メッシュ次元の名前またはインデックスです。デフォルトは None です。

戻り値

ProcessGroup オブジェクトです。

戻り値の型 ProcessGroup

get_local_rank(mesh_dim=None)[source]# DeviceMesh の指定した mesh_dim におけるローカルランクを返します。

パラメーター

mesh_dim (str/python:int, optional) – メッシュ次元の名前またはインデックスです。デフォルトは None です。

戻り値

ローカルランクを表す整数です。

戻り値の型 int

次のプログラムは、SPMD 方式で各プロセス／ランク上で実行されます。この例では、各ホストに 4 つの GPU があるホストが 2 台あります。ランク 0、1、2、3 で mesh_2d.get_local_rank(mesh_dim=0) を呼び出すと 0 が返されます。ランク 4、5、6、7 で mesh_2d.get_local_rank(mesh_dim=0) を呼び出すと 1 が返されます。ランク 0、4 で mesh_2d.get_local_rank(mesh_dim=1) を呼び出すと 0 が返されます。ランク 1、5 では 1、ランク 2、6 では 2、ランク 3、7 では 3 が返されます。

例：

from torch.distributed.device_mesh import DeviceMesh

ホスト間（次元 0）とホスト内（次元 1）のトポロジーを表すため、

デバイスメッシュを (2, 4) として初期化します。

mesh = DeviceMesh(device_type="cuda", mesh=[[0, 1, 2, 3],[4, 5, 6, 7]])

get_rank()[source]# 現在のグローバルランクを返します。

戻り値の型 int

ポイントツーポイント通信#

torch.distributed.send(tensor, dst=None, group=None, tag=0, group_dst=None)[source]# テンソルを同期的に送信します。

警告 NCCL バックエンドでは tag はサポートされていません。

パラメーター

tensor (Tensor) – 送信するテンソルです。

dst (int) – グローバルプロセスグループ内の宛先ランクです（group 引数に関係なく）。宛先ランクは現在のプロセスのランクと同じであってはなりません。

group (ProcessGroup, optional) – 操作対象のプロセスグループです。None の場合、デフォルトのプロセスグループが使われます。

tag (int, optional) – 送信とリモートの recv を対応付けるタグです。

group_dst (int, optional) – group 内の宛先ランクです。dst と group_dst を同時に指定することはできません。

torch.distributed.recv(tensor, src=None, group=None, tag=0, group_src=None)[source]# テンソルを同期的に受信します。

警告 NCCL バックエンドでは tag はサポートされていません。

パラメーター

tensor (Tensor) – 受信データを格納するテンソルです。

src (int, optional) – グローバルプロセスグループ内の送信元ランクです（group 引数に関係なく）。指定しない場合、任意のプロセスから受信します。

group (ProcessGroup, optional) – 操作対象のプロセスグループです。None の場合、デフォルトのプロセスグループが使われます。

tag (int, optional) – recv とリモートの send を対応付けるタグです。

group_src (int, optional) – group 内の宛先ランクです。src と group_src を同時に指定することはできません。

戻り値

送信側のランクです。グループに属していない場合は -1 です。

戻り値の型 int

isend() と irecv() を使うと、分散リクエストオブジェクトが返されます。通常、このオブジェクトの型は指定されていません。手動で作成するものではありませんが、次の 2 つのメソッドをサポートすることが保証されています。

is_completed() - 操作が完了していれば True を返します。

wait() - 操作が完了するまでプロセスをブロックします。is_completed() は、いったん True を返すと、以後も True を返すことが保証されています。

torch.distributed.isend(tensor, dst=None, group=None, tag=0, group_dst=None)[source]# テンソルを非同期に送信します。

警告 リクエストが完了する前にテンソルを変更すると、未定義の動作になります。

警告 NCCL バックエンドでは tag はサポートされていません。

ブロッキングの send と異なり、isend では src == dst ランク、つまり自分自身への送信が可能です。

パラメーター

tensor (Tensor) – 送信するテンソルです。

dst (int) – グローバルプロセスグループ内の宛先ランクです（group 引数に関係なく）。

group (ProcessGroup, optional) – 操作対象のプロセスグループです。None の場合、デフォルトのプロセスグループが使われます。

tag (int, optional) – send とリモートの recv を対応付けるタグです。

group_dst (int, optional) – group 内の宛先ランクです。dst と group_dst を同時に指定することはできません。

戻り値

分散リクエストオブジェクトです。グループに属していない場合は None です。

戻り値の型 Optional[Work]

torch.distributed.irecv(tensor, src=None, group=None, tag=0, group_src=None)[source]# テンソルを非同期に受信します。

警告 NCCL バックエンドでは tag はサポートされていません。

ブロッキングの recv と異なり、irecv では src == dst ランク、つまり自分自身からの受信が可能です。

パラメーター

tensor (Tensor) – 受信データを格納するテンソルです。

src (int, optional) – グローバルプロセスグループ内の送信元ランクです（group 引数に関係なく）。指定しない場合、任意のプロセスから受信します。

group (ProcessGroup, optional) – 操作対象のプロセスグループです。None の場合、デフォルトのプロセスグループが使われます。

tag (int, optional) – recv とリモートの send を対応付けるタグです。

group_src (int, optional) – group 内の宛先ランクです。src と group_src を同時に指定することはできません。

戻り値

分散リクエストオブジェクトです。グループに属していない場合は None です。

戻り値の型 Optional[Work]

torch.distributed.send_object_list(object_list, dst=None, group=None, device=None, group_dst=None, use_batch=False)[source]# object_list 内の pickle 化可能なオブジェクトを同期的に送信します。send() と同様ですが、Python オブジェクトを渡せます。送信するには、object_list 内のすべてのオブジェクトが pickle 化可能でなければなりません。

パラメーター

object_list (List[Any]) – 送信する入力オブジェクトのリストです。各オブジェクトは pickle 化可能である必要があります。受信側も同じサイズのリストを用意する必要があります。

dst (int) – object_list の送信先ランクです。宛先ランクはグローバルプロセスグループに基づきます（group 引数に関係なく）。

group (Optional[ProcessGroup]) – 操作対象のプロセスグループです。None の場合、デフォルトのプロセスグループが使われます。デフォルトは None です。

device (torch.device, optional) – None でない場合、オブジェクトをシリアライズしてテンソルに変換し、送信前にデバイスへ移動します。デフォルトは None です。

group_dst (int, optional) – group 内の宛先ランクです。dst と group_dst のどちらか一方を指定する必要があります。

use_batch (bool, optional) – True の場合、通常の送信操作の代わりにバッチ p2p 操作を使います。これにより 2 ランクのコミュニケーターの初期化を避け、既存のグループ全体のコミュニケーターを利用できます。使い方と前提条件については batch_isend_irecv を参照してください。デフォルトは False です。

戻り値 None。

注記 NCCL ベースのプロセスグループでは、通信を行う前にオブジェクトの内部テンソル表現を GPU デバイスに移動する必要があります。この場合、使用するデバイスは torch.cuda.current_device() で指定されます。torch.cuda.set_device() を使って各ランクが個別の GPU を持つよう設定する責任は、ユーザーにあります。

警告 オブジェクト集合通信には、深刻なパフォーマンスおよびスケーラビリティ上の制約があります。詳細は Object collectives を参照してください。

警告 send_object_list() は pickle モジュールを暗黙的に使用しますが、これは安全ではないことが知られています。悪意のある pickle データを作成し、unpickle 時に任意のコードを実行させることができます。信頼できるデータに対してのみ、この関数を呼び出してください。

警告 GPU テンソルを使って send_object_list() を呼び出すことは、サポートが十分でなく、非効率です。テンソルが pickle 化されるため、GPU から CPU への転送が発生します。代わりに send() の使用を検討してください。

例：::

注：各ランクでのプロセスグループ初期化は省略しています。

import torch.distributed as dist

バックエンドが NCCL でないことを前提とします。

device = torch.device("cpu") if dist.get_rank() == 0:

world_size が 2 であることを前提とします。

objects = ["foo", 12, {1: 2}] # pickle 化可能な任意のオブジェクト dist.send_object_list(objects, dst=1, device=device) else: objects = [None, None, None] dist.recv_object_list(objects, src=0, device=device) objects ['foo', 12, {1: 2}]

torch.distributed.recv_object_list(object_list, src=None, group=None, device=None, group_src=None, use_batch=False)[source]# object_list 内の pickle 化可能なオブジェクトを同期的に受信します。recv() と同様ですが、Python オブジェクトを受信できます。

パラメーター

object_list (List[Any]) – 受信オブジェクトを格納するリストです。送信されるリストと同じサイズのリストを用意する必要があります。

src (int, optional) – object_list の受信元ランクです。送信元ランクはグローバルプロセスグループに基づきます（group 引数に関係なく）。None の場合、任意のランクから受信します。デフォルトは None です。

group (Optional[ProcessGroup]) – 操作対象のプロセスグループです。None の場合、デフォルトのプロセスグループが使われます。デフォルトは None です。

device (torch.device, optional) – None でない場合、このデバイスで受信します。デフォルトは None です。

group_src (int, optional) – group 内の宛先ランクです。src と group_src を同時に指定することはできません。

use_batch (bool, optional) – True の場合、通常の送信操作の代わりにバッチ p2p 操作を使います。これにより 2 ランクのコミュニケーターの初期化を避け、既存のグループ全体のコミュニケーターを利用できます。使い方と前提条件については batch_isend_irecv を参照してください。デフォルトは False です。

戻り値

送信側のランクです。グループに属していない場合は -1 です。グループに属している場合、object_list に src ランクから送信されたオブジェクトが格納されます。

注記 NCCL ベースのプロセスグループでは、通信を行う前にオブジェクトの内部テンソル表現を GPU デバイスに移動する必要があります。この場合、使用するデバイスは torch.cuda.current_device() で指定されます。torch.cuda.set_device() を使って各ランクが個別の GPU を持つよう設定する責任は、ユーザーにあります。

警告 オブジェクト集合通信には、深刻なパフォーマンスおよびスケーラビリティ上の制約があります。詳細は Object collectives を参照してください。

警告 recv_object_list() は pickle モジュールを暗黙的に使用しますが、これは安全ではないことが知られています。悪意のある pickle データを作成し、unpickle 時に任意のコードを実行させることができます。信頼できるデータに対してのみ、この関数を呼び出してください。

警告 GPU テンソルを使って recv_object_list() を呼び出すことは、サポートが十分でなく、非効率です。テンソルが pickle 化されるため、GPU から CPU への転送が発生します。代わりに recv() の使用を検討してください。

例：::

注：各ランクでのプロセスグループ初期化は省略しています。

import torch.distributed as dist

バックエンドが NCCL でないことを前提とします。

device = torch.device("cpu") if dist.get_rank() == 0:

world_size が 2 であることを前提とします。

objects = ["foo", 12, {1: 2}] # pickle 化可能な任意のオブジェクト dist.send_object_list(objects, dst=1, device=device) else: objects = [None, None, None] dist.recv_object_list(objects, src=0, device=device) objects ['foo', 12, {1: 2}]

torch.distributed.batch_isend_irecv(p2p_op_list)[source]# テンソルのバッチを非同期に送受信し、リクエストのリストを返します。p2p_op_list 内の各操作を実行し、対応するリクエストを返します。現在、NCCL、Gloo、UCC バックエンドがサポートされています。

パラメーター

p2p_op_list (list[torch.distributed.distributed_c10d.P2POp]) – ポイントツーポイント操作（各演算子の型は torch.distributed.P2POp）のリストです。リスト内の isend／irecv の順序は重要で、リモート側の対応する isend／irecv と一致させる必要があります。

戻り値

op_list 内の対応する op を呼び出して返された、分散リクエストオブジェクトのリストです。

戻り値の型 list[torch.distributed.distributed_c10d.Work]

例

send_tensor = torch.arange(2, dtype=torch.float32) + 2 * rank recv_tensor = torch.randn(2, dtype=torch.float32) send_op = dist.P2POp(dist.isend, send_tensor, (rank + 1) % world_size) recv_op = dist.P2POp( ... ... dist.irecv, recv_tensor, (rank - 1 + world_size) % world_size ... ... ) reqs = batch_isend_irecv([send_op, recv_op]) for req in reqs: req.wait() recv_tensor tensor([2, 3]) # ランク 0 tensor([0, 1]) # ランク 1

注記 この API を NCCL PG バックエンドで使う場合、ユーザーは torch.cuda.set_device を使って現在の GPU デバイスを設定する必要があります。設定しないと、予期しないハングが発生する可能性があります。また、この API が dist.P2POp に渡したグループで最初の集合通信呼び出しである場合、グループ内のすべてのランクがこの API 呼び出しに参加する必要があります。そうでない場合、動作は未定義です。この API が最初の集合通信呼び出しではない場合、グループ内の一部のランクのみが参加するバッチ P2P 操作が許可されます。

class torch.distributed.P2POp(op, tensor, peer=None, group=None, tag=0, group_peer=None)[source]# batch_isend_irecv 用のポイントツーポイント操作を構築するクラスです。このクラスは P2P 操作の種類、通信バッファー、ピアランク、Process Group、タグを構築します。このクラスのインスタンスは、ポイントツーポイント通信のために batch_isend_irecv に渡します。

パラメーター

op (Callable) – ピアプロセスへのデータ送信またはピアプロセスからのデータ受信を行う関数です。op の型は torch.distributed.isend または torch.distributed.irecv のいずれかです。

tensor (Tensor) – 送信または受信するテンソルです。

peer (int, optional) – 宛先または送信元のランクです。

group (ProcessGroup, optional) – 操作対象のプロセスグループです。None の場合、デフォルトのプロセスグループが使われます。

tag (int, optional) – send と recv を対応付けるタグです。

group_peer (int, optional) – 宛先または送信元のランクです。

同期および非同期の集合通信操作#

すべての集合通信関数は、集合通信に渡された async_op フラグの設定に応じて、次の 2 種類の操作をサポートします。

同期操作 — async_op が False の場合のデフォルトモードです。関数が戻る時点で、集合通信操作が実行されていることが保証されます。CUDA 操作は非同期であるため、CUDA 操作が完了していることまでは保証されません。CPU の集合通信では、集合通信の出力を利用する後続の関数呼び出しは期待どおりに動作します。CUDA の集合通信では、同じ CUDA ストリーム上で出力を利用する関数呼び出しは期待どおりに動作します。異なるストリームで実行する場合、ユーザーは同期に注意する必要があります。ストリーム同期などの CUDA セマンティクスの詳細については、CUDA Semantics を参照してください。次のスクリプトでは、CPU 操作と CUDA 操作のセマンティクスの違いを例で示しています。

非同期操作 — async_op が True の場合です。集合通信関数は分散リクエストオブジェクトを返します。通常、このオブジェクトを手動で作成する必要はありません。次の 2 つのメソッドをサポートすることが保証されています。

is_completed() - CPU の集合通信では、完了している場合に True を返します。CUDA 操作では、操作が CUDA ストリームに正常にキューイングされ、追加の同期なしにデフォルトストリームで出力を利用できる場合に True を返します。

wait() - CPU の集合通信では、操作が完了するまでプロセスをブロックします。CUDA 操作では、操作が完了するまで現在アクティブな CUDA ストリームをブロックします（ただし CPU はブロックしません）。

get_future() - torch._C.Future オブジェクトを返します。NCCL でサポートされており、ピアツーピア操作を除く GLOO と MPI のほとんどの操作でもサポートされています。

注：Future の採用を進め API を統合するにつれ、get_future() の呼び出しが不要になる可能性があります。

例

次のコードは、分散集合通信を使う CUDA 操作のセマンティクスを確認する際の参考になります。集合通信の出力を異なる CUDA ストリームで使う場合に、明示的な同期が必要であることを示しています。

各ランクで実行されます。

dist.init_process_group("nccl", rank=rank, world_size=2) output = torch.tensor([rank]).cuda(rank) s = torch.cuda.Stream() handle = dist.all_reduce(output, async_op=True)

wait は操作のキューイングを保証しますが、完了までは保証しません。

handle.wait()

デフォルト以外のストリームで結果を使用します。

with torch.cuda.stream(s): s.wait_stream(torch.cuda.default_stream()) output.add_(100) if rank == 0: # 明示的な wait_stream 呼び出しを省略した場合、allreduce が加算後に # 値を上書きするかどうかによって、以下の出力は非決定的に 1 または 101 になります。 print(output)

集合通信関数#

torch.distributed.broadcast(tensor, src=None, group=None, async_op=False, group_src=None)[source]# テンソルをグループ全体にブロードキャストします。集合通信に参加するすべてのプロセスで、tensor の要素数は同じでなければなりません。

パラメーター

tensor (Tensor) – src が現在のプロセスのランクである場合に送信するデータで、それ以外の場合は受信データを保存するテンソルです。

src (int) – グローバルプロセスグループ内の送信元ランクです（group 引数に関係なく）。

group (ProcessGroup, optional) – 操作対象のプロセスグループです。None の場合、デフォルトのプロセスグループが使われます。

async_op (bool, optional) – この操作を非同期操作にするかどうかです。

group_src (int) – group 内の送信元ランクです。group_src と src のどちらか一方を指定する必要があります。

戻り値

async_op が True の場合は非同期作業ハンドルを返します。async_op が False の場合、またはグループに属していない場合は None を返します。

torch.distributed.broadcast_object_list(object_list, src=None, group=None, device=None, group_src=None)[source]# object_list 内の pickle 化可能なオブジェクトをグループ全体にブロードキャストします。broadcast() と同様ですが、Python オブジェクトを渡せます。ブロードキャストするには、object_list 内のすべてのオブジェクトが pickle 化可能でなければなりません。

パラメーター

object_list (List[Any]) – ブロードキャストする入力オブジェクトのリストです。各オブジェクトは pickle 化可能である必要があります。ブロードキャストされるのは src ランク上のオブジェクトのみですが、すべてのランクで同じサイズのリストを用意する必要があります。

src (int) – object_list をブロードキャストする送信元ランクです。送信元ランクはグローバルプロセスグループに基づきます（group 引数に関係なく）。

group (Optional[ProcessGroup]) – 操作対象のプロセスグループです。None の場合、デフォルトのプロセスグループが使われます。デフォルトは None です。

device (torch.device, optional) – None でない場合、オブジェクトをシリアライズしてテンソルに変換し、ブロードキャスト前にデバイスへ移動します。デフォルトは None です。

group_src (int) – group 内の送信元ランクです。group_src と src のどちらか一方を指定する必要があります。

戻り値 None。ランクがグループに属している場合、object_list には src ランクからブロードキャストされたオブジェクトが格納されます。

注記 NCCL ベースのプロセスグループでは、通信を行う前にオブジェクトの内部テンソル表現を GPU デバイスに移動する必要があります。この場合、使用するデバイスは torch.cuda.current_device() で指定されます。torch.cuda.set_device() を使って各ランクが個別の GPU を持つよう設定する責任は、ユーザーにあります。

注記 この API は broadcast() 集合通信と少し異なり、async_op ハンドルを提供しないため、ブロッキング呼び出しになります。

警告 オブジェクト集合通信には、深刻なパフォーマンスおよびスケーラビリティ上の制約があります。詳細は Object collectives を参照してください。

警告 broadcast_object_list() は pickle モジュールを暗黙的に使用しますが、これは安全ではないことが知られています。悪意のある pickle データを作成し、unpickle 時に任意のコードを実行させることができます。信頼できるデータに対してのみ、この関数を呼び出してください。

警告 GPU テンソルを使って broadcast_object_list() を呼び出すことは、サポートが十分でなく、非効率です。テンソルが pickle 化されるため、GPU から CPU への転送が発生します。代わりに broadcast() の使用を検討してください。

例：::

注：各ランクでのプロセスグループ初期化は省略しています。

import torch.distributed as dist if dist.get_rank() == 0:

world_size が 3 であることを前提とします。

objects = ["foo", 12, {1: 2}] # pickle 化可能な任意のオブジェクト else: objects = [None, None, None]

バックエンドが NCCL でないことを前提とします。

device = torch.device("cpu") dist.broadcast_object_list(objects, src=0, device=device) objects ['foo', 12, {1: 2}]

torch.distributed.all_reduce(tensor, op=<RedOpType.SUM: 0>, group=None, async_op=False)[source]# すべてのマシン間でテンソルデータを reduction し、全プロセスが最終結果を得ます。呼び出し後、tensor はすべてのプロセスでビット単位まで同一になります。複素数テンソルもサポートされています。

パラメーター

tensor (Tensor) – 集合通信の入力と出力です。この関数はインプレースで操作します。

op (optional) – torch.distributed.ReduceOp 列挙型の値のいずれかです。要素ごとの reduction に使う操作を指定します。

group (ProcessGroup, optional) – 操作対象のプロセスグループです。None の場合、デフォルトのプロセスグループが使われます。

async_op (bool, optional) – この操作を非同期操作にするかどうかです。

戻り値

async_op が True の場合は非同期作業ハンドルを返します。async_op が False の場合、またはグループに属していない場合は None を返します。

例

以下のすべてのテンソルは torch.int64 型です。

2 つのプロセスグループ、2 つのランクがあるものとします。

device = torch.device(f"cuda:{rank}") tensor = torch.arange(2, dtype=torch.int64, device=device) + 1 + 2 * rank tensor tensor([1, 2], device='cuda:0') # ランク 0 tensor([3, 4], device='cuda:1') # ランク 1 dist.all_reduce(tensor, op=ReduceOp.SUM) tensor tensor([4, 6], device='cuda:0') # ランク 0 tensor([4, 6], device='cuda:1') # ランク 1

以下のすべてのテンソルは torch.cfloat 型です。

2 つのプロセスグループ、2 つのランクがあるものとします。

tensor = torch.tensor( ... ... [1 + 1j, 2 + 2j], dtype=torch.cfloat, device=device ... ... ) + 2 * rank * (1 + 1j) tensor tensor([1.+1.j, 2.+2.j], device='cuda:0') # ランク 0 tensor([3.+3.j, 4.+4.j], device='cuda:1') # ランク 1 dist.all_reduce(tensor, op=ReduceOp.SUM) tensor tensor([4.+4.j, 6.+6.j], device='cuda:0') # ランク 0 tensor([4.+4.j, 6.+6.j], device='cuda:1') # ランク 1

torch.distributed.reduce(tensor, dst=None, op=<RedOpType.SUM: 0>, group=None, async_op=False, group_dst=None)[source]# すべてのマシン間でテンソルデータを reduction します。最終結果を受け取るのは、dst ランクのプロセスだけです。

パラメーター

tensor (Tensor) – 集合通信の入力と出力です。この関数はインプレースで操作します。

dst (int) – グローバルプロセスグループ内の宛先ランクです（group 引数に関係なく）。

op (optional) – torch.distributed.ReduceOp 列挙型の値のいずれかです。要素ごとの reduction に使う操作を指定します。

group (ProcessGroup, optional) – 操作対象のプロセスグループです。None の場合、デフォルトのプロセスグループが使われます。

async_op (bool, optional) – この操作を非同期操作にするかどうかです。

group_dst (int) – group 内の宛先ランクです。group_dst と dst のどちらか一方を指定する必要があります。

戻り値

async_op が True の場合は非同期作業ハンドルを返します。async_op が False の場合、またはグループに属していない場合は None を返します。

torch.distributed.all_gather(tensor_list, tensor, group=None, async_op=False)[source]# グループ全体からテンソルを集め、リストに格納します。複素数テンソルとサイズが不揃いなテンソルをサポートしています。

パラメーター

tensor_list (list[Tensor]) – 出力リストです。集合通信の出力に使う、適切なサイズのテンソルを格納する必要があります。サイズが不揃いなテンソルもサポートされています。

tensor (Tensor) – 現在のプロセスからブロードキャストするテンソルです。

group (ProcessGroup, optional) – 操作対象のプロセスグループです。None の場合、デフォルトのプロセスグループが使われます。

async_op (bool, optional) – この操作を非同期操作にするかどうかです。

戻り値

async_op が True の場合は非同期作業ハンドルを返します。async_op が False の場合、またはグループに属していない場合は None を返します。

例

以下のすべてのテンソルは torch.int64 dtype です。

2 つのプロセスグループ、2 つのランクがあるものとします。

device = torch.device(f"cuda:{rank}") tensor_list = [ ... ... torch.zeros(2, dtype=torch.int64, device=device) for _ in range(2) ... ... ] tensor_list [tensor([0, 0], device='cuda:0'), tensor([0, 0], device='cuda:0')] # ランク 0 [tensor([0, 0], device='cuda:1'), tensor([0, 0], device='cuda:1')] # ランク 1 tensor = torch.arange(2, dtype=torch.int64, device=device) + 1 + 2 * rank tensor tensor([1, 2], device='cuda:0') # ランク 0 tensor([3, 4], device='cuda:1') # ランク 1 dist.all_gather(tensor_list, tensor) tensor_list [tensor([1, 2], device='cuda:0'), tensor([3, 4], device='cuda:0')] # ランク 0 [tensor([1, 2], device='cuda:1'), tensor([3, 4], device='cuda:1')] # ランク 1

以下のすべてのテンソルは torch.cfloat dtype です。

2 つのプロセスグループ、2 つのランクがあるものとします。

tensor_list = [ ... ... torch.zeros(2, dtype=torch.cfloat, device=device) for _ in range(2) ... ... ] tensor_list [tensor([0.+0.j, 0.+0.j], device='cuda:0'), tensor([0.+0.j, 0.+0.j], device='cuda:0')] # ランク 0 [tensor([0.+0.j, 0.+0.j], device='cuda:1'), tensor([0.+0.j, 0.+0.j], device='cuda:1')] # ランク 1 tensor = torch.tensor( ... ... [1 + 1j, 2 + 2j], dtype=torch.cfloat, device=device ... ... ) + 2 * rank * (1 + 1j) tensor tensor([1.+1.j, 2.+2.j], device='cuda:0') # ランク 0 tensor([3.+3.j, 4.+4.j], device='cuda:1') # ランク 1 dist.all_gather(tensor_list, tensor) tensor_list [tensor([1.+1.j, 2.+2.j], device='cuda:0'), tensor([3.+3.j, 4.+4.j], device='cuda:0')] # ランク 0 [tensor([1.+1.j, 2.+2.j], device='cuda:1'), tensor([3.+3.j, 4.+4.j], device='cuda:1')] # ランク 1

torch.distributed.all_gather_into_tensor(output_tensor, input_tensor, group=None, async_op=False)[source]# すべてのランクからテンソルを集め、単一の出力テンソルに格納します。この関数では、各プロセスのすべてのテンソルが同じサイズでなければなりません。

パラメーター

output_tensor (Tensor) – すべてのランクのテンソル要素を格納する出力テンソルです。次のいずれかの形式で適切なサイズにする必要があります。（i）すべての入力テンソルを主要次元に沿って連結したもの。「連結」の定義については torch.cat() を参照してください。（ii）すべての入力テンソルを主要次元に沿って積み重ねたもの。「積み重ね」の定義については torch.stack() を参照してください。次の例で、サポートされる出力形式をより具体的に説明します。

input_tensor (Tensor) – 現在のランクから集めるテンソルです。all_gather API とは異なり、この API ではすべてのランクの入力テンソルが同じサイズでなければなりません。

group (ProcessGroup, optional) – 操作対象のプロセスグループです。None の場合、デフォルトのプロセスグループが使われます。

async_op (bool, optional) – この操作を非同期操作にするかどうかです。

戻り値

async_op が True の場合は非同期作業ハンドルを返します。async_op が False の場合、またはグループに属していない場合は None を返します。

例

以下のすべてのテンソルは torch.int64 型で、CUDA デバイス上にあります。

2 つのランクがあるものとします。

device = torch.device(f"cuda:{rank}") tensor_in = torch.arange(2, dtype=torch.int64, device=device) + 1 + 2 * rank tensor_in tensor([1, 2], device='cuda:0') # ランク 0 tensor([3, 4], device='cuda:1') # ランク 1

連結形式の出力

tensor_out = torch.zeros(world_size * 2, dtype=torch.int64, device=device) dist.all_gather_into_tensor(tensor_out, tensor_in) tensor_out tensor([1, 2, 3, 4], device='cuda:0') # ランク 0 tensor([1, 2, 3, 4], device='cuda:1') # ランク 1

積み重ね形式の出力

tensor_out2 = torch.zeros(world_size, 2, dtype=torch.int64, device=device) dist.all_gather_into_tensor(tensor_out2, tensor_in) tensor_out2 tensor([[1, 2], [3, 4]], device='cuda:0') # ランク 0 tensor([[1, 2], [3, 4]], device='cuda:1') # ランク 1

torch.distributed.all_gather_object(object_list, obj, group=None)[source]# グループ全体から pickle 化可能なオブジェクトを集め、リストに格納します。all_gather() と同様ですが、Python オブジェクトを渡せます。オブジェクトは集約できるよう pickle 化可能でなければなりません。

パラメーター

object_list (list[Any]) – 出力リストです。この集合通信のグループサイズに合わせた適切な長さにし、出力を格納できるようにします。

obj (Any) – 現在のプロセスからブロードキャストする、pickle 化可能な Python オブジェクトです。

group (ProcessGroup, optional) – 操作対象のプロセスグループです。None の場合、デフォルトのプロセスグループが使われます。デフォルトは None です。

戻り値 None。呼び出し元のランクがグループに属している場合、集合通信の出力が入力 object_list に格納されます。呼び出し元のランクがグループに属していない場合、渡した object_list は変更されません。

注記 この API は all_gather() 集合通信と少し異なり、async_op ハンドルを提供しないため、ブロッキング呼び出しになります。

注記 NCCL ベースのプロセスグループでは、通信を行う前にオブジェクトの内部テンソル表現を GPU デバイスに移動する必要があります。この場合、使用するデバイスは torch.cuda.current_device() で指定されます。torch.cuda.set_device() を使って各ランクが個別の GPU を持つよう設定する責任は、ユーザーにあります。

警告 オブジェクト集合通信には、深刻なパフォーマンスおよびスケーラビリティ上の制約があります。詳細は Object collectives を参照してください。

警告 all_gather_object() は pickle モジュールを暗黙的に使用しますが、これは安全ではないことが知られています。悪意のある pickle データを作成し、unpickle 時に任意のコードを実行させることができます。信頼できるデータに対してのみ、この関数を呼び出してください。

警告 GPU テンソルを使って all_gather_object() を呼び出すことは、サポートが十分でなく、非効率です。テンソルが pickle 化されるため、GPU から CPU への転送が発生します。代わりに all_gather() の使用を検討してください。

例：::

注：各ランクでのプロセスグループ初期化は省略しています。

import torch.distributed as dist

world_size が 3 であることを前提とします。

gather_objects = ["foo", 12, {1: 2}] # pickle 化可能な任意のオブジェクト output = [None for _ in gather_objects] dist.all_gather_object(output, gather_objects[dist.get_rank()]) output ['foo', 12, {1: 2}]

torch.distributed.gather(tensor, gather_list=None, dst=None, group=None, async_op=False, group_dst=None)[source]# 単一のプロセスにテンソルのリストを集めます。この関数では、各プロセスのすべてのテンソルが同じサイズでなければなりません。

パラメーター

tensor (Tensor) – 入力テンソルです。

gather_list (list[Tensor], optional) – 集約データに使う、適切なサイズで同じサイズのテンソルのリストです（デフォルトは None。宛先ランクでは指定が必須です）。

dst (int, optional) – グローバルプロセスグループ内の宛先ランクです（group 引数に関係なく）。dst と group_dst の両方が None の場合、デフォルトはグローバルランク 0 です。

group (ProcessGroup, optional) – 操作対象のプロセスグループです。None の場合、デフォルトのプロセスグループが使われます。

async_op (bool, optional) – この操作を非同期操作にするかどうかです。

group_dst (int, optional) – group 内の宛先ランクです。dst と group_dst を同時に指定することはできません。

戻り値

async_op が True の場合は非同期作業ハンドルを返します。async_op が False の場合、またはグループに属していない場合は None を返します。

注記 gather_list 内のすべてのテンソルは同じサイズでなければなりません。

例：::

2 つのプロセスグループ、2 つのランクがあるものとします。

tensor_size = 2 device = torch.device(f'cuda:{rank}') tensor = torch.ones(tensor_size, device=device) + rank if dist.get_rank() == 0: gather_list = [torch.zeros_like(tensor, device=device) for i in range(2)] else: gather_list = None dist.gather(tensor, gather_list, dst=0)

ランク 0 が集約データを受け取ります。

gather_list [tensor([1., 1.], device='cuda:0'), tensor([2., 2.], device='cuda:0')] # ランク 0 None # ランク 1

torch.distributed.gather_object(obj, object_gather_list=None, dst=None, group=None, group_dst=None)[source]# グループ全体から pickle 化可能なオブジェクトを単一のプロセスに集めます。gather() と同様ですが、Python オブジェクトを渡せます。オブジェクトは集約できるよう pickle 化可能でなければなりません。

パラメーター

obj (Any) – 入力オブジェクトです。pickle 化可能でなければなりません。

object_gather_list (list[Any]) – 出力リストです。dst ランクでは、この集合通信のグループサイズに合わせた適切な長さにし、出力を格納できるようにします。dst 以外のランクでは None でなければなりません（デフォルトは None）。

dst (int, optional) – グローバルプロセスグループ内の宛先ランクです（group 引数に関係なく）。dst と group_dst の両方が None の場合、デフォルトはグローバルランク 0 です。

group (Optional[ProcessGroup]) – 操作対象のプロセスグループです。None の場合、デフォルトのプロセスグループが使われます。デフォルトは None です。

group_dst (int, optional) – group 内の宛先ランクです。dst と group_dst を同時に指定することはできません。

戻り値 None。dst ランクでは、集合通信の出力が object_gather_list に格納されます。

注記 この API は gather 集合通信と少し異なり、async_op ハンドルを提供しないため、ブロッキング呼び出しになります。

注記 NCCL ベースのプロセスグループでは、通信を行う前にオブジェクトの内部テンソル表現を GPU デバイスに移動する必要があります。この場合、使用するデバイスは torch.cuda.current_device() で指定されます。torch.cuda.set_device() を使って各ランクが個別の GPU を持つよう設定する責任は、ユーザーにあります。

警告 オブジェクト集合通信には、深刻なパフォーマンスおよびスケーラビリティ上の制約があります。詳細は Object collectives を参照してください。

警告 gather_object() は pickle モジュールを暗黙的に使用しますが、これは安全ではないことが知られています。悪意のある pickle データを作成し、unpickle 時に任意のコードを実行させることができます。信頼できるデータに対してのみ、この関数を呼び出してください。

警告 GPU テンソルを使って gather_object() を呼び出すことは、サポートが十分でなく、非効率です。テンソルが pickle 化されるため、GPU から CPU への転送が発生します。代わりに gather() の使用を検討してください。

例：::

注：各ランクでのプロセスグループ初期化は省略しています。

import torch.distributed as dist

world_size が 3 であることを前提とします。

gather_objects = ["foo", 12, {1: 2}] # pickle 化可能な任意のオブジェクト output = [None for _ in gather_objects] dist.gather_object( ... ... gather_objects[dist.get_rank()], ... ... output if dist.get_rank() == 0 else None, ... ... dst=0 ... ... )

ランク 0 上

output ['foo', 12, {1: 2}]

torch.distributed.scatter(tensor, scatter_list=None, src=None, group=None, async_op=False, group_src=None)[source]# テンソルのリストをグループ内のすべてのプロセスに分散します。各プロセスはテンソルを 1 つだけ受け取り、そのデータを tensor 引数に格納します。複素数テンソルもサポートされています。

パラメーター

tensor (Tensor) – 出力テンソルです。

scatter_list (list[Tensor]) – 分散するテンソルのリストです（デフォルトは None。送信元ランクでは指定が必須です）。

src (int) – グローバルプロセスグループ内の送信元ランクです（group 引数に関係なく）。src と group_src の両方が None の場合、デフォルトはグローバルランク 0 です。

group (ProcessGroup, optional) – 操作対象のプロセスグループです。None の場合、デフォルトのプロセスグループが使われます。

async_op (bool, optional) – この操作を非同期操作にするかどうかです。

group_src (int, optional) – group 内の送信元ランクです。src と group_src を同時に指定することはできません。

戻り値

async_op が True の場合は非同期作業ハンドルを返します。async_op が False の場合、またはグループに属していない場合は None を返します。

注記 scatter_list 内のすべてのテンソルは同じサイズでなければなりません。

例：::

注：各ランクでのプロセスグループ初期化は省略しています。

import torch.distributed as dist tensor_size = 2 device = torch.device(f'cuda:{rank}') output_tensor = torch.zeros(tensor_size, device=device) if dist.get_rank() == 0:

world_size が 2 であることを前提とします。

テンソルのみを使い、すべて同じサイズでなければなりません。

t_ones = torch.ones(tensor_size, device=device) t_fives = torch.ones(tensor_size, device=device) * 5 scatter_list = [t_ones, t_fives] else: scatter_list = None dist.scatter(output_tensor, scatter_list, src=0)

ランク i は scatter_list[i] を受け取ります。

output_tensor tensor([1., 1.], device='cuda:0') # ランク 0 tensor([5., 5.], device='cuda:1') # ランク 1

torch.distributed.scatter_object_list(scatter_object_output_list, scatter_object_input_list=None, src=None, group=None, group_src=None)[source]# scatter_object_input_list 内の pickle 化可能なオブジェクトをグループ全体に分散します。scatter() と同様ですが、Python オブジェクトを渡せます。各ランクでは、分散されたオブジェクトが scatter_object_output_list の最初の要素に格納されます。分散するには、scatter_object_input_list 内のすべてのオブジェクトが pickle 化可能でなければなりません。

パラメーター

scatter_object_output_list (List[Any]) – このランクに分散されたオブジェクトを最初の要素に格納する、空でないリストです。

scatter_object_input_list (List[Any], optional) – 分散する入力オブジェクトのリストです。各オブジェクトは pickle 化可能である必要があります。分散されるのは src ランク上のオブジェクトのみであり、src 以外のランクでは引数を None にできます。

src (int) – scatter_object_input_list を分散する送信元ランクです。送信元ランクはグローバルプロセスグループに基づきます（group 引数に関係なく）。src と group_src の両方が None の場合、デフォルトはグローバルランク 0 です。

group (Optional[ProcessGroup]) – 操作対象のプロセスグループです。None の場合、デフォルトのプロセスグループが使われます。デフォルトは None です。

group_src (int, optional) – group 内の送信元ランクです。src と group_src を同時に指定することはできません。

戻り値 None。ランクがグループに属している場合、scatter_object_output_list の最初の要素に、そのランクへ分散されたオブジェクトが設定されます。

注記 この API は scatter 集合通信と少し異なり、async_op ハンドルを提供しないため、ブロッキング呼び出しになります。

警告 オブジェクト集合通信には、深刻なパフォーマンスおよびスケーラビリティ上の制約があります。詳細は Object collectives を参照してください。

警告 scatter_object_list() は pickle モジュールを暗黙的に使用しますが、これは安全ではないことが知られています。悪意のある pickle データを作成し、unpickle 時に任意のコードを実行させることができます。信頼できるデータに対してのみ、この関数を呼び出してください。

警告 GPU テンソルを使って scatter_object_list() を呼び出すことは、サポートが十分でなく、非効率です。テンソルが pickle 化されるため、GPU から CPU への転送が発生します。代わりに scatter() の使用を検討してください。

例：::

注：各ランクでのプロセスグループ初期化は省略しています。

import torch.distributed as dist if dist.get_rank() == 0:

world_size が 3 であることを前提とします。

objects = ["foo", 12, {1: 2}] # pickle 化可能な任意のオブジェクト else:

src 以外のランクでは任意のリストを指定できます。要素は使われません。

objects = [None, None, None] output_list = [None] dist.scatter_object_list(output_list, objects, src=0)

ランク i は objects[i] を受け取ります。例：ランク 2 の場合

output_list [{1: 2}]

torch.distributed.reduce_scatter(output, input_list, op=<RedOpType.SUM: 0>, group=None, async_op=False)[source]# テンソルのリストを reduction してから、グループ内のすべてのプロセスに分散します。

パラメーター

output (Tensor) – 出力テンソルです。

input_list (list[Tensor]) – reduction して分散するテンソルのリストです。

op (optional) – torch.distributed.ReduceOp 列挙型の値のいずれかです。要素ごとの reduction に使う操作を指定します。

group (ProcessGroup, optional) – 操作対象のプロセスグループです。None の場合、デフォルトのプロセスグループが使われます。

async_op (bool, optional) – この操作を非同期操作にするかどうかです。

戻り値

async_op が True の場合は非同期作業ハンドルを返します。async_op が False の場合、またはグループに属していない場合は None を返します。

torch.distributed.reduce_scatter_tensor(output, input, op=<RedOpType.SUM: 0>, group=None, async_op=False)[source]# テンソルを reduction してから、グループ内のすべてのランクに分散します。

パラメーター

output (Tensor) – 出力テンソルです。すべてのランクで同じサイズでなければなりません。

input (Tensor) – reduction して分散する入力テンソルです。サイズは出力テンソルのサイズに world size を掛けた値である必要があります。入力テンソルは次のいずれかの形状をとれます。（i）出力テンソルを主要次元に沿って連結したもの。（ii）出力テンソルを主要次元に沿って積み重ねたもの。「連結」の定義については torch.cat() を、「積み重ね」の定義については torch.stack() を参照してください。

group (ProcessGroup, optional) – 操作対象のプロセスグループです。None の場合、デフォルトのプロセスグループが使われます。

async_op (bool, optional) – この操作を非同期操作にするかどうかです。

戻り値

async_op が True の場合は非同期作業ハンドルを返します。async_op が False の場合、またはグループに属していない場合は None を返します。

例

以下のすべてのテンソルは torch.int64 型で、CUDA デバイス上にあります。

2 つのランクがあるものとします。

device = torch.device(f"cuda:{rank}") tensor_out = torch.zeros(2, dtype=torch.int64, device=device)

連結形式の入力

tensor_in = torch.arange(world_size * 2, dtype=torch.int64, device=device) tensor_in tensor([0, 1, 2, 3], device='cuda:0') # ランク 0 tensor([0, 1, 2, 3], device='cuda:1') # ランク 1 dist.reduce_scatter_tensor(tensor_out, tensor_in) tensor_out tensor([0, 2], device='cuda:0') # ランク 0 tensor([4, 6], device='cuda:1') # ランク 1

積み重ね形式の入力

tensor_in = torch.reshape(tensor_in, (world_size, 2)) tensor_in tensor([[0, 1], [2, 3]], device='cuda:0') # ランク 0 tensor([[0, 1], [2, 3]], device='cuda:1') # ランク 1 dist.reduce_scatter_tensor(tensor_out, tensor_in) tensor_out tensor([0, 2], device='cuda:0') # ランク 0 tensor([4, 6], device='cuda:1') # ランク 1

torch.distributed.all_to_all_single(output, input, output_split_sizes=None, input_split_sizes=None, group=None, async_op=False)[source]# 入力テンソルを分割し、その分割リストをグループ内のすべてのプロセスに分散します。その後、グループ内のすべてのプロセスから受け取ったテンソルを連結し、単一の出力テンソルとして返します。複素数テンソルもサポートされています。

パラメーター

output (Tensor) – 連結された集約済み出力テンソルです。

input (Tensor) – 分散する入力テンソルです。

output_split_sizes – (list[Int], optional)：指定した場合の dim 0 における出力分割サイズです。None または空の場合、出力テンソルの dim 0 は world_size で均等に割り切れる必要があります。

input_split_sizes – (list[Int], optional)：指定した場合の dim 0 における入力分割サイズです。None または空の場合、入力テンソルの dim 0 は world_size で均等に割り切れる必要があります。

group (ProcessGroup, optional) – 操作対象のプロセスグループです。None の場合、デフォルトのプロセスグループが使われます。

async_op (bool, optional) – この操作を非同期操作にするかどうかです。

戻り値

async_op が True の場合は非同期作業ハンドルを返します。async_op が False の場合、またはグループに属していない場合は None を返します。

警告 all_to_all_single は実験的であり、今後変更される可能性があります。

例

input = torch.arange(4) + rank * 4 input tensor([0, 1, 2, 3]) # ランク 0 tensor([4, 5, 6, 7]) # ランク 1 tensor([8, 9, 10, 11]) # ランク 2 tensor([12, 13, 14, 15]) # ランク 3 output = torch.empty([4], dtype=torch.int64) dist.all_to_all_single(output, input) output tensor([0, 4, 8, 12]) # ランク 0 tensor([1, 5, 9, 13]) # ランク 1 tensor([2, 6, 10, 14]) # ランク 2 tensor([3, 7, 11, 15]) # ランク 3

基本的には次の操作と同じです。

scatter_list = list(input.chunk(world_size)) gather_list = list(output.chunk(world_size)) for i in range(world_size): dist.scatter(gather_list[i], scatter_list if i == rank else [], src = i)

不均等分割の別の例

input tensor([0, 1, 2, 3, 4, 5]) # ランク 0 tensor([10, 11, 12, 13, 14, 15, 16, 17, 18]) # ランク 1 tensor([20, 21, 22, 23, 24]) # ランク 2 tensor([30, 31, 32, 33, 34, 35, 36]) # ランク 3 input_splits [2, 2, 1, 1] # ランク 0 [3, 2, 2, 2] # ランク 1 [2, 1, 1, 1] # ランク 2 [2, 2, 2, 1] # ランク 3 output_splits [2, 3, 2, 2] # ランク 0 [2, 2, 1, 2] # ランク 1 [1, 2, 1, 2] # ランク 2 [1, 2, 1, 1] # ランク 3 output = ... dist.all_to_all_single(output, input, output_splits, input_splits) output tensor([ 0, 1, 10, 11, 12, 20, 21, 30, 31]) # ランク 0 tensor([ 2, 3, 13, 14, 22, 32, 33]) # ランク 1 tensor([ 4, 15, 16, 23, 34, 35]) # ランク 2 tensor([ 5, 17, 18, 24, 36]) # ランク 3

torch.cfloat 型のテンソルを使う別の例。

input = torch.tensor( ... ... [1 + 1j, 2 + 2j, 3 + 3j, 4 + 4j], dtype=torch.cfloat ... ... ) + 4 * rank * (1 + 1j) input tensor([1+1j, 2+2j, 3+3j, 4+4j]) # ランク 0 tensor([5+5j, 6+6j, 7+7j, 8+8j]) # ランク 1 tensor([9+9j, 10+10j, 11+11j, 12+12j]) # ランク 2 tensor([13+13j, 14+14j, 15+15j, 16+16j]) # ランク 3 output = torch.empty([4], dtype=torch.int64) dist.all_to_all_single(output, input) output tensor([1+1j, 5+5j, 9+9j, 13+13j]) # ランク 0 tensor([2+2j, 6+6j, 10+10j, 14+14j]) # ランク 1 tensor([3+3j, 7+7j, 11+11j, 15+15j]) # ランク 2 tensor([4+4j, 8+8j, 12+12j, 16+16j]) # ランク 3

torch.distributed.all_to_all(output_tensor_list, input_tensor_list, group=None, async_op=False)[source]# 入力テンソルのリストをグループ内のすべてのプロセスに分散し、出力リストに集約済みテンソルのリストを返します。複素数テンソルもサポートされています。

パラメーター

output_tensor_list (list[Tensor]) – ランクごとに 1 つずつ集約するテンソルのリストです。

input_tensor_list (list[Tensor]) – ランクごとに 1 つずつ分散するテンソルのリストです。

group (ProcessGroup, optional) – 操作対象のプロセスグループです。None の場合、デフォルトのプロセスグループが使われます。

async_op (bool, optional) – この操作を非同期操作にするかどうかです。

戻り値

async_op が True の場合は非同期作業ハンドルを返します。async_op が False の場合、またはグループに属していない場合は None を返します。

警告 all_to_all は実験的であり、今後変更される可能性があります。

例

input = torch.arange(4) + rank * 4 input = list(input.chunk(4)) input [tensor([0]), tensor([1]), tensor([2]), tensor([3])] # ランク 0 [tensor([4]), tensor([5]), tensor([6]), tensor([7])] # ランク 1 [tensor([8]), tensor([9]), tensor([10]), tensor([11])] # ランク 2 [tensor([12]), tensor([13]), tensor([14]), tensor([15])] # ランク 3 output = list(torch.empty([4], dtype=torch.int64).chunk(4)) dist.all_to_all(output, input) output [tensor([0]), tensor([4]), tensor([8]), tensor([12])] # ランク 0 [tensor([1]), tensor([5]), tensor([9]), tensor([13])] # ランク 1 [tensor([2]), tensor([6]), tensor([10]), tensor([14])] # ランク 2 [tensor([3]), tensor([7]), tensor([11]), tensor([15])] # ランク 3

基本的には次の操作と同じです。

scatter_list = input gather_list = output for i in range(world_size): dist.scatter(gather_list[i], scatter_list if i == rank else [], src=i) input tensor([0, 1, 2, 3, 4, 5]) # ランク 0 tensor([10, 11, 12, 13, 14, 15, 16, 17, 18]) # ランク 1 tensor([20, 21, 22, 23, 24]) # ランク 2 tensor([30, 31, 32, 33, 34, 35, 36]) # ランク 3 input_splits [2, 2, 1, 1] # ランク 0 [3, 2, 2, 2] # ランク 1 [2, 1, 1, 1] # ランク 2 [2, 2, 2, 1] # ランク 3 output_splits [2, 3, 2, 2] # ランク 0 [2, 2, 1, 2] # ランク 1 [1, 2, 1, 2] # ランク 2 [1, 2, 1, 1] # ランク 3 input = list(input.split(input_splits)) input [tensor([0, 1]), tensor([2, 3]), tensor([4]), tensor([5])] # ランク 0 [tensor([10, 11, 12]), tensor([13, 14]), tensor([15, 16]), tensor([17, 18])] # ランク 1 [tensor([20, 21]), tensor([22]), tensor([23]), tensor([24])] # ランク 2 [tensor([30, 31]), tensor([32, 33]), tensor([34, 35]), tensor([36])] # ランク 3 output = ... dist.all_to_all(output, input) output [tensor([0, 1]), tensor([10, 11, 12]), tensor([20, 21]), tensor([30, 31])] # ランク 0 [tensor([2, 3]), tensor([13, 14]), tensor([22]), tensor([32, 33])] # ランク 1 [tensor([4]), tensor([15, 16]), tensor([23]), tensor([34, 35])] # ランク 2 [tensor([5]), tensor([17, 18]), tensor([24]), tensor([36])] # ランク 3

torch.cfloat 型のテンソルを使う別の例。

input = torch.tensor( ... ... [1 + 1j, 2 + 2j, 3 + 3j, 4 + 4j], dtype=torch.cfloat ... ... ) + 4 * rank * (1 + 1j) input = list(input.chunk(4)) input [tensor([1+1j]), tensor([2+2j]), tensor([3+3j]), tensor([4+4j])] # ランク 0 [tensor([5+5j]), tensor([6+6j]), tensor([7+7j]), tensor([8+8j])] # ランク 1 [tensor([9+9j]), tensor([10+10j]), tensor([11+11j]), tensor([12+12j])] # ランク 2 [tensor([13+13j]), tensor([14+14j]), tensor([15+15j]), tensor([16+16j])] # ランク 3 output = list(torch.empty([4], dtype=torch.int64).chunk(4)) dist.all_to_all(output, input) output [tensor([1+1j]), tensor([5+5j]), tensor([9+9j]), tensor([13+13j])] # ランク 0 [tensor([2+2j]), tensor([6+6j]), tensor([10+10j]), tensor([14+14j])] # ランク 1 [tensor([3+3j]), tensor([7+7j]), tensor([11+11j]), tensor([15+15j])] # ランク 2 [tensor([4+4j]), tensor([8+8j]), tensor([12+12j]), tensor([16+16j])] # ランク 3

torch.distributed.barrier(group=None, async_op=False, device_ids=None)[source]# すべてのプロセスを同期します。async_op が False の場合、この集合通信はグループ全体がこの関数に入るまでプロセスをブロックします。非同期操作の場合は、非同期作業ハンドルに対して wait() を呼び出すまでプロセスをブロックしません。

パラメーター

group (ProcessGroup, optional) – 操作対象のプロセスグループです。None の場合、デフォルトのプロセスグループが使われます。

async_op (bool, optional) – この操作を非同期操作にするかどうかです。

device_ids ([int], optional) – デバイス／GPU ID のリストです。ID は 1 つだけ指定することが想定されています。

戻り値

async_op が True の場合は非同期作業ハンドルを返します。async_op が False の場合、またはグループに属していない場合は None を返します。

注記 ProcessGroupNCCL では、barrier 集合通信が完了するまで CPU スレッドをブロックするようになりました。

注記 ProcessGroupNCCL は、1 要素テンソルに対する all_reduce として barrier を実装します。このテンソルを割り当てるデバイスを選択する必要があります。デバイスは次の順序で決定されます。（1）None でない場合、barrier の device_ids 引数で渡された最初のデバイス。（2）None でない場合、init_process_group に渡されたデバイス。（3）テンソル入力を受け取る別の集合通信が実行済みの場合、このプロセスグループで最初に使われたデバイス。（4）グローバルランクをローカルデバイス数で割った余りに対応するデバイスインデックス。

torch.distributed.monitored_barrier(group=None, timeout=None, wait_all_ranks=False)[source]# torch.distributed.barrier と同様にプロセスを同期しますが、設定可能なタイムアウトをサポートします。指定されたタイムアウト内にこのバリアを通過しなかったランクを報告できます。具体的には、ゼロ以外のランクは rank 0 からの send／recv が処理されるまでブロックします。rank 0 は他のランクからの send／recv がすべて処理されるまでブロックし、時間内に応答しなかったランクの失敗を報告します。あるランクが monitored_barrier に到達しない場合（たとえばハングした場合）、他のすべてのランクは monitored_barrier で失敗します。この集合通信は、グループ全体が関数を正常に終了するまで、グループ内のすべてのプロセス／ランクをブロックします。デバッグや同期に役立ちますが、パフォーマンスに影響する可能性があるため、デバッグやホスト側で完全な同期が必要な状況に限って使ってください。デバッグのために、アプリケーションの集合通信呼び出しの前にこのバリアを挿入すると、ランク間の非同期状態を確認できます。

注記 この集合通信は GLOO バックエンドでのみサポートされています。

パラメーター

group (ProcessGroup, optional) – 操作対象のプロセスグループです。None の場合、デフォルトのプロセスグループが使われます。

timeout (datetime.timedelta, optional) – monitored_barrier のタイムアウトです。None の場合、デフォルトのプロセスグループのタイムアウトが使われます。

wait_all_ranks (bool, optional) – 失敗したランクをすべて収集するかどうかです。デフォルトは False で、rank 0 の monitored_barrier は最初に検出した失敗ランクで例外を発生させ、速やかに失敗します。wait_all_ranks=True に設定すると、monitored_barrier は失敗したランクをすべて収集し、それらすべての情報を含むエラーを発生させます。

戻り値 None。

例：::

注：各ランクでのプロセスグループ初期化は省略しています。

import torch.distributed as dist if dist.get_rank() != 1: dist.monitored_barrier() # 次の情報を示す例外を発生させます。

ランク 1 が monitored_barrier を呼び出さなかったこと。

wait_all_ranks=True の例

if dist.get_rank() == 0: dist.monitored_barrier(wait_all_ranks=True) # 例外を発生させます。

ランク 1、2、... world_size - 1 が monitored_barrier を呼び出さなかったことを示します。

class torch.distributed.Work# Work オブジェクトは、PyTorch の分散パッケージにおける保留中の非同期操作のハンドルを表します。dist.all_reduce(tensor, async_op=True) など、非ブロッキング集合通信操作によって返されます。

block_current_stream(self: torch._C._distributed_c10d.Work) → None# 現在アクティブな GPU ストリームを、操作が完了するまでブロックします。GPU ベースの集合通信では、これは synchronize と同等です。Gloo を使った CPU 起点の集合通信などでは、操作が完了するまで CUDA ストリームをブロックします。いずれの場合も、このメソッドはただちに戻ります。操作が成功したかどうかを確認するには、Work オブジェクトの結果を非同期に確認してください。

boxed(self: torch._C._distributed_c10d.Work) → object#

exception(self: torch._C._distributed_c10d.Work) → std::__exception_ptr::exception_ptr#

get_future(self: torch._C._distributed_c10d.Work) → torch.Future# Work の完了に関連付けられた torch.futures.Future オブジェクトを返します。たとえば、fut = process_group.allreduce(tensors).get_future() で future オブジェクトを取得できます。

例：:: 次は、get_future API を使って allreduce の完了に関連付けられた Future を取得する、単純な allreduce DDP 通信フックの例です。

def allreduce(process_group: dist.ProcessGroup, bucket: dist.GradBucket): -> torch.futures.Future group_to_use = process_group if process_group is not None else torch.distributed.group.WORLD tensor = bucket.buffer().div_(group_to_use.size()) return torch.distributed.all_reduce(tensor, group=group_to_use, async_op=True).get_future() ddp_model.register_comm_hook(state=None, hook=allreduce)

警告 get_future API は NCCL と、一部の GLOO および MPI バックエンドをサポートします（send／recv などのピアツーピア操作はサポートしません）。torch.futures.Future を返します。上の例では、NCCL バックエンドを使って GPU 上で allreduce が実行されます。fut.wait() は、適切な NCCL ストリームと PyTorch の現在のデバイスストリームを同期した後に戻ります。これにより非同期 CUDA 実行が可能になり、GPU 上の操作全体の完了を待つ必要はありません。CUDAFuture は TORCH_NCCL_BLOCKING_WAIT フラグや NCCL の barrier() をサポートしない点に注意してください。また、fut.then() でコールバック関数を追加した場合、WorkNCCL の NCCL ストリームが ProcessGroupNCCL 専用のコールバックスレッドと同期するまで待機し、コールバックスレッドでコールバックを実行した後、インラインでコールバックを呼び出します。fut.then() は、コールバックの戻り値とコールバックスレッドを記録した CUDAEvent を保持する別の CUDAFuture を返します。CPU の Work では、処理が完了して value() のテンソルが準備できると fut.done() が true を返します。GPU の Work では、操作がキューイングされた場合にのみ fut.done() が true を返します。CPU-GPU 混在の Work（例：GLOO で GPU テンソルを送信する場合）では、テンソルが各ノードに到着した時点で fut.done() が true を返しますが、各 GPU で同期済みとは限りません（GPU Work と同様です）。

get_future_result(self: torch._C._distributed_c10d.Work) → torch.Future# WorkResult 列挙型に対応する int 型の torch.futures.Future オブジェクトを返します。たとえば、fut = process_group.allreduce(tensor).get_future_result() で future オブジェクトを取得できます。

例：:: ユーザーは fut.wait() を使って Work の完了をブロッキングで待ち、fut.value() で WorkResult を取得できます。また、fut.then(call_back_func) を使うと、現在のスレッドをブロックせずに、Work の完了時に呼び出すコールバック関数を登録できます。

警告 get_future_result API は NCCL をサポートします。

is_completed(self: torch._C._distributed_c10d.Work) → bool#

is_success(self: torch._C._distributed_c10d.Work) → bool#

result(self: torch._C._distributed_c10d.Work) → list[torch.Tensor]#

source_rank(self: torch._C._distributed_c10d.Work) → int#

synchronize(self: torch._C._distributed_c10d.Work) → None#

static unbox(arg0: object) → torch._C._distributed_c10d.Work#

wait(self: torch._C._distributed_c10d.Work, timeout: datetime.timedelta = datetime.timedelta(0)) → bool# true／false を返します。

例：:: try:work.wait(timeout) except:# 対応処理

警告 通常、ユーザーがタイムアウトを設定する必要はありません。wait() の呼び出しは synchronize() の呼び出しと同じです。つまり、現在のストリームを NCCL Work の完了までブロックします。ただし、timeout が設定されている場合は、NCCL Work が完了するかタイムアウトになるまで CPU スレッドがブロックされます。タイムアウトになると例外が発生します。

class torch.distributed.ReduceOp# 利用可能な reduction 操作を表す列挙型に似たクラスです。SUM、PRODUCT、MIN、MAX、BAND、BOR、BXOR、PREMUL_SUM が利用できます。NCCL バックエンドでは、BAND、BOR、BXOR の reduction は利用できません。AVG はランク間で合計する前に、値を world size で割ります。AVG は NCCL バックエンドでのみ利用でき、NCCL のバージョン 2.10 以降が必要です。PREMUL_SUM は reduction の前に入力を指定されたスカラーでローカルに乗算します。PREMUL_SUM は NCCL バックエンドでのみ利用でき、NCCL のバージョン 2.11 以降が必要です。ユーザーは torch.distributed._make_nccl_premul_sum を使う必要があります。また、MAX、MIN、PRODUCT は複素数テンソルではサポートされていません。このクラスの値は ReduceOp.SUM のように属性として取得できます。reduce() など、reduction 集合通信の方式を指定する際に使います。このクラスは members プロパティをサポートしません。

class torch.distributed.reduce_op# reduction 操作（SUM、PRODUCT、MIN、MAX）を表す非推奨の列挙型に似たクラスです。代わりに ReduceOp の使用を推奨します。

分散キー・バリューストア#

分散パッケージには、グループ内のプロセス間で情報を共有するため、および torch.distributed.init_process_group() で分散パッケージを初期化するための分散キー・バリューストアが含まれています（init_method を指定する代わりに、ストアを明示的に作成できます）。キー・バリューストアには TCPStore、FileStore、HashStore の 3 種類があります。

class torch.distributed.Store# PyTorch の分散パッケージが提供する 3 つの実装（TCPStore、FileStore、HashStore）など、すべてのストア実装の基底クラスです。

init(self: torch._C._distributed_c10d.Store) → None#

add(self: torch._C._distributed_c10d.Store, arg0: str, arg1: SupportsInt) → int# 指定したキーに対する add の最初の呼び出しは、amount を初期値として、そのキーに関連付けられたカウンターをストア内に作成します。同じキーに対する後続の add 呼び出しは、指定された amount だけカウンターを増やします。set() によってすでにストア内に設定されたキーに対して add() を呼び出すと、例外が発生します。

パラメーター

key (str) – カウンターを増やすストア内のキーです。

amount (int) – カウンターを増やす量です。

例：::

import torch.distributed as dist from datetime import timedelta

例として TCPStore を使用します。他のストア型も使えます。

store = dist.TCPStore("127.0.0.1", 0, 1, True, timedelta(seconds=30)) store.add("first_key", 1) store.add("first_key", 6)

7 を返すはずです。

store.get("first_key")

append(self: torch._C._distributed_c10d.Store, arg0: str, arg1: str) → None# 指定したキーと値のペアをストアに追加します。キーがストアに存在しない場合は作成されます。

パラメーター

key (str) – ストアに追加するキーです。

value (str) – ストアに追加する、key に関連付けられた値です。

例：::

import torch.distributed as dist from datetime import timedelta store = dist.TCPStore("127.0.0.1", 0, 1, True, timedelta(seconds=30)) store.append("first_key", "po") store.append("first_key", "tato")

"potato" を返すはずです。

store.get("first_key")

check(self: torch._C._distributed_c10d.Store, arg0: collections.abc.Sequence[str]) → bool# 指定したキーのリストに値がストアされているかどうかを確認します。通常、この呼び出しはただちに戻りますが、TCPStore が破棄された後に check を呼び出す場合など、まれなデッドロックが発生する場合があります。check() を呼び出すと、指定したキーがストアに保存されているかどうかを確認します。

パラメーター

keys (list[str]) – ストアに保存されているかを調べるキーです。

例：::

import torch.distributed as dist from datetime import timedelta

例として TCPStore を使用します。他のストア型も使えます。

store = dist.TCPStore("127.0.0.1", 0, 1, True, timedelta(seconds=30)) store.add("first_key", 1)

7 を返すはずです。

store.check(["first_key"])

clone(self: torch._C._distributed_c10d.Store) → torch._C._distributed_c10d.Store# ストアを複製し、同じ基盤ストアを参照する新しいオブジェクトを返します。返されたストアは元のオブジェクトと並行して使えます。これは、スレッドごとにストアを複製することで、複数スレッドから安全にストアを使えるようにすることを目的としています。

compare_set(self: torch._C._distributed_c10d.Store, arg0: str, arg1: str, arg2: str) → bytes# 指定されたキーと値のペアをストアに挿入します。挿入前に expected_value と desired_value を比較します。キーの expected_value がストアに存在する場合、または expected_value が空文字列の場合にのみ、desired_value が設定されます。

パラメーター

key (str) – ストアで確認するキーです。

expected_value (str) – 挿入前に確認する、key に関連付けられた値です。

desired_value (str) – ストアに追加する、key に関連付けられた値です。

例：::

import torch.distributed as dist from datetime import timedelta store = dist.TCPStore("127.0.0.1", 0, 1, True, timedelta(seconds=30)) store.set("key", "first_value") store.compare_set("key", "first_value", "second_value")

"second_value" を返すはずです。

store.get("key")

delete_key(self: torch._C._distributed_c10d.Store, arg0: str) → bool# ストアから、指定したキーに関連付けられたキー・値ペアを削除します。キーが正常に削除された場合は true、削除されなかった場合は false を返します。

警告 delete_key API は TCPStore と HashStore でのみサポートされています。FileStore でこの API を使うと例外が発生します。

パラメーター

key (str) – ストアから削除するキーです。

戻り値

key が削除された場合は True、それ以外の場合は False です。

例：::

import torch.distributed as dist from datetime import timedelta

例として TCPStore を使用します。HashStore も使えます。

store = dist.TCPStore("127.0.0.1", 0, 1, True, timedelta(seconds=30)) store.set("first_key")

true を返すはずです。

store.delete_key("first_key")

false を返すはずです。

store.delete_key("bad_key")

get(self: torch._C._distributed_c10d.Store, arg0: str) → bytes# ストア内の指定されたキーに関連付けられた値を取得します。キーがストアに存在しない場合、ストアの初期化時に定義されたタイムアウトまで待機してから、例外を発生させます。

パラメーター

key (str) – このキーに関連付けられた値を返します。

戻り値

key がストア内にある場合は、key に関連付けられた値です。

例：::

import torch.distributed as dist from datetime import timedelta store = dist.TCPStore("127.0.0.1", 0, 1, True, timedelta(seconds=30)) store.set("first_key", "first_value")

"first_value" を返すはずです。

store.get("first_key")

has_extended_api(self: torch._C._distributed_c10d.Store) → bool# ストアが拡張操作をサポートする場合は true を返します。

multi_get(self: torch._C._distributed_c10d.Store, arg0: collections.abc.Sequence[str]) → list[bytes]# keys 内のすべての値を取得します。keys 内のいずれかのキーがストアに存在しない場合、タイムアウトまで待機します。

パラメーター

keys (List[str]) – ストアから取得するキーです。

例：::

import torch.distributed as dist from datetime import timedelta store = dist.TCPStore("127.0.0.1", 0, 1, True, timedelta(seconds=30)) store.set("first_key", "po") store.set("second_key", "tato")

[b"po", b"tato"] を返すはずです。

store.multi_get(["first_key", "second_key"])

multi_set(self: torch._C._distributed_c10d.Store, arg0: collections.abc.Sequence[str], arg1: collections.abc.Sequence[str]) → None# 指定されたキーと値に基づいて、キー・値ペアのリストをストアに挿入します。

パラメーター

keys (List[str]) – 挿入するキーです。

values (List[str]) – 挿入する値です。

例：::

import torch.distributed as dist from datetime import timedelta store = dist.TCPStore("127.0.0.1", 0, 1, True, timedelta(seconds=30)) store.multi_set(["first_key", "second_key"], ["po", "tato"])

b"po" を返すはずです。

store.get("first_key")

num_keys(self: torch._C._distributed_c10d.Store) → int# ストアに設定されたキーの数を返します。すべてのワーカーを調整するために 1 つのキーが使われるため、この数は通常、set() と add() で追加されたキーの数より 1 つ多くなります。

警告 TCPStore で使う場合、num_keys は基盤となるファイルに書き込まれたキーの数を返します。ストアを破棄して同じファイルで別のストアを作成すると、元のキーは保持されます。

戻り値

ストア内にあるキーの数です。

例：::

import torch.distributed as dist from datetime import timedelta

例として TCPStore を使用します。他のストア型も使えます。

store = dist.TCPStore("127.0.0.1", 0, 1, True, timedelta(seconds=30)) store.set("first_key", "first_value")

2 を返すはずです。

store.num_keys()

queue_len(self: torch._C._distributed_c10d.Store, arg0: str) → int# 指定したキューの長さを返します。キューが存在しない場合は 0 を返します。詳細は queue_push を参照してください。

パラメーター

key (str) – 長さを取得するキューのキーです。

queue_pop(self: torch._C._distributed_c10d.Store, key: str, block: bool = True) → bytes# 指定したキューから値を取り出します。キューが空の場合、タイムアウトまで待機します。詳細は queue_push を参照してください。block が False の場合、キューが空だと dist.QueueEmptyError が発生します。

パラメーター

key (str) – 値を取り出すキューのキーです。

block (bool) – キーを待ってブロックするか、ただちに戻るかです。

queue_push(self: torch._C._distributed_c10d.Store, arg0: str, arg1: str) → None# 指定したキューに値を追加します。キューと set／get 操作に同じキーを使うと、予期しない動作が発生する可能性があります。キューでは wait／check 操作がサポートされています。キューで wait を使うと、待機中のワーカー全員ではなく 1 つだけを起こします。

パラメーター

key (str) – 値を追加するキューのキーです。

value (str) – キューに追加する値です。

set(self: torch._C._distributed_c10d.Store, arg0: str, arg1: str) → None# 指定されたキーと値に基づいて、キー・値ペアをストアに挿入します。キーがすでにストアにある場合は、以前の値を新しい値で上書きします。

パラメーター

key (str) – ストアに追加するキーです。

value (str) – ストアに追加する、key に関連付けられた値です。

例：::

import torch.distributed as dist from datetime import timedelta store = dist.TCPStore("127.0.0.1", 0, 1, True, timedelta(seconds=30)) store.set("first_key", "first_value")

"first_value" を返すはずです。

store.get("first_key")

set_timeout(self: torch._C._distributed_c10d.Store, arg0: datetime.timedelta) → None# ストアのデフォルトタイムアウトを設定します。このタイムアウトは、初期化時と wait()、get() で使われます。

パラメーター

timeout (timedelta) – ストアに設定するタイムアウトです。

例：::

import torch.distributed as dist from datetime import timedelta

例として TCPStore を使用します。他のストア型も使えます。

store = dist.TCPStore("127.0.0.1", 0, 1, True, timedelta(seconds=30)) store.set_timeout(timedelta(seconds=10))

10 秒後に例外が発生します。

store.wait(["bad_key"])

property timeout# ストアのタイムアウトを取得します。

wait(*args, **kwargs)# オーバーロードされた関数です。

wait(self: torch._C._distributed_c10d.Store, arg0: collections.abc.Sequence[str]) -> None

keys 内の各キーがストアに追加されるまで待機します。ストア初期化時に設定されたタイムアウトまでにすべてのキーが設定されない場合、wait は例外を発生させます。

パラメーター

keys (list) – ストアに設定されるまで待つキーのリストです。

例：::

import torch.distributed as dist from datetime import timedelta

例として TCPStore を使用します。他のストア型も使えます。

store = dist.TCPStore("127.0.0.1", 0, 1, True, timedelta(seconds=30))

30 秒後に例外が発生します。

store.wait(["bad_key"])

wait(self: torch._C._distributed_c10d.Store, arg0: collections.abc.Sequence[str], arg1: datetime.timedelta) -> None

keys 内の各キーがストアに追加されるまで待機します。指定されたタイムアウトまでにキーが設定されない場合、例外が発生します。

パラメーター

keys (list) – ストアに設定されるまで待つキーのリストです。

timeout (timedelta) – キーが追加されるのを待つ時間です。経過すると例外が発生します。

例：::

import torch.distributed as dist from datetime import timedelta

例として TCPStore を使用します。他のストア型も使えます。

store = dist.TCPStore("127.0.0.1", 0, 1, True, timedelta(seconds=30))

10 秒後に例外が発生します。

store.wait(["bad_key"], timedelta(seconds=10))

class torch.distributed.TCPStore# TCP ベースの分散キー・バリューストア実装です。サーバーストアがデータを保持し、クライアントストアは TCP 経由でサーバーストアに接続して、キー・値ペアを挿入する set()、キー・値ペアを取得する get() などの操作を実行できます。クライアントストアはサーバーとの接続が確立されるまで待機するため、サーバーストアを必ず 1 つ初期化してください。

パラメーター

host_name (str) – サーバーストアを実行するホスト名または IP アドレスです。

port (int) – サーバーストアが受信リクエストを待ち受けるポートです。

world_size (int, optional) – ストアのユーザー総数です（クライアント数 + サーバー 1 つ）。デフォルトは None です（None はストアユーザー数が固定されないことを示します）。

is_master (bool, optional) – サーバーストアを初期化する場合は True、クライアントストアの場合は False です。デフォルトは False です。

timeout (timedelta, optional) – 初期化や get()、wait() などのメソッドでストアが使うタイムアウトです。デフォルトは timedelta(seconds=300) です。

wait_for_workers (bool, optional) – すべてのワーカーがサーバーストアに接続するまで待機するかどうかです。world_size が固定値の場合のみ適用されます。デフォルトは True です。

multi_tenant (bool, optional) – True の場合、現在のプロセス内で同じ host／port を使うすべての TCPStore インスタンスが同じ基盤 TCPServer を使います。デフォルトは False です。

master_listen_fd (int, optional) – 指定した場合、基盤の TCPServer はこのファイル記述子で待ち受けます。これはすでにポートにバインドされているソケットでなければなりません。一時ポートをバインドするには、port を 0 に設定して .port を読み取る方法を推奨します。デフォルトは None です（サーバーが新しいソケットを作成し、指定ポートへのバインドを試みます）。

use_libuv (bool, optional) – True の場合、TCPServer バックエンドに libuv を使います。デフォルトは True です。

例：::

import torch.distributed as dist from datetime import timedelta

プロセス 1（サーバー）で実行

server_store = dist.TCPStore("127.0.0.1", 1234, 2, True, timedelta(seconds=30))

プロセス 2（クライアント）で実行

client_store = dist.TCPStore("127.0.0.1", 1234, 2, False)

初期化後は、サーバーまたはクライアントのどちらからでもストアのメソッドを使えます。

server_store.set("first_key", "first_value") client_store.get("first_key")

init(self: torch._C._distributed_c10d.TCPStore, host_name: str, port: SupportsInt, world_size: SupportsInt | None = None, is_master: bool = False, timeout: datetime.timedelta = datetime.timedelta(seconds=300), wait_for_workers: bool = True, multi_tenant: bool = False, master_listen_fd: SupportsInt | None = None, use_libuv: bool = True) → None# 新しい TCPStore を作成します。

property host# ストアがリクエストを待ち受けるホスト名を取得します。

property libuvBackend# libuv バックエンドを使用している場合は True を返します。

property port# ストアがリクエストを待ち受けるポート番号を取得します。

class torch.distributed.HashStore# 基盤となるハッシュマップに基づく、スレッドセーフなストア実装です。このストアは同じプロセス内（たとえば他のスレッドから）で使えますが、プロセス間では使えません。

例：::

import torch.distributed as dist store = dist.HashStore()

他のスレッドからストアを使えます。

初期化後はストアの任意のメソッドを使えます。

store.set("first_key", "first_value")

init(self: torch._C._distributed_c10d.HashStore) → None# 新しい HashStore を作成します。

class torch.distributed.FileStore# 基盤となるキー・値ペアをファイルに保存するストア実装です。

パラメーター

file_name (str) – キー・値ペアを保存するファイルのパスです。

world_size (int, optional) – ストアを使うプロセスの総数です。デフォルトは -1 です（負の値はストアユーザー数が固定されないことを示します）。

例：::

import torch.distributed as dist store1 = dist.FileStore("/tmp/filestore", 2) store2 = dist.FileStore("/tmp/filestore", 2)

初期化後は、クライアントまたはサーバーのどちらからでもストアのメソッドを使えます。

store1.set("first_key", "first_value") store2.get("first_key")

init(self: torch._C._distributed_c10d.FileStore, file_name: str, world_size: SupportsInt = -1) → None# 新しい FileStore を作成します。

property path# FileStore がキー・値ペアの保存に使うファイルのパスを取得します。

class torch.distributed.PrefixStore# 3 種類のキー・バリューストア（TCPStore、FileStore、HashStore）のいずれかをラップし、ストアに挿入する各キーにプレフィックスを追加するラッパーです。

パラメーター

prefix (str) – ストアに挿入する前に各キーの先頭に追加するプレフィックス文字列です。

store (torch.distributed.store) – 基盤となるキー・バリューストアを構成するストアオブジェクトです。

init(self: torch._C._distributed_c10d.PrefixStore, prefix: str, store: torch._C._distributed_c10d.Store) → None# 新しい PrefixStore を作成します。

property underlying_store# PrefixStore がラップする基盤ストアオブジェクトを取得します。

集合通信のプロファイリング#

ここで説明した集合通信およびポイントツーポイント通信 API のプロファイリングには、torch.profiler（推奨。1.8.1 以降でのみ利用可能）または torch.autograd.profiler を使えます。標準で用意されているすべてのバックエンド（gloo、nccl、mpi）がサポートされており、集合通信の使用状況はプロファイリング出力／トレースに期待どおり表示されます。コードのプロファイリング方法は、通常の torch 演算子と同じです。

import torch import torch.distributed as dist

with torch.profiler(): tensor = torch.randn(20, 10) dist.all_reduce(tensor)

プロファイラー機能の全体像については、プロファイラーのドキュメントを参照してください。

マルチ GPU 集合通信関数#

警告 マルチ GPU 関数（CPU スレッドあたり複数 GPU を意味します）は非推奨です。現在、PyTorch Distributed が推奨するプログラミングモデルは、このドキュメントの API が示すように、スレッドあたり 1 デバイスです。バックエンドの開発者で、スレッドあたり複数デバイスをサポートしたい場合は、PyTorch Distributed のメンテナーにお問い合わせください。

オブジェクト集合通信#

警告 オブジェクト集合通信には、深刻な制約がいくつかあります。ユースケースで安全に使えるかどうかを判断するため、以下をよく読んでください。

オブジェクト集合通信は、pickle 化可能である限り、任意の Python オブジェクトを扱う集合通信に類似した操作です。broadcast、all_gather など、さまざまな集合通信パターンが実装されていますが、いずれもおおむね次の手順に従います。

入力オブジェクトを pickle（生バイト列）に変換し、バイトテンソルに格納する

このバイトテンソルのサイズをピアに通知する（最初の集合通信操作）

実際の集合通信を行うために、適切なサイズのテンソルを割り当てる

オブジェクトデータを通信する（2 番目の集合通信操作）

生データを Python オブジェクトに戻す（unpickle）

オブジェクト集合通信では、予想外のパフォーマンスやメモリー上の特性により、実行時間が長くなったり OOM が発生したりすることがあります。そのため、注意して使う必要があります。よくある問題を以下に示します。

pickle／unpickle 時間の非対称性 — オブジェクトの数、型、サイズによっては、pickle 化に時間がかかることがあります。集合通信が fan-in を持つ場合（例：gather_object）、受信側のランクは送信側のランクが pickle 化したオブジェクトの N 倍の数を unpickle する必要があり、他のランクが次の集合通信でタイムアウトする可能性があります。

非効率なテンソル通信 — テンソルはオブジェクト集合通信 API ではなく、通常の集合通信 API を使って送信してください。オブジェクト集合通信 API でテンソルを送信することもできますが、シリアライズとデシリアライズが行われます（CPU 以外のテンソルの場合は CPU 同期とデバイスからホストへのコピーも含みます）。デバッグやトラブルシューティング以外のほぼすべてのケースで、コードをリファクタリングして非オブジェクト集合通信を使う価値があります。

予期しないテンソルデバイス — オブジェクト集合通信を使ってテンソルを送信する場合、cuda（および他のアクセラレーター）テンソル特有の注意点があります。現在 cuda:3 にあるテンソルを pickle 化してから unpickle すると、実行中のプロセスや、そのプロセスの「デフォルト」CUDA デバイスにかかわらず、cuda:3 上に別のテンソルが作成されます。通常のテンソル集合通信 API では、「出力テンソル」は常に同じローカルデバイス上にあるため、一般にこちらが期待される動作です。テンソルを unpickle すると、そのプロセスで初めて GPU を使う場合に CUDA コンテキストが暗黙的に有効になり、大量の GPU メモリーが無駄に使われる可能性があります。この問題を避けるには、オブジェクト集合通信の入力として渡す前にテンソルを CPU に移動します。

サードパーティーバックエンド#

組み込みの GLOO／MPI／NCCL バックエンドに加え、PyTorch の分散パッケージは、実行時の登録メカニズムを介してサードパーティーバックエンドをサポートしています。C++ Extension を使ってサードパーティーバックエンドを開発する方法の参考資料については、Tutorials - Custom C++ and CUDA Extensions と test/cpp_extensions/cpp_c10d_extension.cpp を参照してください。サードパーティーバックエンドの機能は、各バックエンドの実装によって決まります。新しいバックエンドは c10d::ProcessGroup から派生し、インポート時に torch.distributed.Backend.register_backend() を使ってバックエンド名とインスタンス化インターフェイスを登録します。このバックエンドを手動でインポートし、対応するバックエンド名を指定して torch.distributed.init_process_group() を呼び出すと、torch.distributed パッケージで新しいバックエンドが使われます。

警告 サードパーティーバックエンドのサポートは実験的であり、今後変更される可能性があります。

起動ユーティリティ#

torch.distributed パッケージは、torch.distributed.launch に起動ユーティリティも提供しています。このヘルパーユーティリティを使うと、分散トレーニング用にノードごとに複数のプロセスを起動できます。

モジュール torch.distributed.launch。

torch.distributed.launch は、各トレーニングノード上で複数の分散トレーニングプロセスを起動するモジュールです。

警告 このモジュールは非推奨となり、torchrun に置き換えられる予定です。

このユーティリティは、ノードあたり 1 つ以上のプロセスを起動する単一ノード分散トレーニングに使えます。CPU トレーニングと GPU トレーニングのどちらにも使用できます。GPU トレーニングに使う場合、各分散プロセスは 1 つの GPU を操作します。これにより、単一ノードのトレーニング性能を大幅に向上できます。また、各ノードで複数のプロセスを起動することで、複数ノードの分散トレーニングにも使え、複数ノードのトレーニング性能も大幅に向上できます。これは、GPU への直接アクセスをサポートする複数の InfiniBand インターフェイスを備え、集約通信帯域幅としてすべてを活用できるシステムで特に有効です。

単一ノード分散トレーニングでも複数ノード分散トレーニングでも、このユーティリティは指定された数のプロセスをノードごとに起動します（--nproc-per-node）。GPU トレーニングで使う場合、この数は現在のシステム上の GPU 数（nproc_per_node）以下である必要があります。また、各プロセスは GPU 0 から GPU (nproc_per_node - 1) までのいずれか 1 つを操作します。

このモジュールの使い方：

単一ノードのマルチプロセス分散トレーニング

python -m torch.distributed.launch --nproc-per-node=NUM_GPUS_YOU_HAVE YOUR_TRAINING_SCRIPT.py (--arg1 --arg2 --arg3 and all other arguments of your training script)

複数ノードのマルチプロセス分散トレーニング：（例：2 ノード）

ノード 1：（IP：192.168.1.1、空きポート：1234）

python -m torch.distributed.launch --nproc-per-node=NUM_GPUS_YOU_HAVE --nnodes=2 --node-rank=0 --master-addr="192.168.1.1" --master-port=1234 YOUR_TRAINING_SCRIPT.py (--arg1 --arg2 --arg3 and all other arguments of your training script)

ノード 2：

python -m torch.distributed.launch --nproc-per-node=NUM_GPUS_YOU_HAVE --nnodes=2 --node-rank=1 --master-addr="192.168.1.1" --master-port=1234 YOUR_TRAINING_SCRIPT.py (--arg1 --arg2 --arg3 and all other arguments of your training script)

このモジュールが提供するオプション引数を確認するには、次を実行します。

python -m torch.distributed.launch --help

重要な注意事項：

1. このユーティリティとマルチプロセス分散 GPU トレーニング（単一ノードまたは複数ノード）では、現在 NCCL 分散バックエンドを使用した場合に最高の性能が得られます。そのため、GPU トレーニングには NCCL バックエンドを推奨します。

2. トレーニングプログラムでは、このモジュールから渡されるコマンドライン引数 --local-rank=LOCAL_PROCESS_RANK を解析する必要があります。GPU を使う場合、コードが LOCAL_PROCESS_RANK の GPU デバイス上でのみ実行されるようにしてください。次の方法で実現できます。

local_rank 引数を解析する

import argparse parser = argparse.ArgumentParser() parser.add_argument("--local-rank", "--local_rank", type=int) args = parser.parse_args()

次のいずれかを使ってデバイスを local rank に設定する

torch.cuda.set_device(args.local_rank) # コードの実行前 または with torch.cuda.device(args.local_rank):

実行するコード

...

バージョン 2.0.0 で変更：ランチャーは --local-rank= 引数をスクリプトに渡すようになりました。PyTorch 2.0.0 以降では、以前使われていたアンダースコア付きの --local_rank よりも、ハイフン付きの --local-rank が推奨されます。後方互換性のため、引数解析コードで両方に対応する必要がある場合があります。つまり、引数パーサーに "--local-rank" と "--local_rank" の両方を含めてください。"--local_rank" だけを指定した場合、ランチャーは「error: unrecognized arguments: –local-rank=」というエラーを発生させます。PyTorch 2.0.0 以降のみをサポートするトレーニングコードでは、"--local-rank" だけを含めれば十分です。

3. トレーニングプログラムでは、分散バックエンドを起動するため、開始時に次の関数を呼び出す必要があります。init_method=env:// を強く推奨します。他の初期化方法（例：tcp://）も動作する場合がありますが、このモジュールが公式にサポートするのは env:// です。

torch.distributed.init_process_group(backend='YOUR BACKEND', init_method='env://')

4. トレーニングプログラムでは、通常の分散関数を使うことも、torch.nn.parallel.DistributedDataParallel() モジュールを使うこともできます。GPU を使ってトレーニングし、torch.nn.parallel.DistributedDataParallel() モジュールを使う場合は、次のように設定します。

model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.local_rank], output_device=args.local_rank)

device_ids 引数には、コードが操作する唯一の GPU デバイス ID を設定してください。通常、これはプロセスのローカルランクです。つまり、このユーティリティを使うには、device_ids を [args.local_rank]、output_device を args.local_rank に設定する必要があります。

5. 環境変数 LOCAL_RANK を介して local_rank をサブプロセスに渡す方法もあります。この動作は --use-env=True を指定してスクリプトを起動した場合に有効になります。上のサブプロセスの例を調整し、args.local_rank を os.environ['LOCAL_RANK'] に置き換える必要があります。このフラグを指定すると、ランチャーは --local-rank を渡しません。

警告 local_rank はグローバルに一意ではありません。マシンごとにのみ一意です。したがって、ネットワークファイルシステムへの書き込みなどを行うかどうかの判定に使わないでください。正しく扱わなかった場合に問題が起きる例については、pytorch/pytorch#12042 を参照してください。

Spawn ユーティリティ#

Multiprocessing package - torch.multiprocessing パッケージは、torch.multiprocessing.spawn() にも spawn 関数を提供しています。このヘルパー関数を使うと、複数のプロセスを起動できます。実行する関数を渡すと、N 個のプロセスを起動してその関数を実行します。マルチプロセス分散トレーニングにも使えます。使用方法の参考資料については、PyTorch example - ImageNet implementation を参照してください。

この関数には Python 3.4 以降が必要です。

torch.distributed アプリケーションのデバッグ#

分散アプリケーションのデバッグは、理解しにくいハングやクラッシュ、ランク間の一貫性のない動作が発生するため、難しい場合があります。torch.distributed は、トレーニングアプリケーションをユーザー自身でデバッグするためのツール一式を提供しています。

Python Breakpoint#

分散環境で Python のデバッガーを使うと非常に便利ですが、そのままでは動作しないため、利用していないユーザーも多くいます。PyTorch は pdb をカスタマイズしたラッパーを提供し、デバッグの流れを簡素化します。torch.distributed.breakpoint を使うと簡単にデバッグできます。内部では pdb のブレークポイント動作を 2 点カスタマイズしていますが、それ以外は通常の pdb と同じように動作します。

ユーザーが指定した単一ランクのみにデバッガーを接続します。

デバッグ対象のランクが continue を実行するまで、torch.distributed.barrier() を使って他のすべてのランクを停止させます。

子プロセスの stdin をリダイレクトし、ターミナルに接続します。

使うには、すべてのランクで同じ rank 値を指定して torch.distributed.breakpoint(rank) を呼び出すだけです。

Monitored Barrier#

v1.10 以降、torch.distributed.monitored_barrier() は torch.distributed.barrier() の代替として使えます。指定されたタイムアウト内にすべてのランクが torch.distributed.monitored_barrier() を呼び出さないなどの障害が発生した場合、問題がある可能性のあるランクについて役立つ情報を含むエラーを発生させます。torch.distributed.monitored_barrier() は、確認応答に似た send／recv 通信プリミティブを使ってホスト側のバリアを実装し、rank 0 が時間内にバリアを確認しなかったランクを報告できるようにします。rank 1 が torch.distributed.monitored_barrier() を呼び出さない次の関数を例に考えてみましょう（実際には、アプリケーションのバグや以前の集合通信のハングが原因の場合があります）。

import os from datetime import timedelta import torch import torch.distributed as dist import torch.multiprocessing as mp

def worker(rank): dist.init_process_group("nccl", rank=rank, world_size=2) # monitored barrier では、ホスト側の同期を行うために gloo プロセスグループが必要です。 group_gloo = dist.new_group(backend="gloo") if rank not in [1]: dist.monitored_barrier(group=group_gloo, timeout=timedelta(seconds=2))

if name == "main": os.environ["MASTER_ADDR"] = "localhost" os.environ["MASTER_PORT"] = "29501" mp.spawn(worker, nprocs=2, args=())

rank 0 では次のエラーメッセージが生成されます。ユーザーはこれをもとに、問題がある可能性のあるランクを特定して、さらに調査できます。

RuntimeError: Rank 1 failed to pass monitoredBarrier in 2000 ms Original exception: [gloo/transport/tcp/pair.cc:598] Connection closed by peer [2401:db00:eef0:1100:3560:0:1c05:25d]:8594

TORCH_DISTRIBUTED_DEBUG#

TORCH_CPP_LOG_LEVEL=INFO と組み合わせると、環境変数 TORCH_DISTRIBUTED_DEBUG を使って追加の有用なログ出力や集合通信の同期チェックを有効にし、すべてのランクが適切に同期されていることを確認できます。必要なデバッグレベルに応じて、TORCH_DISTRIBUTED_DEBUG は OFF（デフォルト）、INFO、DETAIL のいずれかに設定できます。最も詳細なオプションである DETAIL はアプリケーションのパフォーマンスに影響する場合があるため、問題のデバッグ時にのみ使用してください。TORCH_DISTRIBUTED_DEBUG=INFO を設定すると、torch.nn.parallel.DistributedDataParallel() でトレーニングするモデルの初期化時に追加のデバッグログが出力されます。TORCH_DISTRIBUTED_DEBUG=DETAIL ではさらに、一部のイテレーションについて実行時のパフォーマンス統計がログに出力されます。統計には、順伝播時間、逆伝播時間、勾配通信時間などのデータが含まれます。

次のアプリケーションを例に説明します。

import os import torch import torch.distributed as dist import torch.multiprocessing as mp

class TwoLinLayerNet(torch.nn.Module): def init(self): super().init() self.a = torch.nn.Linear(10, 10, bias=False) self.b = torch.nn.Linear(10, 1, bias=False)

def forward(self, x):
    a = self.a(x)
    b = self.b(x)
    return (a, b)

def worker(rank): dist.init_process_group("nccl", rank=rank, world_size=2) torch.cuda.set_device(rank) print("init model") model = TwoLinLayerNet().cuda() print("init ddp") ddp_model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[rank]) inp = torch.randn(10, 10).cuda() print("train") for _ in range(20): output = ddp_model(inp) loss = output[0] + output[1] loss.sum().backward()

if name == "main": os.environ["MASTER_ADDR"] = "localhost" os.environ["MASTER_PORT"] = "29501" os.environ["TORCH_CPP_LOG_LEVEL"]="INFO" os.environ[ "TORCH_DISTRIBUTED_DEBUG" ] = "DETAIL" # 実行時のログには DETAIL を設定します。 mp.spawn(worker, nprocs=2, args=())

初期化時には次のログが出力されます。

I0607 16:10:35.739390 515217 logger.cpp:173] [Rank 0]: DDP Initialized with: broadcast_buffers: 1 bucket_cap_bytes: 26214400 find_unused_parameters: 0 gradient_as_bucket_view: 0 is_multi_device_module: 0 iteration: 0 num_parameter_tensors: 2 output_device: 0 rank: 0 total_parameter_size_bytes: 440 world_size: 2 backend_name: nccl bucket_sizes: 440 cuda_visible_devices: N/A device_ids: 0 dtypes: float master_addr: localhost master_port: 29501 module_name: TwoLinLayerNet nccl_async_error_handling: N/A nccl_blocking_wait: N/A nccl_debug: WARN nccl_ib_timeout: N/A nccl_nthreads: N/A nccl_socket_ifname: N/A torch_distributed_debug: INFO

実行時には（TORCH_DISTRIBUTED_DEBUG=DETAIL が設定されている場合）次のログが出力されます。

I0607 16:18:58.085681 544067 logger.cpp:344] [Rank 1 / 2] Training TwoLinLayerNet unused_parameter_size=0 Avg forward compute time: 40838608 Avg backward compute time: 5983335 Avg backward comm. time: 4326421 Avg backward comm/comp overlap time: 4207652

I0607 16:18:58.085693 544066 logger.cpp:344] [Rank 0 / 2] Training TwoLinLayerNet unused_parameter_size=0 Avg forward compute time: 42850427 Avg backward compute time: 3885553 Avg backward comm. time: 2357981 Avg backward comm/comp overlap time: 2234674

さらに、TORCH_DISTRIBUTED_DEBUG=INFO は、モデル内の未使用パラメーターが原因で torch.nn.parallel.DistributedDataParallel() がクラッシュする際のログを改善します。現在、順伝播で使われない可能性があるパラメーターがある場合は、torch.nn.parallel.DistributedDataParallel() の初期化に find_unused_parameters=True を渡す必要があります。また、v1.10 以降では、torch.nn.parallel.DistributedDataParallel() は逆伝播で未使用のパラメーターをサポートしていないため、モデルのすべての出力を損失計算に使う必要があります。特に大規模モデルでは、これらの制約が難しい場合があります。そのため、エラーでクラッシュする際、torch.nn.parallel.DistributedDataParallel() は未使用だったすべてのパラメーターの完全修飾名をログに出力します。

たとえば、上記のアプリケーションで、loss を loss = output[1] と計算するように変更すると、逆伝播で TwoLinLayerNet.a は勾配を受け取らないため、DDP が失敗します。クラッシュ時には、特に大規模モデルで手作業による特定が難しい場合がある、未使用パラメーターに関する情報がユーザーに渡されます。

RuntimeError: Expected to have finished reduction in the prior iteration before starting a new one. This error indicates that your module has parameters that were not used in producing loss. You can enable unused parameter detection by passing the keyword argument find_unused_parameters=True to torch.nn.parallel.DistributedDataParallel, and by making sure all forward function outputs participate in calculating loss. If you already have done the above, then the distributed data parallel module wasn't able to locate the output tensors in the return value of your module's forward function. Please include the loss function and the structure of the return va lue of forward of your module when reporting this issue (e.g. list, dict, iterable).

勾配を受け取らなかったランク 0 のパラメーター：a.weight

勾配を受け取らなかったランク 0 のパラメーターインデックス：0

TORCH_DISTRIBUTED_DEBUG=DETAIL を設定すると、ユーザーが直接、または間接的に（DDP allreduce など）発行するすべての集合通信呼び出しで、一貫性と同期に関する追加チェックが実行されます。これは、torch.distributed.init_process_group() および torch.distributed.new_group() API が返すすべてのプロセスグループをラップする、ラッパープロセスグループを作成することで実現します。その結果、これらの API は通常のプロセスグループと同じように使えるラッパープロセスグループを返しますが、集合通信を基盤となるプロセスグループにディスパッチする前に一貫性チェックを実行します。現在、このチェックには torch.distributed.monitored_barrier() が含まれ、すべてのランクで保留中の集合通信呼び出しが完了したことを確認し、停止しているランクを報告します。次に、すべての集合通信関数が一致し、テンソルの形状も一貫していることを確認して、集合通信自体の整合性を検査します。一致しない場合、アプリケーションがクラッシュしたときに、ハングや情報の少ないエラーメッセージではなく、詳細なエラー報告が表示されます。

入力形状が torch.distributed.all_reduce() で一致しない次の関数を例に考えてみましょう。

import torch import torch.distributed as dist import torch.multiprocessing as mp

def worker(rank): dist.init_process_group("nccl", rank=rank, world_size=2) torch.cuda.set_device(rank) tensor = torch.randn(10 if rank == 0 else 20).cuda() dist.all_reduce(tensor) torch.cuda.synchronize(device=rank)

if name == "main": os.environ["MASTER_ADDR"] = "localhost" os.environ["MASTER_PORT"] = "29501" os.environ["TORCH_CPP_LOG_LEVEL"]="INFO" os.environ["TORCH_DISTRIBUTED_DEBUG"] = "DETAIL" mp.spawn(worker, nprocs=2, args=())

NCCL バックエンドでは、このようなアプリケーションはハングする可能性が高く、複雑な状況では原因の特定が困難です。TORCH_DISTRIBUTED_DEBUG=DETAIL を有効にしてアプリケーションを再実行すると、次のエラーメッセージから根本原因がわかります。

work = default_pg.allreduce([tensor], opts) RuntimeError: Error when verifying shape tensors for collective ALLREDUCE on rank 0. This likely indicates that input shapes into the collective are mismatched across ranks. Got shapes: 10 [ torch.LongTensor{1} ]

注記 実行時のデバッグレベルを細かく制御するには、torch.distributed.set_debug_level()、torch.distributed.set_debug_level_from_env()、torch.distributed.get_debug_level() 関数も使えます。また、TORCH_DISTRIBUTED_DEBUG=DETAIL と TORCH_SHOW_CPP_STACKTRACES=1 を組み合わせると、集合通信の非同期状態が検出されたときに、コールスタック全体がログに出力されます。これらの集合通信の非同期状態チェックは、torch.distributed.init_process_group() および torch.distributed.new_group() API で作成されたプロセスグループを基盤とする c10d 集合通信呼び出しを使うすべてのアプリケーションで機能します。

ログ記録#

torch.distributed.monitored_barrier() と TORCH_DISTRIBUTED_DEBUG による明示的なデバッグサポートに加え、torch.distributed の基盤となる C++ ライブラリもさまざまなレベルでログメッセージを出力します。これらのメッセージは、分散トレーニングジョブの実行状態を把握し、ネットワーク接続の失敗などの問題をトラブルシューティングする際に役立ちます。次の表は、環境変数 TORCH_CPP_LOG_LEVEL と TORCH_DISTRIBUTED_DEBUG の組み合わせによって、ログレベルを調整する方法を示しています。

TORCH_CPP_LOG_LEVEL	TORCH_DISTRIBUTED_DEBUG	有効なログレベル
ERROR	無視	Error
WARNING	無視	Warning
INFO	無視	Info
INFO	INFO	Debug
INFO	DETAIL	Trace（別名：All）

分散コンポーネントは、RuntimeError から派生したカスタム Exception 型を発生させます。

torch.distributed.DistError：すべての分散例外の基底型です。

torch.distributed.DistBackendError：バックエンド固有のエラーが発生したときに発生します。たとえば NCCL バックエンド使用時に、NCCL ライブラリで利用できない GPU をユーザーが使おうとした場合です。

torch.distributed.DistNetworkError：ネットワークライブラリでエラーが発生したときに発生します（例：Connection reset by peer）。

torch.distributed.DistStoreError：Store でエラーが発生したときに発生します（例：TCPStore のタイムアウト）。

class torch.distributed.DistError# 分散ライブラリでエラーが発生した場合に発生する例外です。

class torch.distributed.DistBackendError# 分散処理でバックエンドエラーが発生した場合に発生する例外です。

class torch.distributed.DistNetworkError# 分散処理でネットワークエラーが発生した場合に発生する例外です。

class torch.distributed.DistStoreError# 分散ストアでエラーが発生した場合に発生する例外です。

単一ノードトレーニングを実行している場合、スクリプトに対話形式でブレークポイントを設定すると便利です。単一ランクに簡単にブレークポイントを設定する方法を提供しています。

torch.distributed.breakpoint(rank=0, skip=0, timeout_s=3600)[source]# 単一のランクにのみブレークポイントを設定します。デバッグが完了するまで、他のすべてのランクは待機します。

パラメーター

rank (int) – ブレークポイントを設定するランクです。デフォルト：0

skip (int) – このブレークポイントへの最初の skip 回の呼び出しをスキップします。デフォルト：0

torch.distributed

パターン 3: 初期化# 他のメソッドを呼び出す前に、torch.distributed.init_process_group() または torch.distributed.device_mesh.init_device_mesh() 関数を使ってパッケージを初期化する必要があります。どちらも、すべてのプロセスが参加するまでブロックします。 警告 初期化はスレッドセーフではありません。ランク間で不整合な「UUID」が割り当てられること、および初期化中の競合によってハングが発生することを防ぐため、プロセスグループの作成は単一のスレッドから実行してください。 torch.distributed.is_available()[source]# 分散パッケージが利用可能な場合は True を返します。利用できない場合、torch.distributed は他の API を公開しません。現在、torch.distributed は Linux、MacOS、Windows で利用できます。ソースから PyTorch をビルドする際に有効にするには USE_DISTRIBUTED=1 を設定します。現在、デフォルト値は Linux と Windows では USE_DISTRIBUTED=1、MacOS では USE_DISTRIBUTED=0 です。 戻り値の型 bool torch.distributed.init_process_group(backend=None, init_method=None, timeout=None, world_size=-1, rank=-1, store=None, group_name='', pg_options=None, device_id=None)[source]# デフォルトの分散プロセスグループを初期化します。これにより分散パッケージも初期化されます。プロセスグループを初期化する主な方法は 2 つあります。 store、rank、world_size を明示的に指定する。 init_method（URL 文字列）を指定し、ピアを検出する場所や方法を示す。rank と world_size は任意で指定するか、必要なすべてのパラメーターを URL にエンコードして省略できます。どちらも指定しない場合、init_method は「env://」とみなされます。 パラメーター backend (str または Backend, 任意) – 使用するバックエンド。ビルド時の構成によって、有効な値は mpi、gloo、nccl、ucc、xccl、またはサードパーティのプラグインで登録されたものです。2.6 以降、backend が指定されない場合、c10d は device_id キーワード引数（指定されている場合）で示されたデバイスタイプに登録されたバックエンドを使用します。現在の既定の登録は、cuda では nccl、cpu では gloo、xpu では xccl です。backend も device_id も指定されない場合、c10d は実行時のマシン上のアクセラレーターを検出し、検出されたアクセラレーター（または cpu）に登録されたバックエンドを使用します。この項目は小文字の文字列（例: "gloo"）で指定でき、Backend 属性（例: Backend.GLOO）からもアクセスできます。マシンごとに複数のプロセスを nccl バックエンドで使用する場合、各プロセスは使用するすべての GPU に排他的にアクセスできる必要があります。GPU をプロセス間で共有すると、デッドロックや NCCL の無効な使用を引き起こす可能性があります。ucc バックエンドは実験的です。デバイスの既定のバックエンドは get_default_backend_for_device() で確認できます。 init_method (str, 任意) – プロセスグループの初期化方法を指定する URL。init_method も store も指定されない場合の既定値は「env://」です。store と排他的です。 world_size (int, 任意) – ジョブに参加するプロセス数。store を指定する場合は必須です。 rank (int, 任意) – 現在のプロセスのランク（0 から world_size-1 までの数値である必要があります）。store を指定する場合は必須です。 store (Store, 任意) – すべてのワーカーからアクセス可能なキーと値のストアで、接続やアドレス情報の交換に使用されます。init_method と排他的です。 timeout (timedelta, 任意) – プロセスグループに対して実行される操作のタイムアウト。既定値は NCCL で 10 分、その他のバックエンドで 30 分です。これはコレクティブ操作が非同期に中止され、プロセスがクラッシュするまでの時間です。CUDA の実行は非同期であり、失敗した非同期の NCCL 操作によって後続の CUDA 操作が破損したデータ上で実行される可能性があるため、ユーザーコードの実行を続けることは安全ではなくなります。TORCH_NCCL_BLOCKING_WAIT が設定されている場合、プロセスはブロックし、このタイムアウトまで待機します。 group_name (str, 任意, 非推奨) – グループ名。この引数は無視されます。 pg_options (ProcessGroupOptions, 任意) – 特定のプロセスグループの構築時に渡す必要がある追加オプションを指定するプロセスグループオプション。現在サポートされているオプションは、nccl バックエンド用の ProcessGroupNCCL.Options のみです。is_high_priority_stream を指定すると、計算カーネルが待機しているときに nccl バックエンドが高優先度の CUDA ストリームを使用できます。その他の NCCL 設定オプションについては、https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/api/types.html#ncclconfig-t を参照してください。 device_id (torch.device | int, 任意) – このプロセスが使用する単一の特定のデバイス。バックエンド固有の最適化を可能にします。現在、これは NCCL の場合のみ 2 つの効果があります。通信子（communicator）が即座に形成され（通常の遅延呼び出しではなく、ncclCommInit* を即座に呼び出す）、可能な場合はサブグループが ncclCommSplit を使用して、グループ作成の不要なオーバーヘッドを回避します。NCCL の初期化エラーを早期に知りたい場合にも、このフィールドを使用できます。int が指定された場合、API はコンパイル時のアクセラレーター種別が使用されると仮定します。 注 backend == Backend.MPI を有効にするには、MPI をサポートするシステム上でソースから PyTorch をビルドする必要があります。 注 複数バックエンドのサポートは実験的です。現在、backend を指定しない場合、gloo と nccl の両方のバックエンドが作成されます。CPU テンソルのコレクティブ操作には gloo バックエンドが、CUDA テンソルのコレクティブ操作には nccl バックエンドが使用されます。カスタムバックエンドは、「<device_type>:<backend_name>,<device_type>:<backend_name>」の形式の文字列を渡して指定できます。例: 「cpu:gloo,cuda:custom_backend」。 torch.distributed.device_mesh.init_device_mesh(device_type, mesh_shape, *, mesh_dim_names=None, backend_override=None)[source]# device_type、mesh_shape、mesh_dim_names パラメーターに基づいて DeviceMesh を初期化します。これは、mesh_shape の長さを n とする n 次元配列のレイアウトを持つ DeviceMesh を作成します。mesh_dim_names が指定されている場合、各次元は mesh_dim_names[i] というラベルが付けられます。 注 init_device_mesh は SPMD プログラミングモデルに従います。つまり、同じ PyTorch Python プログラムがクラスター内のすべてのプロセス/ランクで実行されます。mesh_shape（デバイスレイアウトを表す n 次元配列の次元）がすべてのランクで同一であることを確認してください。一貫性のない mesh_shape はハングを引き起こす可能性があります。 注 プロセスグループが見つからない場合、init_device_mesh は分散通信に必要な分散プロセスグループを内部で初期化します。 パラメーター device_type (str) – メッシュのデバイスタイプ。現在サポートされているのは「cpu」「cuda/cuda-like」「xpu」です。「cuda:0」のように GPU インデックスを含むデバイスタイプを渡すことは許可されていません。 mesh_shape (Tuple[int]) – デバイスのレイアウトを表す多次元配列の次元を定義するタプル。 mesh_dim_names (Tuple[str], 任意) – デバイスのレイアウトを表す多次元配列の各次元に割り当てるメッシュ次元名のタプル。その長さは mesh_shape の長さと一致する必要があります。mesh_dim_names の各文字列は一意である必要があります。 backend_override (Dict[int | str, tuple[str, Options] | str | Options], 任意) – 各メッシュ次元に対して作成される ProcessGroup の一部またはすべてに対する上書き設定。各キーは次元のインデックス、または mesh_dim_names が指定されている場合はその名前のいずれかです。各値は、バックエンド名とそのオプションを含むタプル、またはこれら 2 つの要素のいずれか一方のみ（もう一方は既定値に設定されます）を指定できます。 戻り値 デバイスのレイアウトを表す DeviceMesh オブジェクト。 戻り値の型 DeviceMesh 例: >>> from torch.distributed.device_mesh import init_device_mesh >>> >>> mesh_1d = init_device_mesh("cuda", mesh_shape=(8,)) >>> mesh_2d = init_device_mesh("cuda", mesh_shape=(2, 8), mesh_dim_names=("dp", "tp")) torch.distributed.is_initialized()[source]# デフォルトのプロセスグループが初期化済みかどうかを確認します。 戻り値の型 bool torch.distributed.is_mpi_available()[source]# MPI バックエンドが利用可能かどうかを確認します。 戻り値の型 bool torch.distributed.is_nccl_available()[source]# NCCL バックエンドが利用可能かどうかを確認します。 戻り値の型 bool torch.distributed.is_gloo_available()[source]# Gloo バックエンドが利用可能かどうかを確認します。 戻り値の型 bool torch.distributed.distributed_c10d.is_xccl_available()[source]# XCCL バックエンドが利用可能かどうかを確認します。 戻り値の型 bool torch.distributed.is_torchelastic_launched()[source]# このプロセスが torch.distributed.elastic（別名 torchelastic）で起動されたかどうかを確認します。TORCHELASTIC_RUN_ID 環境変数の存在を、現在のプロセスが torchelastic で起動されたかどうかの判定に使用します。これは妥当な判定基準です。TORCHELASTIC_RUN_ID はランデブー ID に対応し、ピア検出のためのジョブ ID を示す常に非 null の値だからです。。 戻り値の型 bool torch.distributed.get_default_backend_for_device(device)[source]# 指定されたデバイスの既定のバックエンドを返します。 パラメーター device (Union[str, torch.device]) – 既定のバックエンドを取得する対象のデバイス。 戻り値 指定されたデバイスの既定のバックエンドを小文字の文字列として返します。 戻り値の型 str 現在、3 つの初期化方法がサポートされています。 TCP 初期化# TCP を使って初期化する方法は 2 つあります。どちらも、すべてのプロセスから到達可能なネットワークアドレスと、希望する world_size が必要です。1 つ目の方法は、ランク 0 のプロセスに属するアドレスを指定します。この初期化方法では、すべてのプロセスのランクを手動で指定する必要があります。最新の分散パッケージではマルチキャストアドレスはサポートされなくなった点に注意してください。group_name も非推奨です。 import torch.distributed as dist # いずれかのマシンのアドレスを使用 dist.init_process_group(backend, init_method='tcp://10.1.1.20:23456', rank=args.rank, world_size=4) 共有ファイルシステム初期化# もう 1 つの初期化方法は、グループ内のすべてのマシンから共有されて見えるファイルシステムと、希望する world_size を利用します。URL は file:// で始まり、共有ファイルシステム上の（既存のディレクトリ内にある）存在しないファイルのパスを含む必要があります。ファイルが存在しない場合、ファイルシステム初期化は自動的にそのファイルを作成しますが、削除はしません。したがって、同じファイルパス/名前に対する次の init_process_group() 呼び出しの前にファイルがクリーンアップされていることを確認するのはユーザーの責任です。最新の分散パッケージでは自動ランク割り当てはサポートされなくなった点、また group_name も非推奨である点に注意してください。 警告 この方法は、ファイルシステムが fcntl によるロックをサポートしていることを前提とします。ほとんどのローカルシステムと NFS はこれをサポートしています。 警告 この方法は常にファイルを作成し、プログラム終了時にできる限りクリーンアップして削除しようとします。言い換えると、ファイル初期化方式で初期化を行うたびに、初期化を成功させるには真新しい空のファイルが必要です。前回の初期化で使用され、クリーンアップされずに残ったファイルを再び使用すると、想定外の動作となり、多くの場合デッドロックや失敗を引き起こします。したがって、この方法はできる限りファイルを削除しようとしますが、自動削除に失敗した場合は、次回の実行時に同じファイルが再利用されないよう、学習の終了時にファイルが削除されていることを確認するのはユーザーの責任です。同じファイル名に対して init_process_group() を複数回呼び出す予定がある場合は、特にこの点が重要です。言い換えると、ファイルが削除/クリーンアップされないまま、そのファイルに対して init_process_group() を再度呼び出すと、失敗が予想されます。ここでの経験則は、init_process_group() が呼び出されるたびに、ファイルが存在しないか空であることを確認することです。 import torch.distributed as dist # rank は必ず指定すること dist.init_process_group(backend, init_method='file:///mnt/nfs/sharedfile', world_size=4, rank=args.rank) 環境変数による初期化# この方法は環境変数から設定を読み取るため、情報の取得方法を完全にカスタマイズできます。設定する変数は次のとおりです。 MASTER_PORT - 必須。ランク 0 のマシン上の空いているポートである必要があります MASTER_ADDR - 必須（ランク 0 を除く）。ランク 0 ノードのアドレス WORLD_SIZE - 必須。ここで設定するか、init 関数の呼び出しで設定できます RANK - 必須。ここで設定するか、init 関数の呼び出しで設定できます ランク 0 のマシンがすべての接続を確立するために使用されます。これは既定の方法であり、init_method を指定する必要はありません（または env:// を指定できます）。 初期化時間の短縮# TORCH_GLOO_LAZY_INIT - 完全なメッシュを使う代わりに、接続をオンデマンドで確立します。all-to-all 以外の操作では、初期化時間を大幅に短縮できます。

torch.distributed.init_process_group()

Pattern 4: Example:

>>> from torch.distributed.device_mesh import init_device_mesh
>>>
>>> mesh_1d = init_device_mesh("cuda", mesh_shape=(8,))
>>> mesh_2d = init_device_mesh("cuda", mesh_shape=(2, 8), mesh_dim_names=("dp", "tp"))

Pattern 5: Groups# デフォルトでは、コレクティブはデフォルトグループ（ワールドとも呼ばれます）上で動作し、すべてのプロセスが分散関数の呼び出しに参加する必要があります。しかし、ワークロードによっては、より細かい粒度の通信が有効な場合があります。ここで分散グループが役立ちます。new_group() 関数を使うと、すべてのプロセスの任意のサブセットを対象とした新しいグループを作成できます。この関数は不透明なグループハンドルを返し、これを各コレクティブ（コレクティブとは、特定の既知のプログラミングパターンで情報を交換する分散関数のことです）にグループ引数として渡すことができます。torch.distributed.new_group(ranks=None, timeout=None, backend=None, pg_options=None, use_local_synchronization=False, group_desc=None, device_id=None)[source]# 新しい分散グループを作成します。この関数は、メイングループ内のすべてのプロセス（つまり分散ジョブに含まれるすべてのプロセス）が、そのグループのメンバーにならない場合でも、この関数に入ることを必要とします。さらに、グループはすべてのプロセスで同じ順序で作成する必要があります。警告 安全な並行利用: NCCL バックエンドで複数のプロセスグループを使用する場合、ユーザーはランク間でコレクティブの実行順序がグローバルに一貫していることを保証する必要があります。プロセス内の複数のスレッドがコレクティブを発行する場合、一貫した順序を保つために明示的な同期が必要です。torch.distributed の通信 API の非同期版を使用すると、work オブジェクトが返され、通信カーネルは別の CUDA ストリーム上にエンキューされるため、通信と計算をオーバーラップさせることができます。あるプロセスグループで 1 つ以上の非同期操作を発行したら、別のプロセスグループを使用する前に work.wait() を呼び出して、他の CUDA ストリームと同期させる必要があります。詳細は、「Using multiple NCCL communicators concurrently」https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/usage/communicators.html#using-multiple-nccl-communicators-concurrently を参照してください。

パラメータ ranks (list[int]) – グループメンバーのランクのリスト。None の場合、すべてのランクに設定されます。デフォルトは None です。 timeout (timedelta, optional) – 詳細とデフォルト値は init_process_group を参照してください。 backend (str or Backend, optional) – 使用するバックエンド。ビルド時の構成によって、有効な値は gloo と nccl です。デフォルトでは、グローバルグループと同じバックエンドを使用します。この値は小文字の文字列（例: "gloo"）で指定する必要があり、Backend の属性（例: Backend.GLOO）を通じて指定することもできます。None が渡された場合、デフォルトのプロセスグループに対応するバックエンドが使用されます。デフォルトは None です。 pg_options (ProcessGroupOptions, optional) – 特定のプロセスグループの構築時に渡す必要がある追加オプションを指定するプロセスグループオプション。たとえば nccl バックエンドでは、is_high_priority_stream を指定することで、プロセスグループが高優先度の CUDA ストリームを使用できるようになります。NCCL の設定に関するその他の利用可能なオプションについては、https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/api/types.html#ncclconfig-t を参照してください。 use_local_synchronization (bool, optional): プロセスグループの作成の最後にグループローカルのバリアを実行します。これは、メンバーではないランクが API を呼び出す必要がなく、バリアにも参加しない点で異なります。 group_desc (str, optional) – プロセスグループを説明する文字列。 device_id (torch.device, optional) – このプロセスを「バインド」する単一の特定のデバイス。この値が指定された場合、new_group の呼び出しは、そのデバイスの通信バックエンドを直ちに初期化しようとします。

戻り値 コレクティブ呼び出しに渡せる分散グループのハンドル。ランクが ranks に含まれない場合は GroupMember.NON_GROUP_MEMBER を返します。

注意 use_local_synchronization は MPI では動作しません。 use_local_synchronization=True は、大規模なクラスターや小規模なプロセスグループで大幅に高速になる場合がありますが、メンバーではないランクがグループのバリア() に参加しなくなるため、クラスターの動作が変わる点に注意が必要です。 use_local_synchronization=True は、各ランクが複数の重複するプロセスグループを作成する場合にデッドロックを引き起こす可能性があります。これを避けるため、すべてのランクが同じグローバルな作成順序に従うようにしてください。

torch.distributed.get_group_rank(group, global_rank)[source]# グローバルランクをグループランクに変換します。global_rank は group に含まれている必要があり、そうでない場合は RuntimeError が発生します。

パラメータ group (ProcessGroup) – 相対ランクを求める対象の ProcessGroup。 global_rank (int) – 照会するグローバルランク。

戻り値 group を基準とした global_rank のグループランク

戻り値の型 int

注意: デフォルトのプロセスグループに対してこの関数を呼び出すと、恒等写像が返されます。

torch.distributed.get_global_rank(group, group_rank)[source]# グループランクをグローバルランクに変換します。group_rank は group に含まれている必要があり、そうでない場合は RuntimeError が発生します。

パラメータ group (ProcessGroup) – グローバルランクを求める対象の ProcessGroup。 group_rank (int) – 照会するグループランク。

戻り値 group を基準とした group_rank のグローバルランク

戻り値の型 int

注意: デフォルトのプロセスグループに対してこの関数を呼び出すと、恒等写像が返されます。

torch.distributed.get_process_group_ranks(group)[source]# group に関連付けられたすべてのランクを取得します。

パラメータ group (Optional[ProcessGroup]) – すべてのランクを取得する対象の ProcessGroup。None の場合、デフォルトのプロセスグループが使用されます。

戻り値 グループランクの順に並んだグローバルランクのリスト

戻り値の型 list[int]

new_group()

Pattern 6: Warning 安全な並行利用: NCCL バックエンドで複数のプロセスグループを使用する場合、ユーザーはランク間でコレクティブの実行順序がグローバルに一貫していることを保証する必要があります。プロセス内の複数のスレッドがコレクティブを発行する場合、一貫した順序を保証するために明示的な同期が必要です。torch.distributed の通信 API の非同期版を使用すると、work オブジェクトが返され、通信カーネルは別の CUDA ストリーム上にエンキューされるため、通信と計算をオーバーラップさせることができます。あるプロセスグループで 1 つ以上の非同期操作を発行した後、別のプロセスグループを使用する前に、work.wait() を呼び出して他の CUDA ストリームと同期させる必要があります。詳細は、NVIDIA の「Using multiple NCCL communicators concurrently」https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/usage/communicators.html#using-multiple-nccl-communicators-concurrently を参照してください。

NCCL

Pattern 7: Note DistributedDataParallel を Distributed RPC Framework と組み合わせて使用する場合は、勾配の計算には常に torch.distributed.autograd.backward() を、パラメーターの最適化には torch.distributed.optim.DistributedOptimizer を使用する必要があります。例: >>> import torch.distributed.autograd as dist_autograd >>> from torch.nn.parallel import DistributedDataParallel as DDP >>> import torch >>> from torch import optim >>> from torch.distributed.optim import DistributedOptimizer >>> import torch.distributed.rpc as rpc >>> from torch.distributed.rpc import RRef >>> >>> t1 = torch.rand((3, 3), requires_grad=True) >>> t2 = torch.rand((3, 3), requires_grad=True) >>> rref = rpc.remote("worker1", torch.add, args=(t1, t2)) >>> ddp_model = DDP(my_model) >>> >>> # オプティマイザーの設定 >>> optimizer_params = [rref] >>> for param in ddp_model.parameters(): >>> optimizer_params.append(RRef(param)) >>> >>> dist_optim = DistributedOptimizer( >>> optim.SGD, >>> optimizer_params, >>> lr=0.05, >>> ) >>> >>> with dist_autograd.context() as context_id: >>> pred = ddp_model(rref.to_here()) >>> loss = loss_func(pred, target) >>> dist_autograd.backward(context_id, [loss]) >>> dist_optim.step(context_id)

torch.distributed.autograd.backward()

Pattern 8: static_graph (bool) – True に設定すると、DDP は学習グラフが静的であることを認識します。静的グラフとは、1) 使用されるパラメーターと未使用のパラメーターの集合が学習ループ全体を通じて変化しないこと、2) 学習グラフの実行方法が学習ループ全体を通じて変化しないこと（つまり、イテレーションに依存する制御フローがないこと）を意味します。static_graph を True に設定すると、DDP は従来サポートされなかったケースに対応します: 1) リエントラントな backward。2) 複数回の活性化チェックポイント。3) モデルに未使用のパラメーターがある場合の活性化チェックポイント。4) forward 関数の外側にあるモデルパラメーターが存在する場合。5) 未使用のパラメーターがある場合のパフォーマンス向上の可能性。static_graph を True に設定すると、DDP は未使用のパラメーターを検出するためにイテレーションごとにグラフを探索しなくなるためです。static_graph を True に設定できるかどうかを確認するone つの方法は、以前のモデル学習の終了時に DDP のログデータを確認することです。ddp_logging_data.get("can_set_static_graph") == True であれば、通常は static_graph = True にも設定できます。例::>>> model_DDP = torch.nn.parallel.DistributedDataParallel(model) >>> # 学習ループ >>> ... >>> ddp_logging_data = model_DDP._get_ddp_logging_data() >>> static_graph = ddp_logging_data.get("can_set_static_graph")

True

参照ファイル

このスキルには、references/ に包括的なドキュメントが含まれています:

• other.md - その他のドキュメント

詳細な情報が必要な場合は、view を使用して特定の参照ファイルを読み取ってください。

このスキルの使い方

初心者向け

基礎的な概念については、getting_started または tutorials の参照ファイルから始めてください。

特定の機能について

詳細な情報は、該当するカテゴリの参照ファイル（api、guides など）を使用してください。

コード例について

上記のクイックリファレンスセクションには、公式ドキュメントから抽出した一般的なパターンが含まれています。

リソース

references/

公式ソースから抽出された整理済みのドキュメントです。これらのファイルには次の内容が含まれます。

• 詳細な説明
• 言語指定付きのコード例
• 元のドキュメントへのリンク
• すばやく移動するための目次

scripts/

一般的な自動化タスク用のヘルパースクリプトをここに追加します。

assets/

テンプレート、ボイラープレート、サンプルプロジェクトをここに追加します。

注意事項

• このスキルは公式ドキュメントから自動生成されました
• リファレンスファイルは、元のドキュメントの構造と例を保持しています
• コード例には、より良いシンタックスハイライトのための言語検出が含まれています
• クイックリファレンスのパターンは、ドキュメント内の一般的な使用例から抽出されています

更新

更新されたドキュメントでこのスキルを最新の状態にするには:

1. 同じ設定でスクレイパーを再実行します
2. スキルが最新の情報で再構築されます