# Distributed communication package - torch.distributed Note Please refer to [PyTorch Distributed Overview](https://pytorch.org/tutorials/beginner/dist_overview.html) for a brief introduction to all features related to distributed training. torch.distributed.elastic.utils.api.get_env_variable_or_raise(*env_name*)[[source]](https://github.com/pytorch/pytorch/blob/v2.13.0/torch/distributed/elastic/utils/api.py#L15) Tries to retrieve environment variable. Raises `ValueError` if no environment variable found. Parameters: **env_name** ([*str*](https://docs.python.org/3/library/stdtypes.html#str)) - Name of the env variable Return type: [str](https://docs.python.org/3/library/stdtypes.html#str) torch.distributed.elastic.utils.distributed.get_free_port()[[source]](https://github.com/pytorch/pytorch/blob/v2.13.0/torch/distributed/elastic/utils/distributed.py#L129) Returns an unused port on localhost. This function finds an unused port on localhost by opening to socket to bind to a port and then closing it. Returns: an unused port on localhost Return type: [int](https://docs.python.org/3/library/functions.html#int) Example ``` >>> get_free_port() 63976 ``` Note The port returned by `get_free_port()` is not reserved and may be taken by another process after this function returns. torch.distributed.elastic.utils.log_level.get_log_level()[[source]](https://github.com/pytorch/pytorch/blob/v2.13.0/torch/distributed/elastic/utils/log_level.py#L10) Return default log level for pytorch. Return type: [str](https://docs.python.org/3/library/stdtypes.html#str) torch.distributed.elastic.utils.logging.get_logger(*name=None*)[[source]](https://github.com/pytorch/pytorch/blob/v2.13.0/torch/distributed/elastic/utils/logging.py#L17) Util function to set up a simple logger that writes into stderr. The loglevel is fetched from the LOGLEVEL env. variable or WARNING as default. The function will use the module name of the caller if no name is provided. Parameters: **name** ([*str*](https://docs.python.org/3/library/stdtypes.html#str)*|**None*) - Name of the logger. If no name provided, the name will be derived from the call stack. Return type: [*Logger*](https://docs.python.org/3/library/logging.html#logging.Logger) torch.distributed.rendezvous.register_rendezvous_handler(*scheme*, *handler*)[[source]](https://github.com/pytorch/pytorch/blob/v2.13.0/torch/distributed/rendezvous.py#L25) Register a new rendezvous handler. Before we can run collective algorithms, participating processes need to find each other and exchange information to be able to communicate. We call this process rendezvous. The outcome of the rendezvous process is a triplet containing a shared key/value store, the rank of the process, and the total number of participating processes. If none of the bundled rendezvous methods apply to your execution environment you can opt to register your own rendezvous handler. Pick a unique name and use the URL scheme to identify it when calling the rendezvous() function. Parameters: - **scheme** ([*str*](https://docs.python.org/3/library/stdtypes.html#str)) - URL scheme to identify your rendezvous handler. - **handler** (*function*) - Handler that is invoked when the rendezvous() function is called with a URL that uses the corresponding scheme. It must be a generator function that yields the triplet. torch.distributed.algorithms.model_averaging.utils.average_parameters(*params*, *process_group*)[[source]](https://github.com/pytorch/pytorch/blob/v2.13.0/torch/distributed/algorithms/model_averaging/utils.py#L21) Averages all the given parameters. For allreduce efficiency, all the parameters are flattened into a contiguous buffer. Thus, it requires extra memory of the same size as the given parameters. torch.distributed.algorithms.model_averaging.utils.average_parameters_or_parameter_groups(*params*, *process_group*)[[source]](https://github.com/pytorch/pytorch/blob/v2.13.0/torch/distributed/algorithms/model_averaging/utils.py#L81) Averages parameters of a model or parameter groups of an optimizer. torch.distributed.algorithms.model_averaging.utils.get_params_to_average(*params*)[[source]](https://github.com/pytorch/pytorch/blob/v2.13.0/torch/distributed/algorithms/model_averaging/utils.py#L52) Return a list of parameters that need to average. This filters out the parameters that do not contain any gradients. :param params: The parameters of a model or parameter groups of an optimizer. ## Backends `torch.distributed` supports four built-in backends, each with different capabilities. The table below shows which functions are available for use with a CPU or GPU for each backend. For NCCL, GPU refers to CUDA GPU while for XCCL to XPU GPU. MPI supports CUDA only if the implementation used to build PyTorch supports it. | Backend | `gloo` | `mpi` | `nccl` | `xccl` | | --- | --- | --- | --- | --- | | Device | CPU | GPU | CPU | GPU | | send | ✓ | ✘ | ✓ | ? | | recv | ✓ | ✘ | ✓ | ? | | broadcast | ✓ | ✓ | ✓ | ? | | all_reduce | ✓ | ✓ | ✓ | ? | | reduce | ✓ | ✓ | ✓ | ? | | all_gather | ✓ | ✓ | ✓ | ? | | gather | ✓ | ✓ | ✓ | ? | | scatter | ✓ | ✓ | ✓ | ? | | reduce_scatter | ✓ | ✓ | ✘ | ✘ | | all_to_all | ✘ | ✘ | ✓ | ? | | barrier | ✓ | ✘ | ✓ | ? | ### Backends that come with PyTorch PyTorch distributed package supports Linux (stable), macOS (stable), and Windows (prototype). By default for Linux, the Gloo and NCCL backends are built and included in PyTorch distributed (NCCL only when building with CUDA). MPI is an optional backend that can only be included if you build PyTorch from source. (e.g. building PyTorch on a host that has MPI installed.) Note As of PyTorch v1.8, Windows supports all collective communications backends but NCCL, If the `init_method` argument of `init_process_group()` points to a file it must adhere to the following schema: - Local file system, `init_method="file:///d:/tmp/some_file"` - Shared file system, `init_method="file://////{machine_name}/{share_folder_name}/some_file"` Same as on Linux platform, you can enable TcpStore by setting environment variables, MASTER_ADDR and MASTER_PORT. ### Which backend to use? In the past, we were often asked: "which backend should I use?". - Rule of thumb - Use the NCCL backend for distributed training with CUDA **GPU**. - Use the XCCL backend for distributed training with XPU **GPU**. - Use the Gloo backend for distributed training with **CPU**. - GPU hosts with InfiniBand interconnect - Use NCCL, since it's the only backend that currently supports InfiniBand and GPUDirect. - GPU hosts with Ethernet interconnect - Use NCCL, since it currently provides the best distributed GPU training performance, especially for multiprocess single-node or multi-node distributed training. If you encounter any problem with NCCL, use Gloo as the fallback option. (Note that Gloo currently runs slower than NCCL for GPUs.) - CPU hosts with InfiniBand interconnect - If your InfiniBand has enabled IP over IB, use Gloo, otherwise, use MPI instead. We are planning on adding InfiniBand support for Gloo in the upcoming releases. - CPU hosts with Ethernet interconnect - Use Gloo, unless you have specific reasons to use MPI. ### Common environment variables #### Choosing the network interface to use By default, both the NCCL and Gloo backends will try to find the right network interface to use. If the automatically detected interface is not correct, you can override it using the following environment variables (applicable to the respective backend): - **NCCL_SOCKET_IFNAME**, for example `export NCCL_SOCKET_IFNAME=eth0` - **GLOO_SOCKET_IFNAME**, for example `export GLOO_SOCKET_IFNAME=eth0` If you're using the Gloo backend, you can specify multiple interfaces by separating them by a comma, like this: `export GLOO_SOCKET_IFNAME=eth0,eth1,eth2,eth3`. The backend will dispatch operations in a round-robin fashion across these interfaces. It is imperative that all processes specify the same number of interfaces in this variable. #### Other NCCL environment variables **Debugging** - in case of NCCL failure, you can set `NCCL_DEBUG=INFO` to print an explicit warning message as well as basic NCCL initialization information. You may also use `NCCL_DEBUG_SUBSYS` to get more details about a specific aspect of NCCL. For example, `NCCL_DEBUG_SUBSYS=COLL` would print logs of collective calls, which may be helpful when debugging hangs, especially those caused by collective type or message size mismatch. In case of topology detection failure, it would be helpful to set `NCCL_DEBUG_SUBSYS=GRAPH` to inspect the detailed detection result and save as reference if further help from NCCL team is needed. **Performance tuning** - NCCL performs automatic tuning based on its topology detection to save users' tuning effort. On some socket-based systems, users may still try tuning `NCCL_SOCKET_NTHREADS` and `NCCL_NSOCKS_PERTHREAD` to increase socket network bandwidth. These two environment variables have been pre-tuned by NCCL for some cloud providers, such as AWS or GCP. For a full list of NCCL environment variables, please refer to [NVIDIA NCCL's official documentation](https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/env.html) You can tune NCCL communicators even further using `torch.distributed.ProcessGroupNCCL.NCCLConfig` and `torch.distributed.ProcessGroupNCCL.Options`. Learn more about them using `help` (e.g. `help(torch.distributed.ProcessGroupNCCL.NCCLConfig)`) in the interpreter. ## Basics The `torch.distributed` package provides PyTorch support and communication primitives for multiprocess parallelism across several computation nodes running on one or more machines. The class [`torch.nn.parallel.DistributedDataParallel()`](generated/torch.nn.parallel.DistributedDataParallel.html#torch.nn.parallel.DistributedDataParallel) builds on this functionality to provide synchronous distributed training as a wrapper around any PyTorch model. This differs from the kinds of parallelism provided by [Multiprocessing package - torch.multiprocessing](multiprocessing.html) and [`torch.nn.DataParallel()`](generated/torch.nn.DataParallel.html#torch.nn.DataParallel) in that it supports multiple network-connected machines and in that the user must explicitly launch a separate copy of the main training script for each process. In the single-machine synchronous case, `torch.distributed` or the [`torch.nn.parallel.DistributedDataParallel()`](generated/torch.nn.parallel.DistributedDataParallel.html#torch.nn.parallel.DistributedDataParallel) wrapper may still have advantages over other approaches to data-parallelism, including [`torch.nn.DataParallel()`](generated/torch.nn.DataParallel.html#torch.nn.DataParallel): - Each process maintains its own optimizer and performs a complete optimization step with each iteration. While this may appear redundant, since the gradients have already been gathered together and averaged across processes and are thus the same for every process, this means that no parameter broadcast step is needed, reducing time spent transferring tensors between nodes. - Each process contains an independent Python interpreter, eliminating the extra interpreter overhead and "GIL-thrashing" that comes from driving several execution threads, model replicas, or GPUs from a single Python process. This is especially important for models that make heavy use of the Python runtime, including models with recurrent layers or many small components. ## Initialization The package needs to be initialized using the `torch.distributed.init_process_group()` or `torch.distributed.device_mesh.init_device_mesh()` function before calling any other methods. Both block until all processes have joined. Warning Initialization is not thread-safe. Process group creation should be performed from a single thread, to prevent inconsistent 'UUID' assignment across ranks, and to prevent races during initialization that can lead to hangs. torch.distributed.is_available()[[source]](https://github.com/pytorch/pytorch/blob/v2.13.0/torch/distributed/__init__.py#L17) Return `True` if the distributed package is available. Otherwise, `torch.distributed` does not expose any other APIs. Currently, `torch.distributed` is available on Linux, MacOS and Windows. Set `USE_DISTRIBUTED=1` to enable it when building PyTorch from source. Currently, the default value is `USE_DISTRIBUTED=1` for Linux and Windows, `USE_DISTRIBUTED=0` for MacOS. Return type: [bool](https://docs.python.org/3/library/functions.html#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*, *_ranks=None*)[[source]](https://github.com/pytorch/pytorch/blob/v2.13.0/torch/distributed/distributed_c10d.py#L1664) Initialize the default distributed process group. This will also initialize the distributed package. There are 2 main ways to initialize a process group: 1. Specify `store`, `rank`, and `world_size` explicitly. 2. Specify `init_method` (a URL string) which indicates where/how to discover peers. Optionally specify `rank` and `world_size`, or encode all required parameters in the URL and omit them. If neither is specified, `init_method` is assumed to be "env://". Parameters: - **backend** ([*str*](https://docs.python.org/3/library/stdtypes.html#str)*or**Backend**,**optional*) - The backend to use. Depending on build-time configurations, valid values include `mpi`, `gloo`, `nccl`, `ucc`, `xccl` or one that is registered by a third-party plugin. Since 2.6, if `backend` is not provided, c10d will use a backend registered for the device type indicated by the device_id kwarg (if provided). The known default registrations today are: `nccl` for `cuda`, `gloo` for `cpu`, `xccl` for `xpu`. If neither `backend` nor `device_id` is provided, c10d will detect the accelerator on the run-time machine and use a backend registered for that detected accelerator (or `cpu`). This field can be given as a lowercase string (e.g., `"gloo"`), which can also be accessed via `Backend` attributes (e.g., `Backend.GLOO`). If using multiple processes per machine with `nccl` backend, each process must have exclusive access to every GPU it uses, as sharing GPUs between processes can result in deadlock or NCCL invalid usage. `ucc` backend is experimental. Default backend for the device can be queried with `get_default_backend_for_device()`. - **init_method** ([*str*](https://docs.python.org/3/library/stdtypes.html#str)*,**optional*) - URL specifying how to initialize the process group. Default is "env://" if no `init_method` or `store` is specified. Mutually exclusive with `store`. - **world_size** ([*int*](https://docs.python.org/3/library/functions.html#int)*,**optional*) - Number of processes participating in the job. Required if `store` is specified. - **rank** ([*int*](https://docs.python.org/3/library/functions.html#int)*,**optional*) - Rank of the current process (it should be a number between 0 and `world_size`-1). Required if `store` is specified. - **store** (*Store**,**optional*) - Key/value store accessible to all workers, used to exchange connection/address information. Mutually exclusive with `init_method`. - **timeout** (*timedelta**,**optional*) - Timeout for operations executed against the process group. Default value is 10 minutes for NCCL and 30 minutes for other backends. This is the duration after which collectives will be aborted asynchronously and the process will crash. This is done since CUDA execution is async and it is no longer safe to continue executing user code since failed async NCCL operations might result in subsequent CUDA operations running on corrupted data. When TORCH_NCCL_BLOCKING_WAIT is set, the process will block and wait for this timeout. - **group_name** ([*str*](https://docs.python.org/3/library/stdtypes.html#str)*,**optional**,**deprecated*) - Group name. This argument is ignored - **pg_options** (*ProcessGroupOptions**,**optional*) - process group options specifying what additional options need to be passed in during the construction of specific process groups. As of now, the only options we support is `ProcessGroupNCCL.Options` for the `nccl` backend, `is_high_priority_stream` can be specified so that the nccl backend can pick up high priority cuda streams when there're compute kernels waiting. For other available options to config nccl, See [https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/api/types.html#ncclconfig-t](https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/api/types.html#ncclconfig-t) - **device_id** ([*torch.device*](tensor_attributes.html#torch.device)*|*[*int*](https://docs.python.org/3/library/functions.html#int)*,**optional*) - a single, specific device this process will work on, allowing for backend-specific optimizations. Currently this has two effects, only under NCCL: the communicator is immediately formed (calling `ncclCommInit*` immediately rather than the normal lazy call) and sub-groups will use `ncclCommSplit` when possible to avoid unnecessary overhead of group creation. If you want to know NCCL initialization error early, you can also use this field. If an int is provided, the API assumes that the accelerator type at compile time will be used. - **_ranks** ([*list*](https://docs.python.org/3/library/stdtypes.html#list)*[*[*int*](https://docs.python.org/3/library/functions.html#int)*]**|**None*) - The ranks in the process group. If provided, the process group name will be the hash of all the ranks in the group. Note To enable `backend == Backend.MPI`, PyTorch needs to be built from source on a system that supports MPI. Note Support for multiple backends is experimental. Currently when no backend is specified, both `gloo` and `nccl` backends will be created. The `gloo` backend will be used for collectives with CPU tensors and the `nccl` backend will be used for collectives with CUDA tensors. A custom backend can be specified by passing in a string with format ":,:", e.g. "cpu:gloo,cuda:custom_backend". torch.distributed.device_mesh.init_device_mesh(*device_type*, *mesh_shape*, ***, *mesh_dim_names=None*, *backend_override=None*)[[source]](https://github.com/pytorch/pytorch/blob/v2.13.0/torch/distributed/device_mesh.py#L1498) Initializes a DeviceMesh based on device_type, mesh_shape, and mesh_dim_names parameters. This creates a DeviceMesh with an n-dimensional array layout, where n is the length of mesh_shape. If mesh_dim_names is provided, each dimension is labeled as mesh_dim_names[i]. Note init_device_mesh follows SPMD programming model, meaning the same PyTorch Python program runs on all processes/ranks in the cluster. Ensure mesh_shape (the dimensions of the nD array describing device layout) is identical across all ranks. Inconsistent mesh_shape may lead to hanging. Note If no process group is found, init_device_mesh will initialize distributed process group/groups required for distributed communications behind the scene. Parameters: - **device_type** ([*str*](https://docs.python.org/3/library/stdtypes.html#str)) - The device type of the mesh. Currently supports: "cpu", "cuda/cuda-like", "xpu". Passing in a device type with a GPU index, such as "cuda:0", is not allowed. - **mesh_shape** (*Tuple**[*[*int*](https://docs.python.org/3/library/functions.html#int)*]*) - A tuple defining the dimensions of the multi-dimensional array describing the layout of devices. - **mesh_dim_names** ([*tuple*](https://docs.python.org/3/library/stdtypes.html#tuple)*[*[*str*](https://docs.python.org/3/library/stdtypes.html#str)*,**...**]**,**optional*) - A tuple of mesh dimension names to assign to each dimension of the multi-dimensional array describing the layout of devices. Its length must match the length of mesh_shape. Each string in mesh_dim_names must be unique. - **backend_override** (*Dict**[*[*int*](https://docs.python.org/3/library/functions.html#int)*|*[*str*](https://docs.python.org/3/library/stdtypes.html#str)*,*[*tuple*](https://docs.python.org/3/library/stdtypes.html#tuple)*[*[*str*](https://docs.python.org/3/library/stdtypes.html#str)*,**Options**]**|*[*str*](https://docs.python.org/3/library/stdtypes.html#str)*|**Options**]**,**optional*) - Overrides for some or all of the ProcessGroups that will be created for each mesh dimension. Each key can be either the index of a dimension or its name (if mesh_dim_names is provided). Each value can be a tuple containing the name of the backend and its options, or just one of these two components (in which case the other will be set to its default value). Returns: A `DeviceMesh` object representing the device layout. Return type: DeviceMesh 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")) ``` torch.distributed.is_initialized()[[source]](https://github.com/pytorch/pytorch/blob/v2.13.0/torch/distributed/distributed_c10d.py#L1394) Check if the default process group has been initialized. Return type: [bool](https://docs.python.org/3/library/functions.html#bool) torch.distributed.is_mpi_available()[[source]](https://github.com/pytorch/pytorch/blob/v2.13.0/torch/distributed/distributed_c10d.py#L1332) Check if the MPI backend is available. Return type: [bool](https://docs.python.org/3/library/functions.html#bool) torch.distributed.is_nccl_available()[[source]](https://github.com/pytorch/pytorch/blob/v2.13.0/torch/distributed/distributed_c10d.py#L1337) Check if the NCCL backend is available. Return type: [bool](https://docs.python.org/3/library/functions.html#bool) torch.distributed.is_gloo_available()[[source]](https://github.com/pytorch/pytorch/blob/v2.13.0/torch/distributed/distributed_c10d.py#L1342) Check if the Gloo backend is available. Return type: [bool](https://docs.python.org/3/library/functions.html#bool) torch.distributed.distributed_c10d.is_xccl_available()[[source]](https://github.com/pytorch/pytorch/blob/v2.13.0/torch/distributed/distributed_c10d.py#L1352) Check if the XCCL backend is available. Return type: [bool](https://docs.python.org/3/library/functions.html#bool) torch.distributed.distributed_c10d.batch_isend_irecv(*p2p_op_list*)[[source]](https://github.com/pytorch/pytorch/blob/v2.13.0/torch/distributed/distributed_c10d.py#L2990) Send or Receive a batch of tensors asynchronously and return a list of requests. Process each of the operations in `p2p_op_list` and return the corresponding requests. NCCL, Gloo, and UCC backend are currently supported. Parameters: **p2p_op_list** ([*list*](https://docs.python.org/3/library/stdtypes.html#list)*[**P2POp**]*) - A list of point-to-point operations(type of each operator is `torch.distributed.P2POp`). All operations in the list are treated as a single batch - the relative ordering of sends vs. receives in the list does not matter and will not cause deadlocks. For example, `[recv_from_1, send_to_1]` is equivalent to `[send_to_1, recv_from_1]`. The ordering **does** matter for multiple operations between the same pair of ranks in the same direction: if rank 0 calls `[send_to_1(A), send_to_1(B)]`, rank 1 must call `[recv_from_0(A), recv_from_0(B)]` in the same order to ensure the correct tensors are matched. Returns: A list of distributed request objects returned by calling the corresponding op in the op_list. Return type: [list](https://docs.python.org/3/library/stdtypes.html#list)[*Work*] Examples ``` >>> 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]) # Rank 0 tensor([0, 1]) # Rank 1 ``` Note Note that when this API is used with the NCCL PG backend, users must set the current GPU device with torch.cuda.set_device, otherwise it will lead to unexpected hang issues. In addition, if this API is the first collective call in the `group` passed to `dist.P2POp`, all ranks of the `group` must participate in this API call; otherwise, the behavior is undefined. If this API call is not the first collective call in the `group`, batched P2P operations involving only a subset of ranks of the `group` are allowed. torch.distributed.distributed_c10d.destroy_process_group(*group=None*)[[source]](https://github.com/pytorch/pytorch/blob/v2.13.0/torch/distributed/distributed_c10d.py#L2361) Destroy a given process group, and deinitialize the distributed package. Parameters: **group** ([*ProcessGroup*](distributed._dist2.html#torch.distributed._dist2.ProcessGroup)*,**optional*) - The process group to be destroyed, if group.WORLD is given, all process groups including the default one will be destroyed. torch.distributed.distributed_c10d.is_backend_available(*backend*)[[source]](https://github.com/pytorch/pytorch/blob/v2.13.0/torch/distributed/distributed_c10d.py#L1369) Check backend availability. Checks if the given backend is available and supports the built-in backends or third-party backends through function `Backend.register_backend`. Parameters: **backend** ([*str*](https://docs.python.org/3/library/stdtypes.html#str)) - Backend name. Returns: Returns true if the backend is available otherwise false. Return type: [bool](https://docs.python.org/3/library/functions.html#bool) torch.distributed.distributed_c10d.irecv(*tensor*, *src=None*, *group=None*, *tag=0*, *group_src=None*)[[source]](https://github.com/pytorch/pytorch/blob/v2.13.0/torch/distributed/distributed_c10d.py#L2655) Receives a tensor asynchronously. Warning `tag` is not supported with the NCCL backend. Unlike recv, which is blocking, irecv allows src == dst rank, i.e. recv from self. Parameters: - **tensor** ([*Tensor*](tensors.html#torch.Tensor)) - Tensor to fill with received data. - **src** ([*int*](https://docs.python.org/3/library/functions.html#int)*,**optional*) - Source rank on global process group (regardless of `group` argument). Will receive from any process if unspecified. - **group** ([*ProcessGroup*](distributed._dist2.html#torch.distributed._dist2.ProcessGroup)*,**optional*) - The process group to work on. If None, the default process group will be used. - **tag** ([*int*](https://docs.python.org/3/library/functions.html#int)*,**optional*) - Tag to match recv with remote send - **group_src** ([*int*](https://docs.python.org/3/library/functions.html#int)*,**optional*) - Destination rank on `group`. Invalid to specify both `src` and `group_src`. Returns: A distributed request object. None, if not part of the group Return type: *Work* | None torch.distributed.distributed_c10d.is_gloo_available()[[source]](https://github.com/pytorch/pytorch/blob/v2.13.0/torch/distributed/distributed_c10d.py#L1342) Check if the Gloo backend is available. Return type: [bool](https://docs.python.org/3/library/functions.html#bool) torch.distributed.distributed_c10d.is_initialized()[[source]](https://github.com/pytorch/pytorch/blob/v2.13.0/torch/distributed/distributed_c10d.py#L1394) Check if the default process group has been initialized. Return type: [bool](https://docs.python.org/3/library/functions.html#bool) torch.distributed.distributed_c10d.is_mpi_available()[[source]](https://github.com/pytorch/pytorch/blob/v2.13.0/torch/distributed/distributed_c10d.py#L1332) Check if the MPI backend is available. Return type: [bool](https://docs.python.org/3/library/functions.html#bool) torch.distributed.distributed_c10d.is_nccl_available()[[source]](https://github.com/pytorch/pytorch/blob/v2.13.0/torch/distributed/distributed_c10d.py#L1337) Check if the NCCL backend is available. Return type: [bool](https://docs.python.org/3/library/functions.html#bool) torch.distributed.distributed_c10d.is_torchelastic_launched()[[source]](https://github.com/pytorch/pytorch/blob/v2.13.0/torch/distributed/distributed_c10d.py#L1399) Check whether this process was launched with `torch.distributed.elastic` (aka torchelastic). The existence of `TORCHELASTIC_RUN_ID` environment variable is used as a proxy to determine whether the current process was launched with torchelastic. This is a reasonable proxy since `TORCHELASTIC_RUN_ID` maps to the rendezvous id which is always a non-null value indicating the job id for peer discovery purposes.. Return type: [bool](https://docs.python.org/3/library/functions.html#bool) torch.distributed.distributed_c10d.is_ucc_available()[[source]](https://github.com/pytorch/pytorch/blob/v2.13.0/torch/distributed/distributed_c10d.py#L1347) Check if the UCC backend is available. Return type: [bool](https://docs.python.org/3/library/functions.html#bool) torch.distributed.is_torchelastic_launched()[[source]](https://github.com/pytorch/pytorch/blob/v2.13.0/torch/distributed/distributed_c10d.py#L1399) Check whether this process was launched with `torch.distributed.elastic` (aka torchelastic). The existence of `TORCHELASTIC_RUN_ID` environment variable is used as a proxy to determine whether the current process was launched with torchelastic. This is a reasonable proxy since `TORCHELASTIC_RUN_ID` maps to the rendezvous id which is always a non-null value indicating the job id for peer discovery purposes.. Return type: [bool](https://docs.python.org/3/library/functions.html#bool) torch.distributed.get_default_backend_for_device(*device*)[[source]](https://github.com/pytorch/pytorch/blob/v2.13.0/torch/distributed/distributed_c10d.py#L1497) Return the default backend for the given device. Parameters: **device** (*Union**[*[*str*](https://docs.python.org/3/library/stdtypes.html#str)*,*[*torch.device*](tensor_attributes.html#torch.device)*]*) - The device to get the default backend for. Returns: The default backend for the given device as a lower case string. Return type: [str](https://docs.python.org/3/library/stdtypes.html#str) --- Currently three initialization methods are supported: ### TCP initialization There are two ways to initialize using TCP, both requiring a network address reachable from all processes and a desired `world_size`. The first way requires specifying an address that belongs to the rank 0 process. This initialization method requires that all processes have manually specified ranks. Note that multicast address is not supported anymore in the latest distributed package. `group_name` is deprecated as well. ``` import torch.distributed as dist # Use address of one of the machines dist.init_process_group(backend, init_method='tcp://10.1.1.20:23456', rank=args.rank, world_size=4) ``` ### Shared file-system initialization Another initialization method makes use of a file system that is shared and visible from all machines in a group, along with a desired `world_size`. The URL should start with `file://` and contain a path to a non-existent file (in an existing directory) on a shared file system. File-system initialization will automatically create that file if it doesn't exist, but will not delete the file. Therefore, it is your responsibility to make sure that the file is cleaned up before the next `init_process_group()` call on the same file path/name. Note that automatic rank assignment is not supported anymore in the latest distributed package and `group_name` is deprecated as well. Warning This method assumes that the file system supports locking using `fcntl` - most local systems and NFS support it. Warning This method will always create the file and try its best to clean up and remove the file at the end of the program. In other words, each initialization with the file init method will need a brand new empty file in order for the initialization to succeed. If the same file used by the previous initialization (which happens not to get cleaned up) is used again, this is unexpected behavior and can often cause deadlocks and failures. Therefore, even though this method will try its best to clean up the file, if the auto-delete happens to be unsuccessful, it is your responsibility to ensure that the file is removed at the end of the training to prevent the same file to be reused again during the next time. This is especially important if you plan to call `init_process_group()` multiple times on the same file name. In other words, if the file is not removed/cleaned up and you call `init_process_group()` again on that file, failures are expected. The rule of thumb here is that, make sure that the file is non-existent or empty every time `init_process_group()` is called. ``` import torch.distributed as dist # rank should always be specified dist.init_process_group(backend, init_method='file:///mnt/nfs/sharedfile', world_size=4, rank=args.rank) ``` ### Environment variable initialization This method will read the configuration from environment variables, allowing one to fully customize how the information is obtained. The variables to be set are: - `MASTER_PORT` - required; has to be a free port on machine with rank 0 - `MASTER_ADDR` - required (except for rank 0); address of rank 0 node - `WORLD_SIZE` - required; can be set either here, or in a call to init function - `RANK` - required; can be set either here, or in a call to init function The machine with rank 0 will be used to set up all connections. This is the default method, meaning that `init_method` does not have to be specified (or can be `env://`). ### Improving initialization time - `TORCH_GLOO_LAZY_INIT` - establishes connections on demand rather than using a full mesh which can greatly improve initialization time for non all2all operations. ## Post-Initialization Once `torch.distributed.init_process_group()` was run, the following functions can be used. To check whether the process group has already been initialized use `torch.distributed.is_initialized()`. *class*torch.distributed.Backend(*name*)[[source]](https://github.com/pytorch/pytorch/blob/v2.13.0/torch/distributed/distributed_c10d.py#L270) An enum-like class for backends. Available backends: GLOO, NCCL, UCC, MPI, XCCL, FAKE, and other registered backends. The values of this class are lowercase strings, e.g., `"gloo"`. They can be accessed as attributes, e.g., `Backend.NCCL`. This class can be directly called to parse the string, e.g., `Backend(backend_str)` will check if `backend_str` is valid, and return the parsed lowercase string if so. It also accepts uppercase strings, e.g., `Backend("GLOO")` returns `"gloo"`. Note The entry `Backend.UNDEFINED` is present but only used as initial value of some fields. Users should neither use it directly nor assume its existence. *classmethod*register_backend(*name*, *func*, *extended_api=False*, *devices=None*)[[source]](https://github.com/pytorch/pytorch/blob/v2.13.0/torch/distributed/distributed_c10d.py#L340) Register a new backend with the given name and instantiating function. This class method is used by 3rd party `ProcessGroup` extension to register new backends. Parameters: - **name** ([*str*](https://docs.python.org/3/library/stdtypes.html#str)) - Backend name of the `ProcessGroup` extension. It should match the one in `init_process_group()`. - **func** (*function*) - Function handler that instantiates the backend. The function should be implemented in the backend extension and takes four arguments, including `store`, `rank`, `world_size`, and `timeout`. - **extended_api** ([*bool*](https://docs.python.org/3/library/functions.html#bool)*,**optional*) - Whether the backend supports extended argument structure. Default: `False`. If set to `True`, the backend will get an instance of `c10d::DistributedBackendOptions`, and a process group options object as defined by the backend implementation. - **device** ([*str*](https://docs.python.org/3/library/stdtypes.html#str)*or*[*list*](https://docs.python.org/3/library/stdtypes.html#list)*of*[*str*](https://docs.python.org/3/library/stdtypes.html#str)*,**optional*) - device type this backend supports, e.g. "cpu", "cuda", etc. If None, assuming both "cpu" and "cuda" Note This support of 3rd party backend is experimental and subject to change. torch.distributed.get_backend(*group=None*)[[source]](https://github.com/pytorch/pytorch/blob/v2.13.0/torch/distributed/distributed_c10d.py#L1471) Return the backend of the given process group. Parameters: **group** ([*ProcessGroup*](distributed._dist2.html#torch.distributed._dist2.ProcessGroup)*,**optional*) - The process group to work on. The default is the general main process group. If another specific group is specified, the calling process must be part of `group`. Returns: The backend of the given process group as a lower case string. Return type: *Backend* torch.distributed.get_backend_config(*group=None*)[[source]](https://github.com/pytorch/pytorch/blob/v2.13.0/torch/distributed/distributed_c10d.py#L1451) Return the backend configuration of the given process group. Parameters: **group** ([*ProcessGroup*](distributed._dist2.html#torch.distributed._dist2.ProcessGroup)*,**optional*) - The process group to work on. The default is the general main process group. If another specific group is specified, the calling process must be part of `group`. Returns: The backend configuration of the given process group as a lower case string. Return type: [str](https://docs.python.org/3/library/stdtypes.html#str) torch.distributed.get_rank(*group=None*)[[source]](https://github.com/pytorch/pytorch/blob/v2.13.0/torch/distributed/distributed_c10d.py#L2552) Return the rank of the current process in the provided `group`, default otherwise. Rank is a unique identifier assigned to each process within a distributed process group. They are always consecutive integers ranging from 0 to `world_size`. Parameters: **group** ([*ProcessGroup*](distributed._dist2.html#torch.distributed._dist2.ProcessGroup)*,**optional*) - The process group to work on. If None, the default process group will be used. Returns: The rank of the process group -1, if not part of the group Return type: [int](https://docs.python.org/3/library/functions.html#int) torch.distributed.get_world_size(*group=None*)[[source]](https://github.com/pytorch/pytorch/blob/v2.13.0/torch/distributed/distributed_c10d.py#L2579) Return the number of processes in the current process group. Parameters: **group** ([*ProcessGroup*](distributed._dist2.html#torch.distributed._dist2.ProcessGroup)*,**optional*) - The process group to work on. If None, the default process group will be used. Returns: The world size of the process group -1, if not part of the group Return type: [int](https://docs.python.org/3/library/functions.html#int) torch.distributed.get_debug_level() → torch._C._distributed_c10d.DebugLevel Gets the debug level of the torch.distributed package. torch.distributed.get_node_local_rank(*fallback_rank=None*)[[source]](https://github.com/pytorch/pytorch/blob/v2.13.0/torch/distributed/distributed_c10d.py#L1563) Return the local rank of the current process relative to the node. Semantically, this is a useful concept for mapping processes to devices. For example, on a node with 8 accelerator you could use the node local rank to decide which accelerator device to bind the process to. In practice, the actual assignment of node local ranks is handled by the process launcher outside of pytorch, and communicated via the LOCAL_RANK environment variable. Torchrun will automatically populate LOCAL_RANK, but other launchers may not. If LOCAL_RANK is unspecified, this API will fall back to the provided kwarg 'fallback_rank' if specified, otherwise it will raise an error. The intent is to allow writing an application that runs either in single or multi device contexts without error. Return type: [int](https://docs.python.org/3/library/functions.html#int) torch.distributed.get_pg_count()[[source]](https://github.com/pytorch/pytorch/blob/v2.13.0/torch/distributed/distributed_c10d.py#L1555) Return the number of process groups. Return type: [int](https://docs.python.org/3/library/functions.html#int) ## Shutdown It is important to clean up resources on exit by calling `destroy_process_group()`. The simplest pattern to follow is to destroy every process group and backend by calling `destroy_process_group()` with the default value of None for the `group` argument, at a point in the training script where communications are no longer needed, usually near the end of main(). The call should be made once per trainer-process, not at the outer process-launcher level. if `destroy_process_group()` is not called by all ranks in a pg within the timeout duration, especially when there are multiple process-groups in the application e.g. for N-D parallelism, hangs on exit are possible. This is because the destructor for ProcessGroupNCCL calls ncclCommAbort, which must be called collectively, but the order of calling ProcessGroupNCCL's destructor if called by python's GC is not deterministic. Calling `destroy_process_group()` helps by ensuring ncclCommAbort is called in a consistent order across ranks, and avoids calling ncclCommAbort during ProcessGroupNCCL's destructor. ### Reinitialization `destroy_process_group` can also be used to destroy individual process groups. One use case could be fault tolerant training, where a process group may be destroyed and then a new one initialized during runtime. In this case, it's critical to synchronize the trainer processes using some means other than torch.distributed primitives _after_ calling destroy and before subsequently initializing. This behavior is currently unsupported/untested, due to the difficulty of achieving this synchronization, and is considered a known issue. Please file a github issue or RFC if this is a use case that's blocking you. --- ## Groups By default collectives operate on the default group (also called the world) and require all processes to enter the distributed function call. However, some workloads can benefit from more fine-grained communication. This is where distributed groups come into play. `new_group()` function can be used to create new groups, with arbitrary subsets of all processes. It returns an opaque group handle that can be given as a `group` argument to all collectives (collectives are distributed functions to exchange information in certain well-known programming patterns). torch.distributed.new_group(*ranks=None*, *timeout=None*, *backend=None*, *pg_options=None*, *use_local_synchronization=False*, *group_desc=None*, *device_id=None*, *sort_ranks=True*)[[source]](https://github.com/pytorch/pytorch/blob/v2.13.0/torch/distributed/distributed_c10d.py#L5744) Create a new distributed group. This function requires that all processes in the main group (i.e. all processes that are part of the distributed job) enter this function, even if they are not going to be members of the group. Additionally, groups should be created in the same order in all processes. Warning Safe concurrent usage: When using multiple process groups with the `NCCL` backend, the user must ensure a globally consistent execution order of collectives across ranks. If multiple threads within a process issue collectives, explicit synchronization is necessary to ensure consistent ordering. When using async variants of torch.distributed communication APIs, a work object is returned and the communication kernel is enqueued on a separate CUDA stream, allowing overlap of communication and computation. Once one or more async ops have been issued on one process group, they must be synchronized with other cuda streams by calling work.wait() before using another process group. See Using multiple NCCL communicators concurrently for more details. Parameters: - **ranks** ([*list*](https://docs.python.org/3/library/stdtypes.html#list)*[*[*int*](https://docs.python.org/3/library/functions.html#int)*]*) - List of ranks of group members. If `None`, will be set to all ranks. Default is `None`. - **timeout** (*timedelta**,**optional*) - see init_process_group for details and default value. - **backend** ([*str*](https://docs.python.org/3/library/stdtypes.html#str)*or**Backend**,**optional*) - The backend to use. Depending on build-time configurations, valid values are `gloo` and `nccl`. By default uses the same backend as the global group. This field should be given as a lowercase string (e.g., `"gloo"`), which can also be accessed via `Backend` attributes (e.g., `Backend.GLOO`). If `None` is passed in, the backend corresponding to the default process group will be used. Default is `None`. - **pg_options** (*ProcessGroupOptions**,**optional*) - process group options specifying what additional options need to be passed in during the construction of specific process groups. i.e. for the `nccl` backend, `is_high_priority_stream` can be specified so that process group can pick up high priority cuda streams. For other available options to config nccl, See [https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/api/types.html#ncclconfig-tuse_local_synchronization](https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/api/types.html#ncclconfig-tuse_local_synchronization) (bool, optional): perform a group-local barrier at the end of the process group creation. This is different in that non-member ranks don't need to call into API and don't join the barrier. - **group_desc** ([*str*](https://docs.python.org/3/library/stdtypes.html#str)*,**optional*) - a string to describe the process group. - **device_id** ([*torch.device*](tensor_attributes.html#torch.device)*,**optional*) - a single, specific device to "bind" this process to, The new_group call will try to initialize a communication backend immediately for the device if this field is given. - **sort_ranks** ([*bool*](https://docs.python.org/3/library/functions.html#bool)*,**optional*) - If `True` (the default), sort the `ranks` list before creating the group. If `False`, preserve the caller-provided rank order so that position in the `ranks` list determines group rank. All processes must pass the identical `ranks` list. Default is `True`. Returns: A handle of distributed group that can be given to collective calls or GroupMember.NON_GROUP_MEMBER if the rank is not part of `ranks`. N.B. use_local_synchronization doesn't work with MPI. N.B. While use_local_synchronization=True can be significantly faster with larger clusters and small process groups, care must be taken since it changes cluster behavior as non-member ranks don't join the group barrier(). N.B. use_local_synchronization=True can lead to deadlocks when each rank creates multiple overlapping process groups. To avoid that, make sure all ranks follow the same global creation order. N.B. When TorchComms is enabled (`torch.distributed.config.use_torchcomms` / `TORCH_DISTRIBUTED_USE_TORCHCOMMS=1`), this function delegates to `split_group()` so subgroup creation goes through the TorchComms path. The delegation raises `NotImplementedError` for arguments that `split_group()` cannot honor (e.g. `use_local_synchronization=True`, `sort_ranks=False`, or an explicit `device_id` that diverges from the default group's bound device). torch.distributed.distributed_c10d.shrink_group(*ranks_to_exclude*, *group=None*, *shrink_flags=0*, *pg_options=None*)[[source]](https://github.com/pytorch/pytorch/blob/v2.13.0/torch/distributed/distributed_c10d.py#L6367) Shrinks a process group by excluding specified ranks. Creates and returns a new, smaller process group comprising only the ranks from the original group that were not in the `ranks_to_exclude` list. Parameters: - **ranks_to_exclude** (*List**[*[*int*](https://docs.python.org/3/library/functions.html#int)*]*) - A list of ranks from the original `group` to exclude from the new group. - **group** ([*ProcessGroup*](distributed._dist2.html#torch.distributed._dist2.ProcessGroup)*,**optional*) - The process group to shrink. If `None`, the default process group is used. Defaults to `None`. - **shrink_flags** ([*int*](https://docs.python.org/3/library/functions.html#int)*,**optional*) - Flags to control the shrinking behavior. Can be `SHRINK_DEFAULT` (default) or `SHRINK_ABORT`. `SHRINK_ABORT` will attempt to terminate ongoing operations in the parent communicator before shrinking. Defaults to `SHRINK_DEFAULT`. - **pg_options** (*ProcessGroupOptions**,**optional*) - Backend-specific options to apply to the shrunken process group. If provided, the backend will use these options when creating the new group. If omitted, the new group inherits defaults from the parent. Returns: a new group comprised of the remaining ranks. If the default group was shrunk, the returned group becomes the new default group. Return type: [ProcessGroup](distributed._dist2.html#torch.distributed._dist2.ProcessGroup) Raises: - [**TypeError**](https://docs.python.org/3/library/exceptions.html#TypeError) - if the group's backend does not support shrinking. - [**ValueError**](https://docs.python.org/3/library/exceptions.html#ValueError) - if `ranks_to_exclude` is invalid (empty, out of bounds, - **duplicates****, or****excludes all ranks****)****.** - - [**RuntimeError**](https://docs.python.org/3/library/exceptions.html#RuntimeError) - if an excluded rank calls this function or the backend - **fails the operation.** - Notes - Only non-excluded ranks should call this function; excluded ranks must not participate in the shrink operation. - Shrinking the default group destroys all other process groups since rank reassignment makes them inconsistent. torch.distributed.get_group_rank(*group*, *global_rank*)[[source]](https://github.com/pytorch/pytorch/blob/v2.13.0/torch/distributed/distributed_c10d.py#L1111) Translate a global rank into a group rank. `global_rank` must be part of `group` otherwise this raises RuntimeError. Parameters: - **group** ([*ProcessGroup*](distributed._dist2.html#torch.distributed._dist2.ProcessGroup)) - ProcessGroup to find the relative rank. - **global_rank** ([*int*](https://docs.python.org/3/library/functions.html#int)) - Global rank to query. Returns: Group rank of `global_rank` relative to `group` Return type: [int](https://docs.python.org/3/library/functions.html#int) N.B. calling this function on the default process group returns identity torch.distributed.get_global_rank(*group*, *group_rank*)[[source]](https://github.com/pytorch/pytorch/blob/v2.13.0/torch/distributed/distributed_c10d.py#L1139) Translate a group rank into a global rank. `group_rank` must be part of group otherwise this raises RuntimeError. Parameters: - **group** ([*ProcessGroup*](distributed._dist2.html#torch.distributed._dist2.ProcessGroup)) - ProcessGroup to find the global rank from. - **group_rank** ([*int*](https://docs.python.org/3/library/functions.html#int)) - Group rank to query. Returns: Global rank of `group_rank` relative to `group` Return type: [int](https://docs.python.org/3/library/functions.html#int) N.B. calling this function on the default process group returns identity torch.distributed.get_process_group_ranks(*group*)[[source]](https://github.com/pytorch/pytorch/blob/v2.13.0/torch/distributed/distributed_c10d.py#L1177) Get all ranks associated with `group`. Parameters: **group** (*Optional**[*[*ProcessGroup*](distributed._dist2.html#torch.distributed._dist2.ProcessGroup)*]*) - ProcessGroup to get all ranks from. If None, the default process group will be used. Returns: List of global ranks ordered by group rank. Return type: [list](https://docs.python.org/3/library/stdtypes.html#list)[[int](https://docs.python.org/3/library/functions.html#int)] torch.distributed.split_group(*parent_pg=None*, *split_ranks=None*, *timeout=None*, *pg_options=None*, *group_desc=None*, *backend=None*)[[source]](https://github.com/pytorch/pytorch/blob/v2.13.0/torch/distributed/distributed_c10d.py#L5516) Create a new process group split from the given parent process group. warning:: This is an experimental API. Only the `NCCL` and custom plugin backends are supported. Other backends will raise an error. Users of this API must guarantee that all ranks in the parent group enter this API call, and the split of the sub groups is the same across all ranks in the parent group. Parameters: - **parent_pg** ([*ProcessGroup*](distributed._dist2.html#torch.distributed._dist2.ProcessGroup)*,**optional*) - The parent process group. If None, the default process group will be used. Users need to guarantee that the parent group is fully initialized (e.g, communicators are initialized) - **split_ranks** ([*list*](https://docs.python.org/3/library/stdtypes.html#list)*[*[*list*](https://docs.python.org/3/library/stdtypes.html#list)*[*[*int*](https://docs.python.org/3/library/functions.html#int)*]**]*) - the split ranks, which is a list of list of ranks. Users need to make sure the validity of the split ranks such that one split (represented by one inner list of ints) does not overlap with any other split. Note that the ranks in each split is the group rank (instead of global rank) in the parent pg. For example, if the parent group has 4 ranks, and split_ranks can be [[0, 1], [2, 3]]. Note [[0,1]] is also a valid split, in which case ranks 2, 3 would return a non-group member. - **timeout** (*timedelta**,**optional*) - see init_process_group for details and default value. - **pg_options** (*ProcessGroupOptions**,**optional*) - Additional options need to be passed in during the construction of specific process groups. i.e.``is_high_priority_stream`` can be specified so that process group can pick up high priority cuda streams. - **group_desc** ([*str*](https://docs.python.org/3/library/stdtypes.html#str)*,**optional*) - a string to describe the process group. - **backend** ([*str*](https://docs.python.org/3/library/stdtypes.html#str)*or**Backend**,**optional*) - selects a subset of the parent process group's per-device backends to retain in the child group. The string follows the same format as `init_process_group` (e.g. `"nccl"` or `"cpu:gloo,cuda:nccl"`). Each device type in the requested backend must be present in the parent group, and the backend name for that device type must exactly match the parent's. The child's default backend's device type must be in the requested set. If `None` (default), the child inherits the full backend configuration of the parent. Returns: ProcessGroup if the current rank is within one split/subgroup given by split_ranks, or None if the current rank is not part of any split_ranks`. Return type: [*ProcessGroup*](distributed._dist2.html#torch.distributed._dist2.ProcessGroup) | None torch.distributed.distributed_c10d.GroupName*= torch.distributed.distributed_c10d.GroupName* NewType creates simple unique types with almost zero runtime overhead. NewType(name, tp) is considered a subtype of tp by static type checkers. At runtime, NewType(name, tp) returns a dummy callable that simply returns its argument. Usage: ``` UserId = NewType('UserId', int) def name_by_id(user_id: UserId) -> str: ... UserId('user') # Fails type check name_by_id(42) # Fails type check name_by_id(UserId(42)) # OK num = UserId(5) + 1 # type: int ``` ## DeviceMesh DeviceMesh is a higher level abstraction that manages process groups (or NCCL communicators). It allows user to easily create inter node and intra node process groups without worrying about how to set up the ranks correctly for different sub process groups, and it helps manage those distributed process group easily. `init_device_mesh()` function can be used to create new DeviceMesh, with a mesh shape describing the device topology. *class*torch.distributed.device_mesh.DeviceMesh(*device_type*, *mesh=None*, ***, *mesh_dim_names=None*, *backend_override=None*, *_init_backend=True*, *_rank=None*, *_layout=None*, *_rank_map=None*, *_root_mesh=None*)[[source]](https://github.com/pytorch/pytorch/blob/v2.13.0/torch/distributed/device_mesh.py#L152) DeviceMesh represents a mesh of devices, where layout of devices could be represented as a n-d dimension array, and each value of the n-d dimensional array is the global id of the default process group ranks. DeviceMesh could be used to setup the N dimensional device connections across the cluster, and manage the ProcessGroups for N dimensional parallelisms. Communications could happen on each dimension of the DeviceMesh separately. DeviceMesh respects the device that user selects already (i.e. if user call torch.cuda.set_device before the DeviceMesh initialization), and will select/set the device for the current process if user does not set the device beforehand. Note that manual device selection should happen BEFORE the DeviceMesh initialization. DeviceMesh can also be used as a context manager when using together with DTensor APIs. Note DeviceMesh follows SPMD programming model, which means the same PyTorch Python program is running on all processes/ranks in the cluster. Therefore, users need to make sure the mesh array (which describes the layout of devices) should be identical across all ranks. Inconsistent mesh will lead to silent hang. Parameters: - **device_type** ([*str*](https://docs.python.org/3/library/stdtypes.html#str)) - The device type of the mesh. Currently supports: "cpu", "cuda/cuda-like". - **mesh** (*ndarray*) - A multi-dimensional array or an integer tensor describing the layout of devices, where the IDs are global IDs of the default process group. - **_rank** ([*int*](https://docs.python.org/3/library/functions.html#int)) - (experimental/internal) The global rank of the current process. If not provided, it will be inferred from the default process group. Returns: A `DeviceMesh` object representing the device layout. Return type: DeviceMesh The following program runs on each process/rank in an SPMD manner. In this example, we have 2 hosts with 4 GPUs each. A reduction over the first dimension of mesh will reduce across columns (0, 4), .. and (3, 7), a reduction over the second dimension of mesh reduces across rows (0, 1, 2, 3) and (4, 5, 6, 7). Example: ``` >>> from torch.distributed.device_mesh import DeviceMesh >>> >>> # Initialize device mesh as (2, 4) to represent the topology >>> # of cross-host(dim 0), and within-host (dim 1). >>> mesh = DeviceMesh(device_type="cuda", mesh=[[0, 1, 2, 3],[4, 5, 6, 7]]) ``` *property*device_type*: [str](https://docs.python.org/3/library/stdtypes.html#str)* Returns the device type of the mesh. *static*from_group(*group*, *device_type*, *mesh=None*, ***, *mesh_dim_names=None*)[[source]](https://github.com/pytorch/pytorch/blob/v2.13.0/torch/distributed/device_mesh.py#L1052) Constructs a `DeviceMesh` with `device_type` from an existing `ProcessGroup` or a list of existing `ProcessGroup`. The constructed device mesh has number of dimensions equal to the number of groups passed. For example, if a single process group is passed in, the resulted DeviceMesh is a 1D mesh. If a list of 2 process groups is passed in, the resulted DeviceMesh is a 2D mesh. If more than one group is passed, then the `mesh` and `mesh_dim_names` arguments are required. The order of the process groups passed in determines the topology of the mesh. For example, the first process group will be the 0th dimension of the DeviceMesh. The mesh tensor passed in must have the same number of dimensions as the number of process groups passed in, and the order of the dimensions in the mesh tensor must match the order in the process groups passed in. Parameters: - **group** ([*ProcessGroup*](distributed._dist2.html#torch.distributed._dist2.ProcessGroup)*or*[*list*](https://docs.python.org/3/library/stdtypes.html#list)*[*[*ProcessGroup*](distributed._dist2.html#torch.distributed._dist2.ProcessGroup)*]*) - the existing ProcessGroup or a list of existing ProcessGroups. - **device_type** ([*str*](https://docs.python.org/3/library/stdtypes.html#str)) - The device type of the mesh. Currently supports: "cpu", "cuda/cuda-like". Passing in a device type with a GPU index, such as "cuda:0", is not allowed. - **mesh** ([*torch.Tensor*](tensors.html#torch.Tensor)*or**ArrayLike**,**optional*) - A multi-dimensional array or an integer tensor describing the layout of devices, where the IDs are global IDs of the default process group. Default is None. - **mesh_dim_names** ([*tuple*](https://docs.python.org/3/library/stdtypes.html#tuple)*[*[*str*](https://docs.python.org/3/library/stdtypes.html#str)*,**...**]**,**optional*) - A tuple of mesh dimension names to assign to each dimension of the multi-dimensional array describing the layout of devices. Its length must match the length of mesh_shape. Each string in mesh_dim_names must be unique. Default is None. Returns: A `DeviceMesh` object representing the device layout. Return type: DeviceMesh get_all_groups()[[source]](https://github.com/pytorch/pytorch/blob/v2.13.0/torch/distributed/device_mesh.py#L831) Returns a list of ProcessGroups for all mesh dimensions. Returns: A list of `ProcessGroup` object. Return type: [list](https://docs.python.org/3/library/stdtypes.html#list)[[*ProcessGroup*](distributed._dist2.html#torch.distributed._dist2.ProcessGroup)] get_coordinate()[[source]](https://github.com/pytorch/pytorch/blob/v2.13.0/torch/distributed/device_mesh.py#L1232) Return the relative indices of this rank relative to all dimensions of the mesh. If this rank is not part of the mesh, return None. Return type: [tuple](https://docs.python.org/3/library/stdtypes.html#tuple)[[int](https://docs.python.org/3/library/functions.html#int), ...] | None get_group(*mesh_dim=None*)[[source]](https://github.com/pytorch/pytorch/blob/v2.13.0/torch/distributed/device_mesh.py#L782) Returns the single ProcessGroup specified by mesh_dim, or, if mesh_dim is not specified and the DeviceMesh is 1-dimensional, returns the only ProcessGroup in the mesh. Parameters: - **mesh_dim** (*str/python:int**,**optional*) - it can be the name of the mesh dimension or the index - **None.** (*of the mesh dimension. Default is*) - Returns: A `ProcessGroup` object. Return type: [*ProcessGroup*](distributed._dist2.html#torch.distributed._dist2.ProcessGroup) get_local_rank(*mesh_dim=None*)[[source]](https://github.com/pytorch/pytorch/blob/v2.13.0/torch/distributed/device_mesh.py#L1169) Returns the local rank of the given mesh_dim of the DeviceMesh. Parameters: - **mesh_dim** (*str/python:int**,**optional*) - it can be the name of the mesh dimension or the index - **None.** (*of the mesh dimension. Default is*) - Returns: An integer denotes the local rank. Return type: [int](https://docs.python.org/3/library/functions.html#int) The following program runs on each process/rank in an SPMD manner. In this example, we have 2 hosts with 4 GPUs each. Calling mesh_2d.get_local_rank(mesh_dim=0) on rank 0, 1, 2, 3 would return 0. Calling mesh_2d.get_local_rank(mesh_dim=0) on rank 4, 5, 6, 7 would return 1. Calling mesh_2d.get_local_rank(mesh_dim=1) on rank 0, 4 would return 0. Calling mesh_2d.get_local_rank(mesh_dim=1) on rank 1, 5 would return 1. Calling mesh_2d.get_local_rank(mesh_dim=1) on rank 2, 6 would return 2. Calling mesh_2d.get_local_rank(mesh_dim=1) on rank 3, 7 would return 3. Example: ``` >>> from torch.distributed.device_mesh import DeviceMesh >>> >>> # Initialize device mesh as (2, 4) to represent the topology >>> # of cross-host(dim 0), and within-host (dim 1). >>> mesh = DeviceMesh(device_type="cuda", mesh=[[0, 1, 2, 3],[4, 5, 6, 7]]) ``` get_rank()[[source]](https://github.com/pytorch/pytorch/blob/v2.13.0/torch/distributed/device_mesh.py#L1163) Returns the current global rank. Return type: [int](https://docs.python.org/3/library/functions.html#int) *property*mesh*: [Tensor](tensors.html#torch.Tensor)* Returns the tensor representing the layout of devices. *property*mesh_dim_names*: [tuple](https://docs.python.org/3/library/stdtypes.html#tuple)[[str](https://docs.python.org/3/library/stdtypes.html#str), ...] | [None](https://docs.python.org/3/library/constants.html#None)* Returns the names of mesh dimensions. ## Point-to-point communication torch.distributed.send(*tensor*, *dst=None*, *group=None*, *tag=0*, *group_dst=None*)[[source]](https://github.com/pytorch/pytorch/blob/v2.13.0/torch/distributed/distributed_c10d.py#L2712) Send a tensor synchronously. Warning `tag` is not supported with the NCCL backend. Parameters: - **tensor** ([*Tensor*](tensors.html#torch.Tensor)) - Tensor to send. - **dst** ([*int*](https://docs.python.org/3/library/functions.html#int)) - Destination rank on global process group (regardless of `group` argument). Destination rank should not be the same as the rank of the current process. - **group** ([*ProcessGroup*](distributed._dist2.html#torch.distributed._dist2.ProcessGroup)*,**optional*) - The process group to work on. If None, the default process group will be used. - **tag** ([*int*](https://docs.python.org/3/library/functions.html#int)*,**optional*) - Tag to match send with remote recv - **group_dst** ([*int*](https://docs.python.org/3/library/functions.html#int)*,**optional*) - Destination rank on `group`. Invalid to specify both `dst` and `group_dst`. torch.distributed.recv(*tensor*, *src=None*, *group=None*, *tag=0*, *group_src=None*)[[source]](https://github.com/pytorch/pytorch/blob/v2.13.0/torch/distributed/distributed_c10d.py#L2756) Receives a tensor synchronously. Warning `tag` is not supported with the NCCL backend. Parameters: - **tensor** ([*Tensor*](tensors.html#torch.Tensor)) - Tensor to fill with received data. - **src** ([*int*](https://docs.python.org/3/library/functions.html#int)*,**optional*) - Source rank on global process group (regardless of `group` argument). Will receive from any process if unspecified. - **group** ([*ProcessGroup*](distributed._dist2.html#torch.distributed._dist2.ProcessGroup)*,**optional*) - The process group to work on. If None, the default process group will be used. - **tag** ([*int*](https://docs.python.org/3/library/functions.html#int)*,**optional*) - Tag to match recv with remote send - **group_src** ([*int*](https://docs.python.org/3/library/functions.html#int)*,**optional*) - Destination rank on `group`. Invalid to specify both `src` and `group_src`. Returns: Sender rank -1, if not part of the group Return type: [int](https://docs.python.org/3/library/functions.html#int) `isend()` and `irecv()` return distributed request objects when used. In general, the type of this object is unspecified as they should never be created manually, but they are guaranteed to support two methods: - `is_completed()` - returns True if the operation has finished - `wait()` - will block the process until the operation is finished. `is_completed()` is guaranteed to return True once it returns. torch.distributed.isend(*tensor*, *dst=None*, *group=None*, *tag=0*, *group_dst=None*)[[source]](https://github.com/pytorch/pytorch/blob/v2.13.0/torch/distributed/distributed_c10d.py#L2598) Send a tensor asynchronously. Warning Modifying `tensor` before the request completes causes undefined behavior. Warning `tag` is not supported with the NCCL backend. Unlike send, which is blocking, isend allows src == dst rank, i.e. send to self. Parameters: - **tensor** ([*Tensor*](tensors.html#torch.Tensor)) - Tensor to send. - **dst** ([*int*](https://docs.python.org/3/library/functions.html#int)) - Destination rank on global process group (regardless of `group` argument) - **group** ([*ProcessGroup*](distributed._dist2.html#torch.distributed._dist2.ProcessGroup)*,**optional*) - The process group to work on. If None, the default process group will be used. - **tag** ([*int*](https://docs.python.org/3/library/functions.html#int)*,**optional*) - Tag to match send with remote recv - **group_dst** ([*int*](https://docs.python.org/3/library/functions.html#int)*,**optional*) - Destination rank on `group`. Invalid to specify both `dst` and `group_dst` Returns: A distributed request object. None, if not part of the group Return type: *Work* | None torch.distributed.irecv(*tensor*, *src=None*, *group=None*, *tag=0*, *group_src=None*)[[source]](https://github.com/pytorch/pytorch/blob/v2.13.0/torch/distributed/distributed_c10d.py#L2655) Receives a tensor asynchronously. Warning `tag` is not supported with the NCCL backend. Unlike recv, which is blocking, irecv allows src == dst rank, i.e. recv from self. Parameters: - **tensor** ([*Tensor*](tensors.html#torch.Tensor)) - Tensor to fill with received data. - **src** ([*int*](https://docs.python.org/3/library/functions.html#int)*,**optional*) - Source rank on global process group (regardless of `group` argument). Will receive from any process if unspecified. - **group** ([*ProcessGroup*](distributed._dist2.html#torch.distributed._dist2.ProcessGroup)*,**optional*) - The process group to work on. If None, the default process group will be used. - **tag** ([*int*](https://docs.python.org/3/library/functions.html#int)*,**optional*) - Tag to match recv with remote send - **group_src** ([*int*](https://docs.python.org/3/library/functions.html#int)*,**optional*) - Destination rank on `group`. Invalid to specify both `src` and `group_src`. Returns: A distributed request object. None, if not part of the group Return type: *Work* | None torch.distributed.send_object_list(*object_list*, *dst=None*, *group=None*, *device=None*, *group_dst=None*, *use_batch=False*)[[source]](https://github.com/pytorch/pytorch/blob/v2.13.0/torch/distributed/distributed_c10d.py#L3664) Sends picklable objects in `object_list` synchronously. Similar to `send()`, but Python objects can be passed in. Note that all objects in `object_list` must be picklable in order to be sent. Parameters: - **object_list** (*List**[**Any**]*) - List of input objects to sent. Each object must be picklable. Receiver must provide lists of equal sizes. - **dst** ([*int*](https://docs.python.org/3/library/functions.html#int)) - Destination rank to send `object_list` to. Destination rank is based on global process group (regardless of `group` argument) - **group** ([*ProcessGroup*](distributed._dist2.html#torch.distributed._dist2.ProcessGroup)*|**None*) - (ProcessGroup, optional): The process group to work on. If None, the default process group will be used. Default is `None`. - **device** (`torch.device`, optional) - If not None, the objects are serialized and converted to tensors which are moved to the `device` before sending. Default is `None`. - **group_dst** ([*int*](https://docs.python.org/3/library/functions.html#int)*,**optional*) - Destination rank on `group`. Must specify one of `dst` and `group_dst` but not both - **use_batch** ([*bool*](https://docs.python.org/3/library/functions.html#bool)*,**optional*) - If True, use batch p2p operations instead of regular send operations. This avoids initializing 2-rank communicators and uses existing entire group communicators. See batch_isend_irecv for usage and assumptions. Default is `False`. Returns: `None`. Note For NCCL-based process groups, internal tensor representations of objects must be moved to the GPU device before communication takes place. In this case, the device used is given by `torch.cuda.current_device()` and it is the user's responsibility to ensure that this is set so that each rank has an individual GPU, via `torch.cuda.set_device()`. Warning Object collectives have a number of serious performance and scalability limitations. See Object collectives for details. Warning `send_object_list()` uses `pickle` module implicitly, which is known to be insecure. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling. Only call this function with data you trust. Warning Calling `send_object_list()` with GPU tensors is not well supported and inefficient as it incurs GPU -> CPU transfer since tensors would be pickled. Please consider using `send()` instead. Example:: ``` >>> # Note: Process group initialization omitted on each rank. >>> import torch.distributed as dist >>> # Assumes backend is not NCCL >>> device = torch.device("cpu") >>> if dist.get_rank() == 0: >>> # Assumes world_size of 2. >>> objects = ["foo", 12, {1: 2}] # any picklable object >>> 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]](https://github.com/pytorch/pytorch/blob/v2.13.0/torch/distributed/distributed_c10d.py#L3782) Receives picklable objects in `object_list` synchronously. Similar to `recv()`, but can receive Python objects. Parameters: - **object_list** (*List**[**Any**]*) - List of objects to receive into. Must provide a list of sizes equal to the size of the list being sent. - **src** ([*int*](https://docs.python.org/3/library/functions.html#int)*,**optional*) - Source rank from which to recv `object_list`. Source rank is based on global process group (regardless of `group` argument) Will receive from any rank if set to None. Default is `None`. - **group** ([*ProcessGroup*](distributed._dist2.html#torch.distributed._dist2.ProcessGroup)*|**None*) - (ProcessGroup, optional): The process group to work on. If None, the default process group will be used. Default is `None`. - **device** (`torch.device`, optional) - If not None, receives on this device. Default is `None`. - **group_src** ([*int*](https://docs.python.org/3/library/functions.html#int)*,**optional*) - Destination rank on `group`. Invalid to specify both `src` and `group_src`. - **use_batch** ([*bool*](https://docs.python.org/3/library/functions.html#bool)*,**optional*) - If True, use batch p2p operations instead of regular send operations. This avoids initializing 2-rank communicators and uses existing entire group communicators. See batch_isend_irecv for usage and assumptions. Default is `False`. Returns: Sender rank. -1 if rank is not part of the group. If rank is part of the group, `object_list` will contain the sent objects from `src` rank. Note For NCCL-based process groups, internal tensor representations of objects must be moved to the GPU device before communication takes place. In this case, the device used is given by `torch.cuda.current_device()` and it is the user's responsibility to ensure that this is set so that each rank has an individual GPU, via `torch.cuda.set_device()`. Warning Object collectives have a number of serious performance and scalability limitations. See Object collectives for details. Warning `recv_object_list()` uses `pickle` module implicitly, which is known to be insecure. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling. Only call this function with data you trust. Warning Calling `recv_object_list()` with GPU tensors is not well supported and inefficient as it incurs GPU -> CPU transfer since tensors would be pickled. Please consider using `recv()` instead. Example:: ``` >>> # Note: Process group initialization omitted on each rank. >>> import torch.distributed as dist >>> # Assumes backend is not NCCL >>> device = torch.device("cpu") >>> if dist.get_rank() == 0: >>> # Assumes world_size of 2. >>> objects = ["foo", 12, {1: 2}] # any picklable object >>> 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]](https://github.com/pytorch/pytorch/blob/v2.13.0/torch/distributed/distributed_c10d.py#L2990) Send or Receive a batch of tensors asynchronously and return a list of requests. Process each of the operations in `p2p_op_list` and return the corresponding requests. NCCL, Gloo, and UCC backend are currently supported. Parameters: **p2p_op_list** ([*list*](https://docs.python.org/3/library/stdtypes.html#list)*[**P2POp**]*) - A list of point-to-point operations(type of each operator is `torch.distributed.P2POp`). All operations in the list are treated as a single batch - the relative ordering of sends vs. receives in the list does not matter and will not cause deadlocks. For example, `[recv_from_1, send_to_1]` is equivalent to `[send_to_1, recv_from_1]`. The ordering **does** matter for multiple operations between the same pair of ranks in the same direction: if rank 0 calls `[send_to_1(A), send_to_1(B)]`, rank 1 must call `[recv_from_0(A), recv_from_0(B)]` in the same order to ensure the correct tensors are matched. Returns: A list of distributed request objects returned by calling the corresponding op in the op_list. Return type: [list](https://docs.python.org/3/library/stdtypes.html#list)[*Work*] Examples ``` >>> 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]) # Rank 0 tensor([0, 1]) # Rank 1 ``` Note Note that when this API is used with the NCCL PG backend, users must set the current GPU device with torch.cuda.set_device, otherwise it will lead to unexpected hang issues. In addition, if this API is the first collective call in the `group` passed to `dist.P2POp`, all ranks of the `group` must participate in this API call; otherwise, the behavior is undefined. If this API call is not the first collective call in the `group`, batched P2P operations involving only a subset of ranks of the `group` are allowed. *class*torch.distributed.P2POp(*op*, *tensor*, *peer=None*, *group=None*, *tag=0*, *group_peer=None*)[[source]](https://github.com/pytorch/pytorch/blob/v2.13.0/torch/distributed/distributed_c10d.py#L563) A class to build point-to-point operations for `batch_isend_irecv`. This class builds the type of P2P operation, communication buffer, peer rank, Process Group, and tag. Instances of this class will be passed to `batch_isend_irecv` for point-to-point communications. Parameters: - **op** (*Callable*) - A function to send data to or receive data from a peer process. The type of `op` is either `torch.distributed.isend` or `torch.distributed.irecv`. - **tensor** ([*Tensor*](tensors.html#torch.Tensor)) - Tensor to send or receive. - **peer** ([*int*](https://docs.python.org/3/library/functions.html#int)*,**optional*) - Destination or source rank. - **group** ([*ProcessGroup*](distributed._dist2.html#torch.distributed._dist2.ProcessGroup)*,**optional*) - The process group to work on. If None, the default process group will be used. - **tag** ([*int*](https://docs.python.org/3/library/functions.html#int)*,**optional*) - Tag to match send with recv. - **group_peer** ([*int*](https://docs.python.org/3/library/functions.html#int)*,**optional*) - Destination or source rank. ## Synchronous and asynchronous collective operations Every collective operation function supports the following two kinds of operations, depending on the setting of the `async_op` flag passed into the collective: **Synchronous operation** - the default mode, when `async_op` is set to `False`. When the function returns, it is guaranteed that the collective operation is performed. In the case of CUDA operations, it is not guaranteed that the CUDA operation is completed, since CUDA operations are asynchronous. For CPU collectives, any further function calls utilizing the output of the collective call will behave as expected. For CUDA collectives, function calls utilizing the output on the same CUDA stream will behave as expected. Users must take care of synchronization under the scenario of running under different streams. For details on CUDA semantics such as stream synchronization, see [CUDA Semantics](https://pytorch.org/docs/stable/notes/cuda.html). See the below script to see examples of differences in these semantics for CPU and CUDA operations. **Asynchronous operation** - when `async_op` is set to True. The collective operation function returns a distributed request object. In general, you don't need to create it manually and it is guaranteed to support two methods: - `is_completed()` - in the case of CPU collectives, returns `True` if completed. In the case of CUDA operations, returns `True` if the operation has been successfully enqueued onto a CUDA stream and the output can be utilized on the default stream without further synchronization. - `wait()` - in the case of CPU collectives, will block the process until the operation is completed. In the case of CUDA collectives, will block the currently active CUDA stream until the operation is completed (but will not block the CPU). - `get_future()` - returns `torch._C.Future` object. Supported for NCCL, also supported for most operations on GLOO and MPI, except for peer to peer operations. Note: as we continue adopting Futures and merging APIs, `get_future()` call might become redundant. **Example** The following code can serve as a reference regarding semantics for CUDA operations when using distributed collectives. It shows the explicit need to synchronize when using collective outputs on different CUDA streams: ``` # Code runs on each rank. 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 ensures the operation is enqueued, but not necessarily complete. handle.wait() # Using result on non-default stream. with torch.cuda.stream(s): s.wait_stream(torch.cuda.default_stream()) output.add_(100) if rank == 0: # if the explicit call to wait_stream was omitted, the output below will be # non-deterministically 1 or 101, depending on whether the allreduce overwrote # the value after the add completed. print(output) ``` ## Collective functions torch.distributed.broadcast(*tensor*, *src=None*, *group=None*, *async_op=False*, *group_src=None*)[[source]](https://github.com/pytorch/pytorch/blob/v2.13.0/torch/distributed/distributed_c10d.py#L3085) Broadcasts the tensor to the whole group. `tensor` must have the same number of elements in all processes participating in the collective. Parameters: - **tensor** ([*Tensor*](tensors.html#torch.Tensor)) - Data to be sent if `src` is the rank of current process, and tensor to be used to save received data otherwise. - **src** ([*int*](https://docs.python.org/3/library/functions.html#int)) - Source rank on global process group (regardless of `group` argument). - **group** ([*ProcessGroup*](distributed._dist2.html#torch.distributed._dist2.ProcessGroup)*,**optional*) - The process group to work on. If None, the default process group will be used. - **async_op** ([*bool*](https://docs.python.org/3/library/functions.html#bool)*,**optional*) - Whether this op should be an async op - **group_src** ([*int*](https://docs.python.org/3/library/functions.html#int)) - Source rank on `group`. Must specify one of `group_src` and `src` but not both. Returns: Async work handle, if async_op is set to True. None, if not async_op or if not part of the group torch.distributed.broadcast_object_list(*object_list*, *src=None*, *group=None*, *device=None*, *group_src=None*)[[source]](https://github.com/pytorch/pytorch/blob/v2.13.0/torch/distributed/distributed_c10d.py#L3924) Broadcasts picklable objects in `object_list` to the whole group. Similar to `broadcast()`, but Python objects can be passed in. Note that all objects in `object_list` must be picklable in order to be broadcasted. Parameters: - **object_list** (*List**[**Any**]*) - List of input objects to broadcast. Each object must be picklable. Only objects on the `src` rank will be broadcast, but each rank must provide lists of equal sizes. - **src** ([*int*](https://docs.python.org/3/library/functions.html#int)) - Source rank from which to broadcast `object_list`. Source rank is based on global process group (regardless of `group` argument) - **group** ([*ProcessGroup*](distributed._dist2.html#torch.distributed._dist2.ProcessGroup)*|**None*) - (ProcessGroup, optional): The process group to work on. If None, the default process group will be used. Default is `None`. - **device** (`torch.device`, optional) - If not None, the objects are serialized and converted to tensors which are moved to the `device` before broadcasting. Default is `None`. - **group_src** ([*int*](https://docs.python.org/3/library/functions.html#int)) - Source rank on `group`. Must not specify one of `group_src` and `src` but not both. Returns: `None`. If rank is part of the group, `object_list` will contain the broadcasted objects from `src` rank. Note For NCCL-based process groups, internal tensor representations of objects must be moved to the GPU device before communication takes place. In this case, the device used is given by `torch.cuda.current_device()` and it is the user's responsibility to ensure that this is set so that each rank has an individual GPU, via `torch.cuda.set_device()`. Note Note that this API differs slightly from the `broadcast()` collective since it does not provide an `async_op` handle and thus will be a blocking call. Warning Object collectives have a number of serious performance and scalability limitations. See Object collectives for details. Warning `broadcast_object_list()` uses `pickle` module implicitly, which is known to be insecure. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling. Only call this function with data you trust. Warning Calling `broadcast_object_list()` with GPU tensors is not well supported and inefficient as it incurs GPU -> CPU transfer since tensors would be pickled. Please consider using `broadcast()` instead. Example:: ``` >>> # Note: Process group initialization omitted on each rank. >>> import torch.distributed as dist >>> if dist.get_rank() == 0: >>> # Assumes world_size of 3. >>> objects = ["foo", 12, {1: 2}] # any picklable object >>> else: >>> objects = [None, None, None] >>> # Assumes backend is not 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=*, *group=None*, *async_op=False*)[[source]](https://github.com/pytorch/pytorch/blob/v2.13.0/torch/distributed/distributed_c10d.py#L3155) Reduces the tensor data across all machines in a way that all get the final result. After the call `tensor` is going to be bitwise identical in all processes. Complex tensors are supported. Parameters: - **tensor** ([*Tensor*](tensors.html#torch.Tensor)) - Input and output of the collective. The function operates in-place. - **op** (*optional*) - One of the values from `torch.distributed.ReduceOp` enum. Specifies an operation used for element-wise reductions. - **group** ([*ProcessGroup*](distributed._dist2.html#torch.distributed._dist2.ProcessGroup)*,**optional*) - The process group to work on. If None, the default process group will be used. - **async_op** ([*bool*](https://docs.python.org/3/library/functions.html#bool)*,**optional*) - Whether this op should be an async op Returns: Async work handle, if async_op is set to True. None, if not async_op or if not part of the group Examples ``` >>> # All tensors below are of torch.int64 type. >>> # We have 2 process groups, 2 ranks. >>> 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') # Rank 0 tensor([3, 4], device='cuda:1') # Rank 1 >>> dist.all_reduce(tensor, op=ReduceOp.SUM) >>> tensor tensor([4, 6], device='cuda:0') # Rank 0 tensor([4, 6], device='cuda:1') # Rank 1 ``` ``` >>> # All tensors below are of torch.cfloat type. >>> # We have 2 process groups, 2 ranks. >>> 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') # Rank 0 tensor([3.+3.j, 4.+4.j], device='cuda:1') # Rank 1 >>> dist.all_reduce(tensor, op=ReduceOp.SUM) >>> tensor tensor([4.+4.j, 6.+6.j], device='cuda:0') # Rank 0 tensor([4.+4.j, 6.+6.j], device='cuda:1') # Rank 1 ``` torch.distributed.all_reduce_coalesced(*tensors*, *op=*, *group=None*, *async_op=False*)[[source]](https://github.com/pytorch/pytorch/blob/v2.13.0/torch/distributed/distributed_c10d.py#L3256) WARNING: at this time individual shape checking is not implemented across nodes. For example, if the rank 0 node passes [torch.rand(4), torch.rand(2)] and the rank 1 node passes [torch.rand(2), torch.rand(2), torch.rand(2)], the allreduce operation will proceed without complaint and return erroneous outputs. This lack of shape checking results in significant performance improvements but users of this function should take extra care to ensure that each node passes in tensors whose shapes match across nodes. Reduces each tensor in tensors (residing on the same device) across all machines in such a way that all get the final result. After the call each tensor in tensors is going to bitwise identical in all processes. Complex tensors are supported. Parameters: - **tensors** (*Union**[**List**[*[*Tensor*](tensors.html#torch.Tensor)*]**,*[*Tensor*](tensors.html#torch.Tensor)*]*) - Input and output of the collective. The function operates in-place. - **op** (*Optional**[**ReduceOp**]*) - One of the values from `torch.distributed.ReduceOp` enum. Specifies an operation used for element-wise reductions. - **group** ([*ProcessGroup*](distributed._dist2.html#torch.distributed._dist2.ProcessGroup)*,**optional*) - The process group to work on. If None, the default process group will be used. - **async_op** (*Optional**[*[*bool*](https://docs.python.org/3/library/functions.html#bool)*]*) - Whether this op should be an async op. Returns: Async work handle, if async_op is set to True. None, if not async_op or if not part of the group. torch.distributed.reduce(*tensor*, *dst=None*, *op=*, *group=None*, *async_op=False*, *group_dst=None*)[[source]](https://github.com/pytorch/pytorch/blob/v2.13.0/torch/distributed/distributed_c10d.py#L3336) Reduces the tensor data across all machines. Only the process with rank `dst` is going to receive the final result. Parameters: - **tensor** ([*Tensor*](tensors.html#torch.Tensor)) - Input and output of the collective. The function operates in-place. - **dst** ([*int*](https://docs.python.org/3/library/functions.html#int)) - Destination rank on global process group (regardless of `group` argument) - **op** (*optional*) - One of the values from `torch.distributed.ReduceOp` enum. Specifies an operation used for element-wise reductions. - **group** ([*ProcessGroup*](distributed._dist2.html#torch.distributed._dist2.ProcessGroup)*,**optional*) - The process group to work on. If None, the default process group will be used. - **async_op** ([*bool*](https://docs.python.org/3/library/functions.html#bool)*,**optional*) - Whether this op should be an async op - **group_dst** ([*int*](https://docs.python.org/3/library/functions.html#int)) - Destination rank on `group`. Must specify one of `group_dst` and `dst` but not both. Returns: Async work handle, if async_op is set to True. None, if not async_op or if not part of the group torch.distributed.all_gather(*tensor_list*, *tensor*, *group=None*, *async_op=False*)[[source]](https://github.com/pytorch/pytorch/blob/v2.13.0/torch/distributed/distributed_c10d.py#L4191) Gathers tensors from the whole group in a list. Complex and uneven sized tensors are supported. Parameters: - **tensor_list** ([*list*](https://docs.python.org/3/library/stdtypes.html#list)*[*[*Tensor*](tensors.html#torch.Tensor)*]*) - Output list. It should contain correctly-sized tensors to be used for output of the collective. Uneven sized tensors are supported. - **tensor** ([*Tensor*](tensors.html#torch.Tensor)) - Tensor to be broadcast from current process. - **group** ([*ProcessGroup*](distributed._dist2.html#torch.distributed._dist2.ProcessGroup)*,**optional*) - The process group to work on. If None, the default process group will be used. - **async_op** ([*bool*](https://docs.python.org/3/library/functions.html#bool)*,**optional*) - Whether this op should be an async op Returns: Async work handle, if async_op is set to True. None, if not async_op or if not part of the group Examples ``` >>> # All tensors below are of torch.int64 dtype. >>> # We have 2 process groups, 2 ranks. >>> 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')] # Rank 0 [tensor([0, 0], device='cuda:1'), tensor([0, 0], device='cuda:1')] # Rank 1 >>> tensor = torch.arange(2, dtype=torch.int64, device=device) + 1 + 2 * rank >>> tensor tensor([1, 2], device='cuda:0') # Rank 0 tensor([3, 4], device='cuda:1') # Rank 1 >>> dist.all_gather(tensor_list, tensor) >>> tensor_list [tensor([1, 2], device='cuda:0'), tensor([3, 4], device='cuda:0')] # Rank 0 [tensor([1, 2], device='cuda:1'), tensor([3, 4], device='cuda:1')] # Rank 1 ``` ``` >>> # All tensors below are of torch.cfloat dtype. >>> # We have 2 process groups, 2 ranks. >>> 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')] # Rank 0 [tensor([0.+0.j, 0.+0.j], device='cuda:1'), tensor([0.+0.j, 0.+0.j], device='cuda:1')] # Rank 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') # Rank 0 tensor([3.+3.j, 4.+4.j], device='cuda:1') # Rank 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')] # Rank 0 [tensor([1.+1.j, 2.+2.j], device='cuda:1'), tensor([3.+3.j, 4.+4.j], device='cuda:1')] # Rank 1 ``` torch.distributed.all_gather_single(*output_tensor*, *input_tensor*, *group=None*, *async_op=False*)[[source]](https://github.com/pytorch/pytorch/blob/v2.13.0/torch/distributed/distributed_c10d.py#L4291) Gather tensors from all ranks and put them in a single output tensor. This function requires all tensors to be the same size on each process. Parameters: - **output_tensor** ([*Tensor*](tensors.html#torch.Tensor)) - Output tensor to accommodate tensor elements from all ranks. It must be correctly sized to have one of the following forms: (i) a concatenation of all the input tensors along the primary dimension; for definition of "concatenation", see `torch.cat()`; (ii) a stack of all the input tensors along the primary dimension; for definition of "stack", see `torch.stack()`. Examples below may better explain the supported output forms. - **input_tensor** ([*Tensor*](tensors.html#torch.Tensor)) - Tensor to be gathered from current rank. Different from the `all_gather` API, the input tensors in this API must have the same size across all ranks. - **group** ([*ProcessGroup*](distributed._dist2.html#torch.distributed._dist2.ProcessGroup)*,**optional*) - The process group to work on. If None, the default process group will be used. - **async_op** ([*bool*](https://docs.python.org/3/library/functions.html#bool)*,**optional*) - Whether this op should be an async op Returns: Async work handle, if async_op is set to True. None, if not async_op or if not part of the group Examples ``` >>> # All tensors below are of torch.int64 dtype and on CUDA devices. >>> # We have two ranks. >>> 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') # Rank 0 tensor([3, 4], device='cuda:1') # Rank 1 >>> # Output in concatenation form >>> tensor_out = torch.zeros(world_size * 2, dtype=torch.int64, device=device) >>> dist.all_gather_single(tensor_out, tensor_in) >>> tensor_out tensor([1, 2, 3, 4], device='cuda:0') # Rank 0 tensor([1, 2, 3, 4], device='cuda:1') # Rank 1 >>> # Output in stack form >>> tensor_out2 = torch.zeros(world_size, 2, dtype=torch.int64, device=device) >>> dist.all_gather_single(tensor_out2, tensor_in) >>> tensor_out2 tensor([[1, 2], [3, 4]], device='cuda:0') # Rank 0 tensor([[1, 2], [3, 4]], device='cuda:1') # Rank 1 ``` torch.distributed.all_gather_object(*object_list*, *obj*, *group=None*)[[source]](https://github.com/pytorch/pytorch/blob/v2.13.0/torch/distributed/distributed_c10d.py#L3438) Gathers picklable objects from the whole group into a list. Similar to `all_gather()`, but Python objects can be passed in. Note that the object must be picklable in order to be gathered. Parameters: - **object_list** ([*list*](https://docs.python.org/3/library/stdtypes.html#list)*[**Any**]*) - Output list. It should be correctly sized as the size of the group for this collective and will contain the output. - **obj** (*Any*) - Pickable Python object to be broadcast from current process. - **group** ([*ProcessGroup*](distributed._dist2.html#torch.distributed._dist2.ProcessGroup)*,**optional*) - The process group to work on. If None, the default process group will be used. Default is `None`. Returns: None. If the calling rank is part of this group, the output of the collective will be populated into the input `object_list`. If the calling rank is not part of the group, the passed in `object_list` will be unmodified. Note Note that this API differs slightly from the `all_gather()` collective since it does not provide an `async_op` handle and thus will be a blocking call. Note For NCCL-based processed groups, internal tensor representations of objects must be moved to the GPU device before communication takes place. In this case, the device used is given by `torch.cuda.current_device()` and it is the user's responsibility to ensure that this is set so that each rank has an individual GPU, via `torch.cuda.set_device()`. Warning Object collectives have a number of serious performance and scalability limitations. See Object collectives for details. Warning `all_gather_object()` uses `pickle` module implicitly, which is known to be insecure. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling. Only call this function with data you trust. Warning Calling `all_gather_object()` with GPU tensors is not well supported and inefficient as it incurs GPU -> CPU transfer since tensors would be pickled. Please consider using `all_gather()` instead. Example:: ``` >>> # Note: Process group initialization omitted on each rank. >>> import torch.distributed as dist >>> # Assumes world_size of 3. >>> gather_objects = ["foo", 12, {1: 2}] # any picklable object >>> output = [None for _ in gather_objects] >>> dist.all_gather_object(output, gather_objects[dist.get_rank()]) >>> output ['foo', 12, {1: 2}] ``` torch.distributed.all_gather_coalesced(*output_tensor_lists*, *input_tensor_list*, *group=None*, *async_op=False*)[[source]](https://github.com/pytorch/pytorch/blob/v2.13.0/torch/distributed/distributed_c10d.py#L4446) Gathers input tensors from the whole group in a list in a coalesced manner. Complex tensors are supported. Parameters: - **output_tensor_lists** ([*list*](https://docs.python.org/3/library/stdtypes.html#list)*[*[*list*](https://docs.python.org/3/library/stdtypes.html#list)*[*[*Tensor*](tensors.html#torch.Tensor)*]**]*) - Output list. It should contain correctly-sized tensors to be used for output of the collective. - **input_tensor_list** ([*list*](https://docs.python.org/3/library/stdtypes.html#list)*[*[*Tensor*](tensors.html#torch.Tensor)*]*) - Tensors to be broadcast from current process. At least one tensor has to be non empty. - **group** ([*ProcessGroup*](distributed._dist2.html#torch.distributed._dist2.ProcessGroup)*,**optional*) - The process group to work on. If None, the default process group will be used. - **async_op** ([*bool*](https://docs.python.org/3/library/functions.html#bool)*,**optional*) - Whether this op should be an async op. Returns: Async work handle, if async_op is set to True. None, if not async_op or if not part of the group Example: ``` # we have 2 process groups, 2 ranks. # rank 0 passes: input_tensor_list = [[[1, 1], [1, 1]], [2], [3, 3]] output_tensor_lists = [ [[[-1, -1], [-1, -1]], [-1], [-1, -1]], [[[-1, -1], [-1, -1]], [-1], [-1, -1]], ] # rank 1 passes: input_tensor_list = [[[3, 3], [3, 3]], [5], [1, 1]] output_tensor_lists = [ [[[-1, -1], [-1, -1]], [-1], [-1, -1]], [[[-1, -1], [-1, -1]], [-1], [-1, -1]], ] # both rank 0 and 1 get: output_tensor_lists = [ [[[1, 1], [1, 1]], [2], [3, 3]], [[[3, 3], [3, 3]], [5], [1, 1]], ] ``` WARNING: at this time individual shape checking is not implemented across nodes. For example, if the rank 0 node passes [torch.rand(4), torch.rand(2)] and the rank 1 node passes [torch.rand(2), torch.rand(2), torch.rand(2)], the all_gather_coalesced operation will proceed without complaint and return erroneous outputs. This lack of shape checking results in significant performance improvements but users of this function should take extra care to ensure that each node passes in tensors whose shapes match across nodes. torch.distributed.gather(*tensor*, *gather_list=None*, *dst=None*, *group=None*, *async_op=False*, *group_dst=None*)[[source]](https://github.com/pytorch/pytorch/blob/v2.13.0/torch/distributed/distributed_c10d.py#L4567) Gathers a list of tensors in a single process. This function requires all tensors to be the same size on each process. Parameters: - **tensor** ([*Tensor*](tensors.html#torch.Tensor)) - Input tensor. - **gather_list** ([*list*](https://docs.python.org/3/library/stdtypes.html#list)*[*[*Tensor*](tensors.html#torch.Tensor)*]**,**optional*) - List of appropriately, same-sized tensors to use for gathered data (default is None, must be specified on the destination rank) - **dst** ([*int*](https://docs.python.org/3/library/functions.html#int)*,**optional*) - Destination rank on global process group (regardless of `group` argument). (If both `dst` and `group_dst` are None, default is global rank 0) - **group** ([*ProcessGroup*](distributed._dist2.html#torch.distributed._dist2.ProcessGroup)*,**optional*) - The process group to work on. If None, the default process group will be used. - **async_op** ([*bool*](https://docs.python.org/3/library/functions.html#bool)*,**optional*) - Whether this op should be an async op - **group_dst** ([*int*](https://docs.python.org/3/library/functions.html#int)*,**optional*) - Destination rank on `group`. Invalid to specify both `dst` and `group_dst` Returns: Async work handle, if async_op is set to True. None, if not async_op or if not part of the group Note Note that all Tensors in gather_list must have the same size. Example:: ``` >>> # We have 2 process groups, 2 ranks. >>> 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) >>> # Rank 0 gets gathered data. >>> gather_list [tensor([1., 1.], device='cuda:0'), tensor([2., 2.], device='cuda:0')] # Rank 0 None # Rank 1 ``` torch.distributed.gather_object(*obj*, *object_gather_list=None*, *dst=None*, *group=None*, *group_dst=None*)[[source]](https://github.com/pytorch/pytorch/blob/v2.13.0/torch/distributed/distributed_c10d.py#L3533) Gathers picklable objects from the whole group in a single process. Similar to `gather()`, but Python objects can be passed in. Note that the object must be picklable in order to be gathered. Parameters: - **obj** (*Any*) - Input object. Must be picklable. - **object_gather_list** ([*list*](https://docs.python.org/3/library/stdtypes.html#list)*[**Any**]*) - Output list. On the `dst` rank, it should be correctly sized as the size of the group for this collective and will contain the output. Must be `None` on non-dst ranks. (default is `None`) - **dst** ([*int*](https://docs.python.org/3/library/functions.html#int)*,**optional*) - Destination rank on global process group (regardless of `group` argument). (If both `dst` and `group_dst` are None, default is global rank 0) - **group** ([*ProcessGroup*](distributed._dist2.html#torch.distributed._dist2.ProcessGroup)*|**None*) - (ProcessGroup, optional): The process group to work on. If None, the default process group will be used. Default is `None`. - **group_dst** ([*int*](https://docs.python.org/3/library/functions.html#int)*,**optional*) - Destination rank on `group`. Invalid to specify both `dst` and `group_dst` Returns: None. On the `dst` rank, `object_gather_list` will contain the output of the collective. Note Note that this API differs slightly from the gather collective since it does not provide an async_op handle and thus will be a blocking call. Note For NCCL-based processed groups, internal tensor representations of objects must be moved to the GPU device before communication takes place. In this case, the device used is given by `torch.cuda.current_device()` and it is the user's responsibility to ensure that this is set so that each rank has an individual GPU, via `torch.cuda.set_device()`. Warning Object collectives have a number of serious performance and scalability limitations. See Object collectives for details. Warning `gather_object()` uses `pickle` module implicitly, which is known to be insecure. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling. Only call this function with data you trust. Warning Calling `gather_object()` with GPU tensors is not well supported and inefficient as it incurs GPU -> CPU transfer since tensors would be pickled. Please consider using `gather()` instead. Example:: ``` >>> # Note: Process group initialization omitted on each rank. >>> import torch.distributed as dist >>> # Assumes world_size of 3. >>> gather_objects = ["foo", 12, {1: 2}] # any picklable object >>> output = [None for _ in gather_objects] >>> dist.gather_object( ... gather_objects[dist.get_rank()], ... output if dist.get_rank() == 0 else None, ... dst=0 ... ) >>> # On rank 0 >>> output ['foo', 12, {1: 2}] ``` torch.distributed.scatter(*tensor*, *scatter_list=None*, *src=None*, *group=None*, *async_op=False*, *group_src=None*)[[source]](https://github.com/pytorch/pytorch/blob/v2.13.0/torch/distributed/distributed_c10d.py#L4671) Scatters a list of tensors to all processes in a group. Each process will receive exactly one tensor and store its data in the `tensor` argument. Complex tensors are supported. Parameters: - **tensor** ([*Tensor*](tensors.html#torch.Tensor)) - Output tensor. - **scatter_list** ([*list*](https://docs.python.org/3/library/stdtypes.html#list)*[*[*Tensor*](tensors.html#torch.Tensor)*]*) - List of tensors to scatter (default is None, must be specified on the source rank) - **src** ([*int*](https://docs.python.org/3/library/functions.html#int)) - Source rank on global process group (regardless of `group` argument). (If both `src` and `group_src` are None, default is global rank 0) - **group** ([*ProcessGroup*](distributed._dist2.html#torch.distributed._dist2.ProcessGroup)*,**optional*) - The process group to work on. If None, the default process group will be used. - **async_op** ([*bool*](https://docs.python.org/3/library/functions.html#bool)*,**optional*) - Whether this op should be an async op - **group_src** ([*int*](https://docs.python.org/3/library/functions.html#int)*,**optional*) - Source rank on `group`. Invalid to specify both `src` and `group_src` Returns: Async work handle, if async_op is set to True. None, if not async_op or if not part of the group Note Note that all Tensors in scatter_list must have the same size. Example:: ``` >>> # Note: Process group initialization omitted on each rank. >>> 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: >>> # Assumes world_size of 2. >>> # Only tensors, all of which must be the same size. >>> 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) >>> # Rank i gets scatter_list[i]. >>> output_tensor tensor([1., 1.], device='cuda:0') # Rank 0 tensor([5., 5.], device='cuda:1') # Rank 1 ``` torch.distributed.scatter_object_list(*scatter_object_output_list*, *scatter_object_input_list=None*, *src=None*, *group=None*, *group_src=None*)[[source]](https://github.com/pytorch/pytorch/blob/v2.13.0/torch/distributed/distributed_c10d.py#L4056) Scatters picklable objects in `scatter_object_input_list` to the whole group. Similar to `scatter()`, but Python objects can be passed in. On each rank, the scattered object will be stored as the first element of `scatter_object_output_list`. Note that all objects in `scatter_object_input_list` must be picklable in order to be scattered. Parameters: - **scatter_object_output_list** (*List**[**Any**]*) - Non-empty list whose first element will store the object scattered to this rank. - **scatter_object_input_list** (*List**[**Any**]**,**optional*) - List of input objects to scatter. Each object must be picklable. Only objects on the `src` rank will be scattered, and the argument can be `None` for non-src ranks. - **src** ([*int*](https://docs.python.org/3/library/functions.html#int)) - Source rank from which to scatter `scatter_object_input_list`. Source rank is based on global process group (regardless of `group` argument). (If both `src` and `group_src` are None, default is global rank 0) - **group** ([*ProcessGroup*](distributed._dist2.html#torch.distributed._dist2.ProcessGroup)*|**None*) - (ProcessGroup, optional): The process group to work on. If None, the default process group will be used. Default is `None`. - **group_src** ([*int*](https://docs.python.org/3/library/functions.html#int)*,**optional*) - Source rank on `group`. Invalid to specify both `src` and `group_src` Returns: `None`. If rank is part of the group, `scatter_object_output_list` will have its first element set to the scattered object for this rank. Note Note that this API differs slightly from the scatter collective since it does not provide an `async_op` handle and thus will be a blocking call. Warning Object collectives have a number of serious performance and scalability limitations. See Object collectives for details. Warning `scatter_object_list()` uses `pickle` module implicitly, which is known to be insecure. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling. Only call this function with data you trust. Warning Calling `scatter_object_list()` with GPU tensors is not well supported and inefficient as it incurs GPU -> CPU transfer since tensors would be pickled. Please consider using `scatter()` instead. Example:: ``` >>> # Note: Process group initialization omitted on each rank. >>> import torch.distributed as dist >>> if dist.get_rank() == 0: >>> # Assumes world_size of 3. >>> objects = ["foo", 12, {1: 2}] # any picklable object >>> else: >>> # Can be any list on non-src ranks, elements are not used. >>> objects = [None, None, None] >>> output_list = [None] >>> dist.scatter_object_list(output_list, objects, src=0) >>> # Rank i gets objects[i]. For example, on rank 2: >>> output_list [{1: 2}] ``` torch.distributed.reduce_scatter(*output*, *input_list*, *op=*, *group=None*, *async_op=False*)[[source]](https://github.com/pytorch/pytorch/blob/v2.13.0/torch/distributed/distributed_c10d.py#L4789) Reduces, then scatters a list of tensors to all processes in a group. Parameters: - **output** ([*Tensor*](tensors.html#torch.Tensor)) - Output tensor. - **input_list** ([*list*](https://docs.python.org/3/library/stdtypes.html#list)*[*[*Tensor*](tensors.html#torch.Tensor)*]*) - List of tensors to reduce and scatter. - **op** (*optional*) - One of the values from `torch.distributed.ReduceOp` enum. Specifies an operation used for element-wise reductions. - **group** ([*ProcessGroup*](distributed._dist2.html#torch.distributed._dist2.ProcessGroup)*,**optional*) - The process group to work on. If None, the default process group will be used. - **async_op** ([*bool*](https://docs.python.org/3/library/functions.html#bool)*,**optional*) - Whether this op should be an async op. Returns: Async work handle, if async_op is set to True. None, if not async_op or if not part of the group. torch.distributed.reduce_scatter_single(*output*, *input*, *op=*, *group=None*, *async_op=False*)[[source]](https://github.com/pytorch/pytorch/blob/v2.13.0/torch/distributed/distributed_c10d.py#L4846) Reduces, then scatters a tensor to all ranks in a group. Parameters: - **output** ([*Tensor*](tensors.html#torch.Tensor)) - Output tensor. It should have the same size across all ranks. - **input** ([*Tensor*](tensors.html#torch.Tensor)) - Input tensor to be reduced and scattered. Its size should be output tensor size times the world size. The input tensor can have one of the following shapes: (i) a concatenation of the output tensors along the primary dimension, or (ii) a stack of the output tensors along the primary dimension. For definition of "concatenation", see `torch.cat()`. For definition of "stack", see `torch.stack()`. - **group** ([*ProcessGroup*](distributed._dist2.html#torch.distributed._dist2.ProcessGroup)*,**optional*) - The process group to work on. If None, the default process group will be used. - **async_op** ([*bool*](https://docs.python.org/3/library/functions.html#bool)*,**optional*) - Whether this op should be an async op. Returns: Async work handle, if async_op is set to True. None, if not async_op or if not part of the group. Examples ``` >>> # All tensors below are of torch.int64 dtype and on CUDA devices. >>> # We have two ranks. >>> device = torch.device(f"cuda:{rank}") >>> tensor_out = torch.zeros(2, dtype=torch.int64, device=device) >>> # Input in concatenation form >>> tensor_in = torch.arange(world_size * 2, dtype=torch.int64, device=device) >>> tensor_in tensor([0, 1, 2, 3], device='cuda:0') # Rank 0 tensor([0, 1, 2, 3], device='cuda:1') # Rank 1 >>> dist.reduce_scatter_single(tensor_out, tensor_in) >>> tensor_out tensor([0, 2], device='cuda:0') # Rank 0 tensor([4, 6], device='cuda:1') # Rank 1 >>> # Input in stack form >>> tensor_in = torch.reshape(tensor_in, (world_size, 2)) >>> tensor_in tensor([[0, 1], [2, 3]], device='cuda:0') # Rank 0 tensor([[0, 1], [2, 3]], device='cuda:1') # Rank 1 >>> dist.reduce_scatter_single(tensor_out, tensor_in) >>> tensor_out tensor([0, 2], device='cuda:0') # Rank 0 tensor([4, 6], device='cuda:1') # Rank 1 ``` torch.distributed.all_to_all_single(*output*, *input*, *output_split_sizes=None*, *input_split_sizes=None*, *group=None*, *async_op=False*)[[source]](https://github.com/pytorch/pytorch/blob/v2.13.0/torch/distributed/distributed_c10d.py#L4995) Split input tensor and then scatter the split list to all processes in a group. Later the received tensors are concatenated from all the processes in the group and returned as a single output tensor. Complex tensors are supported. Parameters: - **output** ([*Tensor*](tensors.html#torch.Tensor)) - Gathered concatenated output tensor. - **input** ([*Tensor*](tensors.html#torch.Tensor)) - Input tensor to scatter. - **output_split_sizes** - (list[Int], optional): Output split sizes for dim 0 if specified None or empty, dim 0 of `output` tensor must divide equally by `world_size`. - **input_split_sizes** - (list[Int], optional): Input split sizes for dim 0 if specified None or empty, dim 0 of `input` tensor must divide equally by `world_size`. - **group** ([*ProcessGroup*](distributed._dist2.html#torch.distributed._dist2.ProcessGroup)*,**optional*) - The process group to work on. If None, the default process group will be used. - **async_op** ([*bool*](https://docs.python.org/3/library/functions.html#bool)*,**optional*) - Whether this op should be an async op. Returns: Async work handle, if async_op is set to True. None, if not async_op or if not part of the group. Warning all_to_all_single is experimental and subject to change. Examples ``` >>> input = torch.arange(4) + rank * 4 >>> input tensor([0, 1, 2, 3]) # Rank 0 tensor([4, 5, 6, 7]) # Rank 1 tensor([8, 9, 10, 11]) # Rank 2 tensor([12, 13, 14, 15]) # Rank 3 >>> output = torch.empty([4], dtype=torch.int64) >>> dist.all_to_all_single(output, input) >>> output tensor([0, 4, 8, 12]) # Rank 0 tensor([1, 5, 9, 13]) # Rank 1 tensor([2, 6, 10, 14]) # Rank 2 tensor([3, 7, 11, 15]) # Rank 3 ``` ``` >>> # Essentially, it is similar to following operation: >>> 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) ``` ``` >>> # Another example with uneven split >>> input tensor([0, 1, 2, 3, 4, 5]) # Rank 0 tensor([10, 11, 12, 13, 14, 15, 16, 17, 18]) # Rank 1 tensor([20, 21, 22, 23, 24]) # Rank 2 tensor([30, 31, 32, 33, 34, 35, 36]) # Rank 3 >>> input_splits [2, 2, 1, 1] # Rank 0 [3, 2, 2, 2] # Rank 1 [2, 1, 1, 1] # Rank 2 [2, 2, 2, 1] # Rank 3 >>> output_splits [2, 3, 2, 2] # Rank 0 [2, 2, 1, 2] # Rank 1 [1, 2, 1, 2] # Rank 2 [1, 2, 1, 1] # Rank 3 >>> output = ... >>> dist.all_to_all_single(output, input, output_splits, input_splits) >>> output tensor([ 0, 1, 10, 11, 12, 20, 21, 30, 31]) # Rank 0 tensor([ 2, 3, 13, 14, 22, 32, 33]) # Rank 1 tensor([ 4, 15, 16, 23, 34, 35]) # Rank 2 tensor([ 5, 17, 18, 24, 36]) # Rank 3 ``` ``` >>> # Another example with tensors of torch.cfloat type. >>> 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]) # Rank 0 tensor([5+5j, 6+6j, 7+7j, 8+8j]) # Rank 1 tensor([9+9j, 10+10j, 11+11j, 12+12j]) # Rank 2 tensor([13+13j, 14+14j, 15+15j, 16+16j]) # Rank 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]) # Rank 0 tensor([2+2j, 6+6j, 10+10j, 14+14j]) # Rank 1 tensor([3+3j, 7+7j, 11+11j, 15+15j]) # Rank 2 tensor([4+4j, 8+8j, 12+12j, 16+16j]) # Rank 3 ``` torch.distributed.all_to_all(*output_tensor_list*, *input_tensor_list*, *group=None*, *async_op=False*)[[source]](https://github.com/pytorch/pytorch/blob/v2.13.0/torch/distributed/distributed_c10d.py#L5144) Scatters list of input tensors to all processes in a group and return gathered list of tensors in output list. Complex tensors are supported. Parameters: - **output_tensor_list** ([*list*](https://docs.python.org/3/library/stdtypes.html#list)*[*[*Tensor*](tensors.html#torch.Tensor)*]*) - List of tensors to be gathered one per rank. - **input_tensor_list** ([*list*](https://docs.python.org/3/library/stdtypes.html#list)*[*[*Tensor*](tensors.html#torch.Tensor)*]*) - List of tensors to scatter one per rank. - **group** ([*ProcessGroup*](distributed._dist2.html#torch.distributed._dist2.ProcessGroup)*,**optional*) - The process group to work on. If None, the default process group will be used. - **async_op** ([*bool*](https://docs.python.org/3/library/functions.html#bool)*,**optional*) - Whether this op should be an async op. Returns: Async work handle, if async_op is set to True. None, if not async_op or if not part of the group. Warning all_to_all is experimental and subject to change. Examples ``` >>> input = torch.arange(4) + rank * 4 >>> input = list(input.chunk(4)) >>> input [tensor([0]), tensor([1]), tensor([2]), tensor([3])] # Rank 0 [tensor([4]), tensor([5]), tensor([6]), tensor([7])] # Rank 1 [tensor([8]), tensor([9]), tensor([10]), tensor([11])] # Rank 2 [tensor([12]), tensor([13]), tensor([14]), tensor([15])] # Rank 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])] # Rank 0 [tensor([1]), tensor([5]), tensor([9]), tensor([13])] # Rank 1 [tensor([2]), tensor([6]), tensor([10]), tensor([14])] # Rank 2 [tensor([3]), tensor([7]), tensor([11]), tensor([15])] # Rank 3 ``` ``` >>> # Essentially, it is similar to following operation: >>> 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]) # Rank 0 tensor([10, 11, 12, 13, 14, 15, 16, 17, 18]) # Rank 1 tensor([20, 21, 22, 23, 24]) # Rank 2 tensor([30, 31, 32, 33, 34, 35, 36]) # Rank 3 >>> input_splits [2, 2, 1, 1] # Rank 0 [3, 2, 2, 2] # Rank 1 [2, 1, 1, 1] # Rank 2 [2, 2, 2, 1] # Rank 3 >>> output_splits [2, 3, 2, 2] # Rank 0 [2, 2, 1, 2] # Rank 1 [1, 2, 1, 2] # Rank 2 [1, 2, 1, 1] # Rank 3 >>> input = list(input.split(input_splits)) >>> input [tensor([0, 1]), tensor([2, 3]), tensor([4]), tensor([5])] # Rank 0 [tensor([10, 11, 12]), tensor([13, 14]), tensor([15, 16]), tensor([17, 18])] # Rank 1 [tensor([20, 21]), tensor([22]), tensor([23]), tensor([24])] # Rank 2 [tensor([30, 31]), tensor([32, 33]), tensor([34, 35]), tensor([36])] # Rank 3 >>> output = ... >>> dist.all_to_all(output, input) >>> output [tensor([0, 1]), tensor([10, 11, 12]), tensor([20, 21]), tensor([30, 31])] # Rank 0 [tensor([2, 3]), tensor([13, 14]), tensor([22]), tensor([32, 33])] # Rank 1 [tensor([4]), tensor([15, 16]), tensor([23]), tensor([34, 35])] # Rank 2 [tensor([5]), tensor([17, 18]), tensor([24]), tensor([36])] # Rank 3 ``` ``` >>> # Another example with tensors of torch.cfloat type. >>> 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])] # Rank 0 [tensor([5+5j]), tensor([6+6j]), tensor([7+7j]), tensor([8+8j])] # Rank 1 [tensor([9+9j]), tensor([10+10j]), tensor([11+11j]), tensor([12+12j])] # Rank 2 [tensor([13+13j]), tensor([14+14j]), tensor([15+15j]), tensor([16+16j])] # Rank 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])] # Rank 0 [tensor([2+2j]), tensor([6+6j]), tensor([10+10j]), tensor([14+14j])] # Rank 1 [tensor([3+3j]), tensor([7+7j]), tensor([11+11j]), tensor([15+15j])] # Rank 2 [tensor([4+4j]), tensor([8+8j]), tensor([12+12j]), tensor([16+16j])] # Rank 3 ``` torch.distributed.barrier(*group=None*, *async_op=False*, *device_ids=None*, *timeout=None*)[[source]](https://github.com/pytorch/pytorch/blob/v2.13.0/torch/distributed/distributed_c10d.py#L5283) Synchronize all processes. This collective blocks processes until the whole group enters this function, if async_op is False, or if async work handle is called on wait(). Parameters: - **group** ([*ProcessGroup*](distributed._dist2.html#torch.distributed._dist2.ProcessGroup)*,**optional*) - The process group to work on. If None, the default process group will be used. - **async_op** ([*bool*](https://docs.python.org/3/library/functions.html#bool)*,**optional*) - Whether this op should be an async op - **device_ids** (*[*[*int*](https://docs.python.org/3/library/functions.html#int)*]**,**optional*) - List of device/GPU ids. Only one id is expected. - **timeout** ([*datetime.timedelta*](https://docs.python.org/3/library/datetime.html#datetime.timedelta)*,**optional*) - Timeout for barrier. If `None`, the default process group timeout will be used. Returns: Async work handle, if async_op is set to True. None, if not async_op or if not part of the group Note ProcessGroupNCCL now blocks the cpu thread till the completion of the barrier collective. Note ProcessGroupNCCL implements barrier as an all_reduce of a 1-element tensor. A device must be chosen for allocating this tensor. The device choice is made by checking in this order (1) the first device passed to device_ids arg of barrier if not None, (2) the device passed to init_process_group if not None, (3) the device that was first used with this process group, if another collective with tensor inputs has been performed, (4) the device index indicated by the global rank mod local device count. torch.distributed.monitored_barrier(*group=None*, *timeout=None*, *wait_all_ranks=False*)[[source]](https://github.com/pytorch/pytorch/blob/v2.13.0/torch/distributed/distributed_c10d.py#L5360) Synchronize processes similar to `torch.distributed.barrier`, but consider a configurable timeout. It is able to report ranks that did not pass this barrier within the provided timeout. Specifically, for non-zero ranks, will block until a send/recv is processed from rank 0. Rank 0 will block until all send /recv from other ranks are processed, and will report failures for ranks that failed to respond in time. Note that if one rank does not reach the monitored_barrier (for example due to a hang), all other ranks would fail in monitored_barrier. This collective will block all processes/ranks in the group, until the whole group exits the function successfully, making it useful for debugging and synchronizing. However, it can have a performance impact and should only be used for debugging or scenarios that require full synchronization points on the host-side. For debugging purposes, this barrier can be inserted before the application's collective calls to check if any ranks are desynchronized. Note Note that this collective is only supported with the GLOO backend. Parameters: - **group** ([*ProcessGroup*](distributed._dist2.html#torch.distributed._dist2.ProcessGroup)*,**optional*) - The process group to work on. If `None`, the default process group will be used. - **timeout** ([*datetime.timedelta*](https://docs.python.org/3/library/datetime.html#datetime.timedelta)*,**optional*) - Timeout for monitored_barrier. If `None`, the default process group timeout will be used. - **wait_all_ranks** ([*bool*](https://docs.python.org/3/library/functions.html#bool)*,**optional*) - Whether to collect all failed ranks or not. By default, this is `False` and `monitored_barrier` on rank 0 will throw on the first failed rank it encounters in order to fail fast. By setting `wait_all_ranks=True` `monitored_barrier` will collect all failed ranks and throw an error containing information about all failed ranks. Returns: `None`. Example:: ``` >>> # Note: Process group initialization omitted on each rank. >>> import torch.distributed as dist >>> if dist.get_rank() != 1: >>> dist.monitored_barrier() # Raises exception indicating that >>> # rank 1 did not call into monitored_barrier. >>> # Example with wait_all_ranks=True >>> if dist.get_rank() == 0: >>> dist.monitored_barrier(wait_all_ranks=True) # Raises exception >>> # indicating that ranks 1, 2, ... world_size - 1 did not call into >>> # monitored_barrier. ``` *class*torch.distributed.Work A Work object represents the handle to a pending asynchronous operation in PyTorch's distributed package. It is returned by non-blocking collective operations, such as dist.all_reduce(tensor, async_op=True). block_current_stream(*self: torch._C._distributed_c10d.Work*) → [None](https://docs.python.org/3/library/constants.html#None) Blocks the currently active GPU stream on the operation to complete. For GPU based collectives this is equivalent to synchronize. For CPU initiated collectives such as with Gloo this will block the CUDA stream until the operation is complete. This returns immediately in all cases. To check whether an operation was successful you should check the Work object result asynchronously. boxed(*self: torch._C._distributed_c10d.Work*) → [object](https://docs.python.org/3/library/functions.html#object) exception(*self: torch._C._distributed_c10d.Work*) → [object](https://docs.python.org/3/library/functions.html#object) get_future(*self: torch._C._distributed_c10d.Work*) → torch.Future Returns: A `torch.futures.Future` object which is associated with the completion of the `Work`. As an example, a future object can be retrieved by `fut = process_group.allreduce(tensors).get_future()`. Example:: Below is an example of a simple allreduce DDP communication hook that uses `get_future` API to retrieve a Future associated with the completion of `allreduce`. ``` >>> 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) ``` Warning `get_future` API supports NCCL, and partially GLOO and MPI backends (no support for peer-to-peer operations like send/recv) and will return a `torch.futures.Future`. In the example above, `allreduce` work will be done on GPU using NCCL backend, `fut.wait()` will return after synchronizing the appropriate NCCL streams with PyTorch's current device streams to ensure we can have asynchronous CUDA execution and it does not wait for the entire operation to complete on GPU. Note that `CUDAFuture` does not support `TORCH_NCCL_BLOCKING_WAIT` flag or NCCL's `barrier()`. In addition, if a callback function was added by `fut.then()`, it will wait until `WorkNCCL`'s NCCL streams synchronize with `ProcessGroupNCCL`'s dedicated callback stream and invoke the callback inline after running the callback on the callback stream. `fut.then()` will return another `CUDAFuture` that holds the return value of the callback and a `CUDAEvent` that recorded the callback stream. > 1. For CPU work, `fut.done()` returns true when work has been completed and value() > tensors are ready. > 2. For GPU work, `fut.done()` returns true only whether the operation has been enqueued. > 3. For mixed CPU-GPU work (e.g. sending GPU tensors with GLOO), `fut.done()` returns > true when tensors have arrived on respective nodes, but not yet necessarily synched on > respective GPUs (similarly to GPU work). get_future_result(*self: torch._C._distributed_c10d.Work*) → torch.Future Returns: A `torch.futures.Future` object of int type which maps to the enum type of WorkResult As an example, a future object can be retrieved by `fut = process_group.allreduce(tensor).get_future_result()`. Example:: users can use `fut.wait()` to blocking wait for the completion of the work and get the WorkResult by `fut.value()`. Also, users can use `fut.then(call_back_func)` to register a callback function to be called when the work is completed, without blocking the current thread. Warning `get_future_result` API supports NCCL is_completed(*self: torch._C._distributed_c10d.Work*) → [bool](https://docs.python.org/3/library/functions.html#bool) is_success(*self: torch._C._distributed_c10d.Work*) → [bool](https://docs.python.org/3/library/functions.html#bool) result(*self: torch._C._distributed_c10d.Work*) → [list](https://docs.python.org/3/library/stdtypes.html#list)[[torch.Tensor](tensors.html#torch.Tensor)] source_rank(*self: torch._C._distributed_c10d.Work*) → [int](https://docs.python.org/3/library/functions.html#int) synchronize(*self: torch._C._distributed_c10d.Work*) → [None](https://docs.python.org/3/library/constants.html#None) *static*unbox(*arg0: [object](https://docs.python.org/3/library/functions.html#object)*) → torch._C._distributed_c10d.Work wait(*self: torch._C._distributed_c10d.Work*, *timeout: [datetime.timedelta](https://docs.python.org/3/library/datetime.html#datetime.timedelta) = datetime.timedelta(0)*) → [bool](https://docs.python.org/3/library/functions.html#bool) Returns: true/false. Example:: try: work.wait(timeout) except: # some handling Warning In normal cases, users do not need to set the timeout. calling wait() is the same as calling synchronize(): Letting the current stream block on the completion of the NCCL work. However, if timeout is set, it will block the CPU thread until the NCCL work is completed or timed out. If timeout, exception will be thrown. *class*torch.distributed.ReduceOp An enum-like class for available reduction operations: `SUM`, `PRODUCT`, `MIN`, `MAX`, `BAND`, `BOR`, `BXOR`, and `PREMUL_SUM`. `BAND`, `BOR`, and `BXOR` reductions are not available when using the `NCCL` backend. `AVG` divides values by the world size before summing across ranks. `AVG` is only available with the `NCCL` backend, and only for NCCL versions 2.10 or later. `PREMUL_SUM` multiplies inputs by a given scalar locally before reduction. `PREMUL_SUM` is available with the `NCCL` backend (NCCL versions 2.11 or later) and the `XCCL` backend. It can be used by calling `ReduceOp.PREMUL_SUM(factor)` where factor is a float or a single-element Tensor. Additionally, `MAX`, `MIN` and `PRODUCT` are not supported for complex tensors. The values of this class can be accessed as attributes, e.g., `ReduceOp.SUM`. They are used in specifying strategies for reduction collectives, e.g., `reduce()`. This class does not support `__members__` property. *class*RedOpType Members: SUM AVG PRODUCT MIN MAX BAND BOR BXOR PREMUL_SUM *property*name *property*factor The factor of the PREMUL_SUM ReduceOp. *class*torch.distributed.reduce_op Deprecated enum-like class for reduction operations: `SUM`, `PRODUCT`, `MIN`, and `MAX`. `ReduceOp` is recommended to use instead. ## Distributed Key-Value Store The distributed package comes with a distributed key-value store, which can be used to share information between processes in the group as well as to initialize the distributed package in `torch.distributed.init_process_group()` (by explicitly creating the store as an alternative to specifying `init_method`.) There are 3 choices for Key-Value Stores: `TCPStore`, `FileStore`, and `HashStore`. *class*torch.distributed.Store Base class for all store implementations, such as the 3 provided by PyTorch distributed: (`TCPStore`, `FileStore`, and `HashStore`). __init__(*self: torch._C._distributed_c10d.Store*) → [None](https://docs.python.org/3/library/constants.html#None) add(*self: torch._C._distributed_c10d.Store*, *arg0: [str](https://docs.python.org/3/library/stdtypes.html#str)*, *arg1: [SupportsInt](https://docs.python.org/3/library/typing.html#typing.SupportsInt) | [SupportsIndex](https://docs.python.org/3/library/typing.html#typing.SupportsIndex)*) → [int](https://docs.python.org/3/library/functions.html#int) The first call to add for a given `key` creates a counter associated with `key` in the store, initialized to `amount`. Subsequent calls to add with the same `key` increment the counter by the specified `amount`. Calling `add()` with a key that has already been set in the store by `set()` will result in an exception. Parameters: - **key** ([*str*](https://docs.python.org/3/library/stdtypes.html#str)) - The key in the store whose counter will be incremented. - **amount** ([*int*](https://docs.python.org/3/library/functions.html#int)) - The quantity by which the counter will be incremented. Example:: ``` >>> import torch.distributed as dist >>> from datetime import timedelta >>> # Using TCPStore as an example, other store types can also be used >>> store = dist.TCPStore("127.0.0.1", 0, 1, True, timedelta(seconds=30)) >>> store.add("first_key", 1) >>> store.add("first_key", 6) >>> # Should return 7 >>> store.get("first_key") ``` append(*self: torch._C._distributed_c10d.Store*, *arg0: [str](https://docs.python.org/3/library/stdtypes.html#str)*, *arg1: [str](https://docs.python.org/3/library/stdtypes.html#str)*) → [None](https://docs.python.org/3/library/constants.html#None) Append the key-value pair into the store based on the supplied `key` and `value`. If `key` does not exist in the store, it will be created. Parameters: - **key** ([*str*](https://docs.python.org/3/library/stdtypes.html#str)) - The key to be appended to the store. - **value** ([*str*](https://docs.python.org/3/library/stdtypes.html#str)) - The value associated with `key` to be added to the store. Example:: ``` >>> 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") >>> # Should return "potato" >>> store.get("first_key") ``` barrier(*self: torch._C._distributed_c10d.Store*, *key: [str](https://docs.python.org/3/library/stdtypes.html#str)*, *world_size: [SupportsInt](https://docs.python.org/3/library/typing.html#typing.SupportsInt) | [SupportsIndex](https://docs.python.org/3/library/typing.html#typing.SupportsIndex)*, *timeout: [datetime.timedelta](https://docs.python.org/3/library/datetime.html#datetime.timedelta) | [None](https://docs.python.org/3/library/constants.html#None) = None*) → [None](https://docs.python.org/3/library/constants.html#None) Barrier operation that blocks until `world_size` workers have called it with the same `key`. If `timeout` is not specified, the store's default timeout is used. Parameters: - **key** ([*str*](https://docs.python.org/3/library/stdtypes.html#str)) - The unique key for this barrier instance. - **world_size** ([*int*](https://docs.python.org/3/library/functions.html#int)) - The number of workers that must call barrier before it unblocks. - **timeout** (*timedelta**,**optional*) - Time to wait before throwing an exception. Defaults to store timeout. Example:: ``` >>> import torch.distributed as dist >>> from datetime import timedelta >>> store = dist.TCPStore("127.0.0.1", 0, 1, True, timedelta(seconds=30)) >>> # This will return immediately since world_size=1 >>> store.barrier("my_barrier", 1) >>> store.barrier("my_barrier2", 1, timedelta(seconds=10)) ``` check(*self: torch._C._distributed_c10d.Store*, *arg0: [collections.abc.Sequence](https://docs.python.org/3/library/collections.abc.html#collections.abc.Sequence)[[str](https://docs.python.org/3/library/stdtypes.html#str)]*) → [bool](https://docs.python.org/3/library/functions.html#bool) The call to check whether a given list of `keys` have value stored in the store. This call immediately returns in normal cases but still suffers from some edge deadlock cases, e.g, calling check after TCPStore has been destroyed. Calling `check()` with a list of keys that one wants to check whether stored in the store or not. Parameters: **keys** ([*list*](https://docs.python.org/3/library/stdtypes.html#list)*[*[*str*](https://docs.python.org/3/library/stdtypes.html#str)*]*) - The keys to query whether stored in the store. Example:: ``` >>> import torch.distributed as dist >>> from datetime import timedelta >>> # Using TCPStore as an example, other store types can also be used >>> store = dist.TCPStore("127.0.0.1", 0, 1, True, timedelta(seconds=30)) >>> store.add("first_key", 1) >>> # Should return 7 >>> store.check(["first_key"]) ``` clone(*self: torch._C._distributed_c10d.Store*) → torch._C._distributed_c10d.Store Clones the store and returns a new object that points to the same underlying store. The returned store can be used concurrently with the original object. This is intended to provide a safe way to use a store from multiple threads by cloning one store per thread. compare_set(*self: torch._C._distributed_c10d.Store*, *arg0: [str](https://docs.python.org/3/library/stdtypes.html#str)*, *arg1: [str](https://docs.python.org/3/library/stdtypes.html#str)*, *arg2: [str](https://docs.python.org/3/library/stdtypes.html#str)*) → [bytes](https://docs.python.org/3/library/stdtypes.html#bytes) Inserts the key-value pair into the store based on the supplied `key` and performs comparison between `expected_value` and `desired_value` before inserting. `desired_value` will only be set if `expected_value` for the `key` already exists in the store or if `expected_value` is an empty string. Parameters: - **key** ([*str*](https://docs.python.org/3/library/stdtypes.html#str)) - The key to be checked in the store. - **expected_value** ([*str*](https://docs.python.org/3/library/stdtypes.html#str)) - The value associated with `key` to be checked before insertion. - **desired_value** ([*str*](https://docs.python.org/3/library/stdtypes.html#str)) - The value associated with `key` to be added to the store. Example:: ``` >>> 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") >>> # Should return "second_value" >>> store.get("key") ``` delete_key(*self: torch._C._distributed_c10d.Store*, *arg0: [str](https://docs.python.org/3/library/stdtypes.html#str)*) → [bool](https://docs.python.org/3/library/functions.html#bool) Deletes the key-value pair associated with `key` from the store. Returns true if the key was successfully deleted, and false if it was not. Warning The `delete_key` API is only supported by the `TCPStore` and `HashStore`. Using this API with the `FileStore` will result in an exception. Parameters: **key** ([*str*](https://docs.python.org/3/library/stdtypes.html#str)) - The key to be deleted from the store Returns: True if `key` was deleted, otherwise False. Example:: ``` >>> import torch.distributed as dist >>> from datetime import timedelta >>> # Using TCPStore as an example, HashStore can also be used >>> store = dist.TCPStore("127.0.0.1", 0, 1, True, timedelta(seconds=30)) >>> store.set("first_key") >>> # This should return true >>> store.delete_key("first_key") >>> # This should return false >>> store.delete_key("bad_key") ``` get(*self: torch._C._distributed_c10d.Store*, *arg0: [str](https://docs.python.org/3/library/stdtypes.html#str)*) → [bytes](https://docs.python.org/3/library/stdtypes.html#bytes) Retrieves the value associated with the given `key` in the store. If `key` is not present in the store, the function will wait for `timeout`, which is defined when initializing the store, before throwing an exception. Parameters: **key** ([*str*](https://docs.python.org/3/library/stdtypes.html#str)) - The function will return the value associated with this key. Returns: Value associated with `key` if `key` is in the store. Example:: ``` >>> 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") >>> # Should return "first_value" >>> store.get("first_key") ``` has_extended_api(*self: torch._C._distributed_c10d.Store*) → [bool](https://docs.python.org/3/library/functions.html#bool) Returns true if the store supports extended operations. list_keys(*self: torch._C._distributed_c10d.Store*) → [list](https://docs.python.org/3/library/stdtypes.html#list)[[str](https://docs.python.org/3/library/stdtypes.html#str)] Returns a list of all keys in the store. multi_get(*self: torch._C._distributed_c10d.Store*, *arg0: [collections.abc.Sequence](https://docs.python.org/3/library/collections.abc.html#collections.abc.Sequence)[[str](https://docs.python.org/3/library/stdtypes.html#str)]*) → [list](https://docs.python.org/3/library/stdtypes.html#list)[[bytes](https://docs.python.org/3/library/stdtypes.html#bytes)] Retrieve all values in `keys`. If any key in `keys` is not present in the store, the function will wait for `timeout` Parameters: **keys** (*List**[*[*str*](https://docs.python.org/3/library/stdtypes.html#str)*]*) - The keys to be retrieved from the store. Example:: ``` >>> 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") >>> # Should return [b"po", b"tato"] >>> store.multi_get(["first_key", "second_key"]) ``` multi_set(*self: torch._C._distributed_c10d.Store*, *arg0: [collections.abc.Sequence](https://docs.python.org/3/library/collections.abc.html#collections.abc.Sequence)[[str](https://docs.python.org/3/library/stdtypes.html#str)]*, *arg1: [collections.abc.Sequence](https://docs.python.org/3/library/collections.abc.html#collections.abc.Sequence)[[str](https://docs.python.org/3/library/stdtypes.html#str)]*) → [None](https://docs.python.org/3/library/constants.html#None) Inserts a list key-value pair into the store based on the supplied `keys` and `values` Parameters: - **keys** (*List**[*[*str*](https://docs.python.org/3/library/stdtypes.html#str)*]*) - The keys to insert. - **values** (*List**[*[*str*](https://docs.python.org/3/library/stdtypes.html#str)*]*) - The values to insert. Example:: ``` >>> 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"]) >>> # Should return b"po" >>> store.get("first_key") ``` num_keys(*self: torch._C._distributed_c10d.Store*) → [int](https://docs.python.org/3/library/functions.html#int) Returns the number of keys set in the store. Note that this number will typically be one greater than the number of keys added by `set()` and `add()` since one key is used to coordinate all the workers using the store. Warning When used with the `TCPStore`, `num_keys` returns the number of keys written to the underlying file. If the store is destructed and another store is created with the same file, the original keys will be retained. Returns: The number of keys present in the store. Example:: ``` >>> import torch.distributed as dist >>> from datetime import timedelta >>> # Using TCPStore as an example, other store types can also be used >>> store = dist.TCPStore("127.0.0.1", 0, 1, True, timedelta(seconds=30)) >>> store.set("first_key", "first_value") >>> # This should return 2 >>> store.num_keys() ``` queue_len(*self: torch._C._distributed_c10d.Store*, *arg0: [str](https://docs.python.org/3/library/stdtypes.html#str)*) → [int](https://docs.python.org/3/library/functions.html#int) Returns the length of the specified queue. If the queue doesn't exist it returns 0. See queue_push for more details. Parameters: **key** ([*str*](https://docs.python.org/3/library/stdtypes.html#str)) - The key of the queue to get the length. queue_pop(*self: torch._C._distributed_c10d.Store*, *key: [str](https://docs.python.org/3/library/stdtypes.html#str)*, *block: [bool](https://docs.python.org/3/library/functions.html#bool) = True*) → [bytes](https://docs.python.org/3/library/stdtypes.html#bytes) Pops a value from the specified queue or waits until timeout if the queue is empty. See queue_push for more details. If block is False, a dist.QueueEmptyError will be raised if the queue is empty. Parameters: - **key** ([*str*](https://docs.python.org/3/library/stdtypes.html#str)) - The key of the queue to pop from. - **block** ([*bool*](https://docs.python.org/3/library/functions.html#bool)) - Whether to block waiting for the key or immediately return. queue_push(*self: torch._C._distributed_c10d.Store*, *arg0: [str](https://docs.python.org/3/library/stdtypes.html#str)*, *arg1: [str](https://docs.python.org/3/library/stdtypes.html#str)*) → [None](https://docs.python.org/3/library/constants.html#None) Pushes a value into the specified queue. Using the same key for queues and set/get operations may result in unexpected behavior. wait/check operations are supported for queues. wait with queues will only wake one waiting worker rather than all. Parameters: - **key** ([*str*](https://docs.python.org/3/library/stdtypes.html#str)) - The key of the queue to push to. - **value** ([*str*](https://docs.python.org/3/library/stdtypes.html#str)) - The value to push into the queue. set(*self: torch._C._distributed_c10d.Store*, *arg0: [str](https://docs.python.org/3/library/stdtypes.html#str)*, *arg1: [str](https://docs.python.org/3/library/stdtypes.html#str)*) → [None](https://docs.python.org/3/library/constants.html#None) Inserts the key-value pair into the store based on the supplied `key` and `value`. If `key` already exists in the store, it will overwrite the old value with the new supplied `value`. Parameters: - **key** ([*str*](https://docs.python.org/3/library/stdtypes.html#str)) - The key to be added to the store. - **value** ([*str*](https://docs.python.org/3/library/stdtypes.html#str)) - The value associated with `key` to be added to the store. Example:: ``` >>> 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") >>> # Should return "first_value" >>> store.get("first_key") ``` set_timeout(*self: torch._C._distributed_c10d.Store*, *arg0: [datetime.timedelta](https://docs.python.org/3/library/datetime.html#datetime.timedelta)*) → [None](https://docs.python.org/3/library/constants.html#None) Sets the store's default timeout. This timeout is used during initialization and in `wait()` and `get()`. Parameters: **timeout** (*timedelta*) - timeout to be set in the store. Example:: ``` >>> import torch.distributed as dist >>> from datetime import timedelta >>> # Using TCPStore as an example, other store types can also be used >>> store = dist.TCPStore("127.0.0.1", 0, 1, True, timedelta(seconds=30)) >>> store.set_timeout(timedelta(seconds=10)) >>> # This will throw an exception after 10 seconds >>> store.wait(["bad_key"]) ``` *property*timeout Gets the timeout of the store. wait(**args*, ***kwargs*) Overloaded function. 1. wait(self: torch._C._distributed_c10d.Store, arg0: collections.abc.Sequence[str]) -> None Waits for each key in `keys` to be added to the store. If not all keys are set before the `timeout` (set during store initialization), then `wait` will throw an exception. Parameters: **keys** ([*list*](https://docs.python.org/3/library/stdtypes.html#list)) - List of keys on which to wait until they are set in the store. Example:: ``` >>> import torch.distributed as dist >>> from datetime import timedelta >>> # Using TCPStore as an example, other store types can also be used >>> store = dist.TCPStore("127.0.0.1", 0, 1, True, timedelta(seconds=30)) >>> # This will throw an exception after 30 seconds >>> store.wait(["bad_key"]) ``` 1. wait(self: torch._C._distributed_c10d.Store, arg0: collections.abc.Sequence[str], arg1: datetime.timedelta) -> None Waits for each key in `keys` to be added to the store, and throws an exception if the keys have not been set by the supplied `timeout`. Parameters: - **keys** ([*list*](https://docs.python.org/3/library/stdtypes.html#list)) - List of keys on which to wait until they are set in the store. - **timeout** (*timedelta*) - Time to wait for the keys to be added before throwing an exception. Example:: ``` >>> import torch.distributed as dist >>> from datetime import timedelta >>> # Using TCPStore as an example, other store types can also be used >>> store = dist.TCPStore("127.0.0.1", 0, 1, True, timedelta(seconds=30)) >>> # This will throw an exception after 10 seconds >>> store.wait(["bad_key"], timedelta(seconds=10)) ``` *class*torch.distributed.TCPStore A TCP-based distributed key-value store implementation. The server store holds the data, while the client stores can connect to the server store over TCP and perform actions such as `set()` to insert a key-value pair, `get()` to retrieve a key-value pair, etc. There should always be one server store initialized because the client store(s) will wait for the server to establish a connection. Parameters: - **host_name** ([*str*](https://docs.python.org/3/library/stdtypes.html#str)) - The hostname or IP Address the server store should run on. - **port** ([*int*](https://docs.python.org/3/library/functions.html#int)) - The port on which the server store should listen for incoming requests. - **world_size** ([*int*](https://docs.python.org/3/library/functions.html#int)*,**optional*) - The total number of store users (number of clients + 1 for the server). Default is None (None indicates a non-fixed number of store users). - **is_master** ([*bool*](https://docs.python.org/3/library/functions.html#bool)*,**optional*) - True when initializing the server store and False for client stores. Default is False. - **timeout** (*timedelta**,**optional*) - Timeout used by the store during initialization and for methods such as `get()` and `wait()`. Default is timedelta(seconds=300) - **wait_for_workers** ([*bool*](https://docs.python.org/3/library/functions.html#bool)*,**optional*) - Whether to wait for all the workers to connect with the server store. This is only applicable when world_size is a fixed value. Default is True. - **multi_tenant** ([*bool*](https://docs.python.org/3/library/functions.html#bool)*,**optional*) - If True, all `TCPStore` instances in the current process with the same host/port will use the same underlying `TCPServer`. Default is False. - **master_listen_fd** ([*int*](https://docs.python.org/3/library/functions.html#int)*,**optional*) - If specified, the underlying `TCPServer` will listen on this file descriptor, which must be a socket already bound to `port`. To bind an ephemeral port we recommend setting the port to 0 and reading `.port`. Default is None (meaning the server creates a new socket and attempts to bind it to `port`). - **use_libuv** ([*bool*](https://docs.python.org/3/library/functions.html#bool)*,**optional*) - If True, use libuv for `TCPServer` backend. Default is True. Example:: ``` >>> import torch.distributed as dist >>> from datetime import timedelta >>> # Run on process 1 (server) >>> server_store = dist.TCPStore("127.0.0.1", 1234, 2, True, timedelta(seconds=30)) >>> # Run on process 2 (client) >>> client_store = dist.TCPStore("127.0.0.1", 1234, 2, False) >>> # Use any of the store methods from either the client or server after initialization >>> server_store.set("first_key", "first_value") >>> client_store.get("first_key") ``` __init__(*self: torch._C._distributed_c10d.TCPStore*, *host_name: [str](https://docs.python.org/3/library/stdtypes.html#str)*, *port: [SupportsInt](https://docs.python.org/3/library/typing.html#typing.SupportsInt) | [SupportsIndex](https://docs.python.org/3/library/typing.html#typing.SupportsIndex)*, *world_size: [SupportsInt](https://docs.python.org/3/library/typing.html#typing.SupportsInt) | [SupportsIndex](https://docs.python.org/3/library/typing.html#typing.SupportsIndex) | [None](https://docs.python.org/3/library/constants.html#None) = None*, *is_master: [bool](https://docs.python.org/3/library/functions.html#bool) = False*, *timeout: [datetime.timedelta](https://docs.python.org/3/library/datetime.html#datetime.timedelta) = datetime.timedelta(seconds=300)*, *wait_for_workers: [bool](https://docs.python.org/3/library/functions.html#bool) = True*, *multi_tenant: [bool](https://docs.python.org/3/library/functions.html#bool) = False*, *master_listen_fd: [SupportsInt](https://docs.python.org/3/library/typing.html#typing.SupportsInt) | [SupportsIndex](https://docs.python.org/3/library/typing.html#typing.SupportsIndex) | [None](https://docs.python.org/3/library/constants.html#None) = None*, *use_libuv: [bool](https://docs.python.org/3/library/functions.html#bool) = True*) → [None](https://docs.python.org/3/library/constants.html#None) Creates a new TCPStore. *property*host Gets the hostname on which the store listens for requests. *property*libuvBackend Returns True if it's using the libuv backend. *property*port Gets the port number on which the store listens for requests. *class*torch.distributed.HashStore A thread-safe store implementation based on an underlying hashmap. This store can be used within the same process (for example, by other threads), but cannot be used across processes. Example:: ``` >>> import torch.distributed as dist >>> store = dist.HashStore() >>> # store can be used from other threads >>> # Use any of the store methods after initialization >>> store.set("first_key", "first_value") ``` __init__(*self: torch._C._distributed_c10d.HashStore*) → [None](https://docs.python.org/3/library/constants.html#None) Creates a new HashStore. *class*torch.distributed.FileStore A store implementation that uses a file to store the underlying key-value pairs. Parameters: - **file_name** ([*str*](https://docs.python.org/3/library/stdtypes.html#str)) - path of the file in which to store the key-value pairs - **world_size** ([*int*](https://docs.python.org/3/library/functions.html#int)*,**optional*) - The total number of processes using the store. Default is -1 (a negative value indicates a non-fixed number of store users). Example:: ``` >>> import torch.distributed as dist >>> store1 = dist.FileStore("/tmp/filestore", 2) >>> store2 = dist.FileStore("/tmp/filestore", 2) >>> # Use any of the store methods from either the client or server after initialization >>> store1.set("first_key", "first_value") >>> store2.get("first_key") ``` __init__(*self: torch._C._distributed_c10d.FileStore*, *file_name: [str](https://docs.python.org/3/library/stdtypes.html#str)*, *world_size: [SupportsInt](https://docs.python.org/3/library/typing.html#typing.SupportsInt) | [SupportsIndex](https://docs.python.org/3/library/typing.html#typing.SupportsIndex) = -1*) → [None](https://docs.python.org/3/library/constants.html#None) Creates a new FileStore. *property*path Gets the path of the file used by FileStore to store key-value pairs. *class*torch.distributed.PrefixStore A wrapper around any of the 3 key-value stores (`TCPStore`, `FileStore`, and `HashStore`) that adds a prefix to each key inserted to the store. Parameters: - **prefix** ([*str*](https://docs.python.org/3/library/stdtypes.html#str)) - The prefix string that is prepended to each key before being inserted into the store. - **store** (*torch.distributed.store*) - A store object that forms the underlying key-value store. __init__(*self: torch._C._distributed_c10d.PrefixStore*, *prefix: [str](https://docs.python.org/3/library/stdtypes.html#str)*, *store: torch._C._distributed_c10d.Store*) → [None](https://docs.python.org/3/library/constants.html#None) Creates a new PrefixStore. *property*underlying_store Gets the underlying store object that PrefixStore wraps around. ## Profiling Collective Communication Note that you can use `torch.profiler` (recommended, only available after 1.8.1) or `torch.autograd.profiler` to profile collective communication and point-to-point communication APIs mentioned here. All out-of-the-box backends (`gloo`, `nccl`, `mpi`) are supported and collective communication usage will be rendered as expected in profiling output/traces. Profiling your code is the same as any regular torch operator: ``` import torch import torch.distributed as dist with torch.profiler(): tensor = torch.randn(20, 10) dist.all_reduce(tensor) ``` Please refer to the [profiler documentation](https://pytorch.org/docs/main/profiler.html) for a full overview of profiler features. torch.distributed.distributed_c10d.record_comm(*name*)[[source]](https://github.com/pytorch/pytorch/blob/v2.13.0/torch/distributed/distributed_c10d.py#L6886) Context manager to set a custom profiling name for communication collectives. When active, all c10d collectives issued within this context will use `name` as their profiling title in the Work base class, overriding the default backend-specific name (e.g. `nccl:all_reduce`). This works across all backends without per-backend or per-collective changes. Parameters: **name** ([*str*](https://docs.python.org/3/library/stdtypes.html#str)) - The profiling name to associate with collectives. Example:: ``` >>> with dist.record_comm("FSDP::all_gather (layer1)"): ... dist.all_gather_into_tensor(output, input, group=pg) ``` ## Optimization with Symmetric Memory ### Copy Engine Collectives When NCCL collective operations are performed on symmetric memory tensors with the zero-CTA policy, data movement is offloaded to the GPU's copy engines (DMA engines) instead of using CUDA streaming multiprocessors (SMs). This frees up SMs for compute work, enabling better overlap of communication and computation. For setup instructions, requirements, and examples, see [Copy Engine Collectives](symmetric_memory.html#copy-engine-collectives) in the Symmetric Memory documentation. ### Higher-Precision Reduction When NCCL collectives such as `reduce_scatter` and `all_reduce` operate on symmetric memory tensors, NCCL's symmetric kernel implementation automatically performs internal reduction with higher precision (e.g., BF16/FP16 in → FP32 accumulate → BF16/FP16 out). This improves numerical accuracy without any code changes to the collective call. For details on scope, supported domains, and version requirements, see [Higher-Precision Reduction](symmetric_memory.html#higher-precision-reduction) in the Symmetric Memory documentation. ## Multi-GPU collective functions Warning The multi-GPU functions (which stand for multiple GPUs per CPU thread) are deprecated. As of today, PyTorch Distributed's preferred programming model is one device per thread, as exemplified by the APIs in this document. If you are a backend developer and want to support multiple devices per thread, please contact PyTorch Distributed's maintainers. ## Object collectives Warning Object collectives have a number of serious limitations. Read further to determine if they are safe to use for your use case. Object collectives are a set of collective-like operations that work on arbitrary Python objects, as long as they can be pickled. There are various collective patterns implemented (e.g. broadcast, all_gather, ...) but they each roughly follow this pattern: 1. convert the input object into a pickle (raw bytes), then shove it into a byte tensor 2. communicate the size of this byte tensor to peers (first collective operation) 3. allocate appropriately sized tensor to perform the real collective 4. communicate the object data (second collective operation) 5. convert raw data back into Python (unpickle) Object collectives sometimes have surprising performance or memory characteristics that lead to long runtimes or OOMs, and thus they should be used with caution. Here are some common issues. **Asymmetric pickle/unpickle time** - Pickling objects can be slow, depending on the number, type and size of the objects. When the collective has a fan-in (e.g. gather_object), the receiving rank(s) must unpickle N times more objects than the sending rank(s) had to pickle, which can cause other ranks to time out on their next collective. **Inefficient tensor communication** - Tensors should be sent via regular collective APIs, not object collective APIs. It is possible to send Tensors via object collective APIs, but they will be serialized and deserialized (including a CPU-sync and device-to-host copy in the case of non-CPU tensors), and in almost every case other than debugging or troubleshooting code, it would be worth the trouble to refactor the code to use non-object collectives instead. **Unexpected tensor devices** - If you still want to send tensors via object collectives, there is another aspect specific to cuda (and possibly other accelerators) tensors. If you pickle a tensor that is currently on `cuda:3`, and then unpickle it, you will get another tensor on `cuda:3` *regardless of which process you are on, or which CUDA device is the 'default' device for that process*. With regular tensor collective APIs, 'output tensors' will always be on the same, local device, which is generally what you'd expect. Unpickling a tensor will implicitly activate a CUDA context if it is the first time a GPU is used by the process, which can waste significant amounts of GPU memory. This issue can be avoided by moving tensors to CPU before passing them as inputs to an object collective. ## Third-party backends Besides the builtin GLOO/MPI/NCCL backends, PyTorch distributed supports third-party backends through a run-time register mechanism. For references on how to develop a third-party backend through C++ Extension, please refer to [Tutorials - Custom C++ and CUDA Extensions](https://pytorch.org/tutorials/advanced/cpp_extension.html) and `test/cpp_extensions/cpp_c10d_extension.cpp`. The capability of third-party backends are decided by their own implementations. The new backend derives from `c10d::ProcessGroup` and registers the backend name and the instantiating interface through `torch.distributed.Backend.register_backend()` when imported. When manually importing this backend and invoking `torch.distributed.init_process_group()` with the corresponding backend name, the `torch.distributed` package runs on the new backend. Warning The support of third-party backend is experimental and subject to change. ## TorchComms backend [TorchComms](https://github.com/meta-pytorch/torchcomms) is an optional communication backend for `torch.distributed`. When enabled, it overrides the normal backend instantiation in `init_process_group()` so that all process groups are created through TorchComms instead of the built-in `ProcessGroup` implementations. Note TorchComms is experimental and must be installed separately. The `torchcomms` package must be importable for the flags below to take effect. ### Enabling TorchComms Set the `TORCH_DISTRIBUTED_USE_TORCHCOMMS` environment variable before calling `init_process_group()`: ``` export TORCH_DISTRIBUTED_USE_TORCHCOMMS=1 ``` Or set the config flag programmatically: ``` import torch.distributed.config as dist_config dist_config.use_torchcomms = True ``` The `backend` argument to `init_process_group()` (e.g. `"nccl"`, `"gloo"`) is still respected - it is forwarded to TorchComms, which selects the corresponding vendor plugin. No other application code changes are required; all `torch.distributed` collective APIs continue to work as before. ### Behavior when enabled When TorchComms is enabled, `init_process_group()` changes its backend instantiation path for every device/backend pair in the process group (except the `fake` backend, which is always handled natively): 1. A TorchComms communicator is created via `torchcomms.new_comm()` using the requested backend string and device. 2. The communicator is wrapped in a `_BackendWrapper` that implements the `c10d::Backend` C++ interface, making it a drop-in replacement for the native `ProcessGroup` backends. 3. A `FlightRecorderHook` is automatically registered on the communicator. The hook respects the `TORCH_FR_BUFFER_SIZE` and `TORCH_NCCL_TRACE_BUFFER_SIZE` environment variables for configuring the trace buffer size. 4. `destroy_process_group()` calls `finalize()` on TorchComms communicators during cleanup. 5. `split_group()` creates sub-communicators through TorchComms' native splitting rather than constructing a new process group from scratch. ### Eager initialization TorchComms communicators are eagerly initialized during `init_process_group()` and only support a single backend device per group. The `device_id` argument must be specified at initialization time: ``` dist.init_process_group(backend="nccl", device_id=torch.device("cuda", local_rank)) ``` ### Point-to-point operation concurrency Each TorchComms process group maps 1:1 to a single underlying communicator. Point-to-point operations (`send`/`recv`) issued on the same group and stream are **not guaranteed to run concurrently**. Code that relies on concurrent point-to-point operations must either: - Use the batched P2P APIs (`batch_isend_irecv()`), or - Issue the operations on separate groups or communicators. ## Launch utility The `torch.distributed` package also provides a launch utility in `torch.distributed.launch`. This helper utility can be used to launch multiple processes per node for distributed training. Module `torch.distributed.launch`. `torch.distributed.launch` is a module that spawns up multiple distributed training processes on each of the training nodes. Warning This module is going to be deprecated in favor of [torchrun](elastic/run.html#launcher-api). The utility can be used for single-node distributed training, in which one or more processes per node will be spawned. The utility can be used for either CPU training or GPU training. If the utility is used for GPU training, each distributed process will be operating on a single GPU. This can achieve well-improved single-node training performance. It can also be used in multi-node distributed training, by spawning up multiple processes on each node for well-improved multi-node distributed training performance as well. This will especially be beneficial for systems with multiple Infiniband interfaces that have direct-GPU support, since all of them can be utilized for aggregated communication bandwidth. In both cases of single-node distributed training or multi-node distributed training, this utility will launch the given number of processes per node (`--nproc-per-node`). If used for GPU training, this number needs to be less or equal to the number of GPUs on the current system (`nproc_per_node`), and each process will be operating on a single GPU from *GPU 0 to GPU (nproc_per_node - 1)*. **How to use this module:** 1. Single-Node multi-process distributed training ``` 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) ``` 1. Multi-Node multi-process distributed training: (e.g. two nodes) Node 1: *(IP: 192.168.1.1, and has a free port: 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) ``` Node 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) ``` 1. To look up what optional arguments this module offers: ``` python -m torch.distributed.launch --help ``` **Important Notices:** 1. This utility and multi-process distributed (single-node or multi-node) GPU training currently only achieves the best performance using the NCCL distributed backend. Thus NCCL backend is the recommended backend to use for GPU training. 2. In your training program, you must parse the command-line argument: `--local-rank=LOCAL_PROCESS_RANK`, which will be provided by this module. If your training program uses GPUs, you should ensure that your code only runs on the GPU device of LOCAL_PROCESS_RANK. This can be done by: Parsing the local_rank argument ``` >>> import argparse >>> parser = argparse.ArgumentParser() >>> parser.add_argument("--local-rank", "--local_rank", type=int) >>> args = parser.parse_args() ``` Set your device to local rank using either ``` >>> torch.cuda.set_device(args.local_rank) # before your code runs ``` or ``` >>> with torch.cuda.device(args.local_rank): >>> # your code to run >>> ... ``` Changed in version 2.0.0: The launcher will passes the `--local-rank=` argument to your script. From PyTorch 2.0.0 onwards, the dashed `--local-rank` is preferred over the previously used underscored `--local_rank`. For backward compatibility, it may be necessary for users to handle both cases in their argument parsing code. This means including both `"--local-rank"` and `"--local_rank"` in the argument parser. If only `"--local_rank"` is provided, the launcher will trigger an error: "error: unrecognized arguments: -local-rank=". For training code that only supports PyTorch 2.0.0+, including `"--local-rank"` should be sufficient. 3. In your training program, you are supposed to call the following function at the beginning to start the distributed backend. It is strongly recommended that `init_method=env://`. Other init methods (e.g. `tcp://`) may work, but `env://` is the one that is officially supported by this module. ``` >>> torch.distributed.init_process_group(backend='YOUR BACKEND', >>> init_method='env://') ``` 4. In your training program, you can either use regular distributed functions or use [`torch.nn.parallel.DistributedDataParallel()`](generated/torch.nn.parallel.DistributedDataParallel.html#torch.nn.parallel.DistributedDataParallel) module. If your training program uses GPUs for training and you would like to use [`torch.nn.parallel.DistributedDataParallel()`](generated/torch.nn.parallel.DistributedDataParallel.html#torch.nn.parallel.DistributedDataParallel) module, here is how to configure it. ``` >>> model = torch.nn.parallel.DistributedDataParallel(model, >>> device_ids=[args.local_rank], >>> output_device=args.local_rank) ``` Please ensure that `device_ids` argument is set to be the only GPU device id that your code will be operating on. This is generally the local rank of the process. In other words, the `device_ids` needs to be `[args.local_rank]`, and `output_device` needs to be `args.local_rank` in order to use this utility 5. Another way to pass `local_rank` to the subprocesses via environment variable `LOCAL_RANK`. This behavior is enabled when you launch the script with `--use-env=True`. You must adjust the subprocess example above to replace `args.local_rank` with `os.environ['LOCAL_RANK']`; the launcher will not pass `--local-rank` when you specify this flag. Warning `local_rank` is NOT globally unique: it is only unique per process on a machine. Thus, don't use it to decide if you should, e.g., write to a networked filesystem. See [pytorch/pytorch#12042](https://github.com/pytorch/pytorch/issues/12042) for an example of how things can go wrong if you don't do this correctly. torch.distributed.launch.launch(*args*)[[source]](https://github.com/pytorch/pytorch/blob/v2.13.0/torch/distributed/launch.py#L183) torch.distributed.launch.main(*args=None*)[[source]](https://github.com/pytorch/pytorch/blob/v2.13.0/torch/distributed/launch.py#L191) torch.distributed.launch.parse_args(*args*)[[source]](https://github.com/pytorch/pytorch/blob/v2.13.0/torch/distributed/launch.py#L168) ## Spawn utility The [Multiprocessing package - torch.multiprocessing](multiprocessing.html#multiprocessing-doc) package also provides a `spawn` function in [`torch.multiprocessing.spawn()`](multiprocessing.html#module-torch.multiprocessing.spawn). This helper function can be used to spawn multiple processes. It works by passing in the function that you want to run and spawns N processes to run it. This can be used for multiprocess distributed training as well. For references on how to use it, please refer to [PyTorch example - ImageNet implementation](https://github.com/pytorch/examples/tree/master/imagenet) Note that this function requires Python 3.4 or higher. ## Debugging `torch.distributed` applications Debugging distributed applications can be challenging due to hard to understand hangs, crashes, or inconsistent behavior across ranks. `torch.distributed` provides a suite of tools to help debug training applications in a self-serve fashion: ### Python Breakpoint It is extremely convenient to use python's debugger in a distributed environment, but because it does not work out of the box many people do not use it at all. PyTorch offers a customized wrapper around pdb that streamlines the process. `torch.distributed.breakpoint` makes this process easy. Internally, it customizes `pdb`'s breakpoint behavior in two ways but otherwise behaves as normal `pdb`. 1. Attaches the debugger only on one rank (specified by the user). 2. Ensures all other ranks stop, by using a `torch.distributed.barrier()` that will release once the debugged rank issues a `continue` 3. Reroutes stdin from the child process such that it connects to your terminal. To use it, simply issue `torch.distributed.breakpoint(rank)` on all ranks, using the same value for `rank` in each case. ### Monitored Barrier As of v1.10, `torch.distributed.monitored_barrier()` exists as an alternative to `torch.distributed.barrier()` which fails with helpful information about which rank may be faulty when crashing, i.e. not all ranks calling into `torch.distributed.monitored_barrier()` within the provided timeout. `torch.distributed.monitored_barrier()` implements a host-side barrier using `send`/`recv` communication primitives in a process similar to acknowledgements, allowing rank 0 to report which rank(s) failed to acknowledge the barrier in time. As an example, consider the following function where rank 1 fails to call into `torch.distributed.monitored_barrier()` (in practice this could be due to an application bug or hang in a previous collective): ``` 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 requires gloo process group to perform host-side sync. 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=()) ``` The following error message is produced on rank 0, allowing the user to determine which rank(s) may be faulty and investigate further: ``` 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` With `TORCH_CPP_LOG_LEVEL=INFO`, the environment variable `TORCH_DISTRIBUTED_DEBUG` can be used to trigger additional useful logging and collective synchronization checks to ensure all ranks are synchronized appropriately. `TORCH_DISTRIBUTED_DEBUG` can be set to either `OFF` (default), `INFO`, or `DETAIL` depending on the debugging level required. Please note that the most verbose option, `DETAIL` may impact the application performance and thus should only be used when debugging issues. Setting `TORCH_DISTRIBUTED_DEBUG=INFO` will result in additional debug logging when models trained with [`torch.nn.parallel.DistributedDataParallel()`](generated/torch.nn.parallel.DistributedDataParallel.html#torch.nn.parallel.DistributedDataParallel) are initialized, and `TORCH_DISTRIBUTED_DEBUG=DETAIL` will additionally log runtime performance statistics a select number of iterations. These runtime statistics include data such as forward time, backward time, gradient communication time, etc. As an example, given the following application: ``` 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" # set to DETAIL for runtime logging. mp.spawn(worker, nprocs=2, args=()) ``` The following logs are rendered at initialization time: ``` 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 ``` The following logs are rendered during runtime (when `TORCH_DISTRIBUTED_DEBUG=DETAIL` is set): ``` 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 ``` In addition, `TORCH_DISTRIBUTED_DEBUG=INFO` enhances crash logging in [`torch.nn.parallel.DistributedDataParallel()`](generated/torch.nn.parallel.DistributedDataParallel.html#torch.nn.parallel.DistributedDataParallel) due to unused parameters in the model. Currently, `find_unused_parameters=True` must be passed into [`torch.nn.parallel.DistributedDataParallel()`](generated/torch.nn.parallel.DistributedDataParallel.html#torch.nn.parallel.DistributedDataParallel) initialization if there are parameters that may be unused in the forward pass, and as of v1.10, all model outputs are required to be used in loss computation as [`torch.nn.parallel.DistributedDataParallel()`](generated/torch.nn.parallel.DistributedDataParallel.html#torch.nn.parallel.DistributedDataParallel) does not support unused parameters in the backwards pass. These constraints are challenging especially for larger models, thus when crashing with an error, [`torch.nn.parallel.DistributedDataParallel()`](generated/torch.nn.parallel.DistributedDataParallel.html#torch.nn.parallel.DistributedDataParallel) will log the fully qualified name of all parameters that went unused. For example, in the above application, if we modify `loss` to be instead computed as `loss = output[1]`, then `TwoLinLayerNet.a` does not receive a gradient in the backwards pass, and thus results in `DDP` failing. On a crash, the user is passed information about parameters which went unused, which may be challenging to manually find for large models: ``` 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). Parameters which did not receive grad for rank 0: a.weight Parameter indices which did not receive grad for rank 0: 0 ``` Setting `TORCH_DISTRIBUTED_DEBUG=DETAIL` will trigger additional consistency and synchronization checks on every collective call issued by the user either directly or indirectly (such as DDP `allreduce`). This is done by creating a wrapper process group that wraps all process groups returned by `torch.distributed.init_process_group()` and `torch.distributed.new_group()` APIs. As a result, these APIs will return a wrapper process group that can be used exactly like a regular process group, but performs consistency checks before dispatching the collective to an underlying process group. Currently, these checks include a `torch.distributed.monitored_barrier()`, which ensures all ranks complete their outstanding collective calls and reports ranks which are stuck. Next, the collective itself is checked for consistency by ensuring all collective functions match and are called with consistent tensor shapes. If this is not the case, a detailed error report is included when the application crashes, rather than a hang or uninformative error message. As an example, consider the following function which has mismatched input shapes into `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=()) ``` With the `NCCL` backend, such an application would likely result in a hang which can be challenging to root-cause in nontrivial scenarios. If the user enables `TORCH_DISTRIBUTED_DEBUG=DETAIL` and reruns the application, the following error message reveals the root cause: ``` 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} ] ``` Note For fine-grained control of the debug level during runtime the functions `torch.distributed.set_debug_level()`, `torch.distributed.set_debug_level_from_env()`, and `torch.distributed.get_debug_level()` can also be used. In addition, `TORCH_DISTRIBUTED_DEBUG=DETAIL` can be used in conjunction with `TORCH_SHOW_CPP_STACKTRACES=1` to log the entire callstack when a collective desynchronization is detected. These collective desynchronization checks will work for all applications that use `c10d` collective calls backed by process groups created with the `torch.distributed.init_process_group()` and `torch.distributed.new_group()` APIs. ### torch.distributed.debug HTTP Server The `torch.distributed.debug` module provides a HTTP server that can be used to debug distributed applications. The server can be started by calling `torch.distributed.debug.start_debug_server()`. This allows users to collect data across all workers at runtime. torch.distributed.debug.start_debug_server(*port=25999*, *worker_port=0*, *start_method=None*, *dump_dir=None*, *dump_interval=60.0*, *enabled_dumps=None*, *handlers=None*, *fetch_timeout=60.0*)[[source]](https://github.com/pytorch/pytorch/blob/v2.13.0/torch/distributed/debug/__init__.py#L27) Start the debug server stack on all workers. The frontend debug server is only started on rank0 while the per rank worker servers are started on all ranks. This server provides an HTTP frontend that allows for debugging slow and deadlocked distributed jobs across all ranks simultaneously. This collects data such as stack traces, FlightRecorder events, and performance profiles. This depends on dependencies which are not installed by default. Dependencies: - Jinja2 - aiohttp WARNING: This is intended to only be used in trusted network environments. The debug server is not designed to be secure and should not be exposed to the public internet. See SECURITY.md for more details. WARNING: This is an experimental feature and may change at any time. Parameters: - **port** ([*int*](https://docs.python.org/3/library/functions.html#int)) - The port to start the frontend debug server on. - **worker_port** ([*int*](https://docs.python.org/3/library/functions.html#int)) - The port to start the worker server on. Defaults to 0, which will cause the worker server to bind to an ephemeral port. - **start_method** ([*str*](https://docs.python.org/3/library/stdtypes.html#str)*|**None*) - The multiprocessing start method to use for the frontend server process. One of "fork", "spawn", or "forkserver". If None, uses the default start method. Using "spawn" is recommended when using CUDA or when fork safety is a concern. - **dump_dir** ([*str*](https://docs.python.org/3/library/stdtypes.html#str)*|**None*) - Directory to write periodic debug dumps to. If None, periodic dumping is disabled. - **dump_interval** ([*float*](https://docs.python.org/3/library/functions.html#float)) - Seconds between periodic dumps. Defaults to 60. - **enabled_dumps** ([*set*](https://docs.python.org/3/library/stdtypes.html#set)*[*[*str*](https://docs.python.org/3/library/stdtypes.html#str)*]**|**None*) - Set of handler dump filenames to enable (e.g. {"stacks", "fr_trace", "tcpstore"}). If None, all handlers that implement dump() are enabled. - **handlers** ([*list*](https://docs.python.org/3/library/stdtypes.html#list)*[**DebugHandler**]**|**None*) - List of debug handlers to use. If None, uses the default handlers. See torch.distributed.debug._handlers for the default handlers. - **fetch_timeout** ([*float*](https://docs.python.org/3/library/functions.html#float)) - Timeout in seconds for fetching data from individual workers. Defaults to 60. Workers that don't respond within this time will be reported as unavailable. torch.distributed.debug.stop_debug_server()[[source]](https://github.com/pytorch/pytorch/blob/v2.13.0/torch/distributed/debug/__init__.py#L130) Shutdown the debug server and stop the frontend debug server process. ## Logging In addition to explicit debugging support via `torch.distributed.monitored_barrier()` and `TORCH_DISTRIBUTED_DEBUG`, the underlying C++ library of `torch.distributed` also outputs log messages at various levels. These messages can be helpful to understand the execution state of a distributed training job and to troubleshoot problems such as network connection failures. The following matrix shows how the log level can be adjusted via the combination of `TORCH_CPP_LOG_LEVEL` and `TORCH_DISTRIBUTED_DEBUG` environment variables. | `TORCH_CPP_LOG_LEVEL` | `TORCH_DISTRIBUTED_DEBUG` | Effective Log Level | | --- | --- | --- | | `ERROR` | ignored | Error | | `WARNING` | ignored | Warning | | `INFO` | ignored | Info | | `INFO` | `INFO` | Debug | | `INFO` | `DETAIL` | Trace (a.k.a. All) | Distributed components raise custom Exception types derived from `RuntimeError`: - `torch.distributed.DistError`: This is the base type of all distributed exceptions. - `torch.distributed.DistBackendError`: This exception is thrown when a backend-specific error occurs. For example, if the `NCCL` backend is used and the user attempts to use a GPU that is not available to the `NCCL` library. - `torch.distributed.DistNetworkError`: This exception is thrown when networking libraries encounter errors (ex: Connection reset by peer) - `torch.distributed.DistStoreError`: This exception is thrown when the Store encounters an error (ex: TCPStore timeout) *class*torch.distributed.DistError Exception raised when an error occurs in the distributed library *class*torch.distributed.DistBackendError Exception raised when a backend error occurs in distributed *class*torch.distributed.DistNetworkError Exception raised when a network error occurs in distributed *class*torch.distributed.DistStoreError Exception raised when an error occurs in the distributed store If you are running single node training, it may be convenient to interactively breakpoint your script. We offer a way to conveniently breakpoint a single rank: torch.distributed.breakpoint(*rank=0*, *skip=0*, *timeout_s=3600*)[[source]](https://github.com/pytorch/pytorch/blob/v2.13.0/torch/distributed/__init__.py#L101) Set a breakpoint, but only on a single rank. All other ranks will wait for you to be done with the breakpoint before continuing. Parameters: - **rank** ([*int*](https://docs.python.org/3/library/functions.html#int)) - Which rank to break on. Default: `0` - **skip** ([*int*](https://docs.python.org/3/library/functions.html#int)) - Skip the first `skip` calls to this breakpoint. Default: `0`. torch.distributed.collective_utils.all_gather_object_enforce_type(*pg*, *object_list*, *obj*, *type_checker=>*)[[source]](https://github.com/pytorch/pytorch/blob/v2.13.0/torch/distributed/collective_utils.py#L205) Similar to plain all_gather_object but with additional type checking AFTER gather is done to ensure basic consistency. If check does not pass, all ranks will fail with exception. This is generally to prevent conditional logic leading to unexpected messages being received. This is considered fatal code error, but due to logic stacks this might happen implicitly in practice. The default check does not check sub type (considered different) or covariance (considered same) but users can pass in custom checker if more complicated check is needed. torch.distributed.launcher.api.launch_agent(*config*, *entrypoint*, *args*, *health_check_server=None*)[[source]](https://github.com/pytorch/pytorch/blob/v2.13.0/torch/distributed/launcher/api.py#L241) Return type: [dict](https://docs.python.org/3/library/stdtypes.html#dict)[[int](https://docs.python.org/3/library/functions.html#int), [*Any*](https://docs.python.org/3/library/typing.html#typing.Any)]