Scatters tensor across multiple GPUs.
Note
Exactly one of devices and out must be specified. When out is specified, chunk_sizes must not be specified and will be inferred from sizes of out.
a tuple containing chunks of tensor, placed on devices.
a tuple containing out tensors, each containing a chunk of tensor.


Returns cublasHandle_t pointer to current cuBLAS handle


Returns the index of a currently selected device.


Returns the currently selected Stream for a given device.
device (torch.device or int, optional) – selected device. Returns the currently selected Stream for the current device, given by current_device(), if device is None (default).


Returns the default Stream for a given device.
device (torch.device or int, optional) – selected device. Returns the default Stream for the current device, given by current_device(), if device is None (default).


Context-manager that changes the selected device.
device (torch.device or int) – device index to select. It’s a no-op if this argument is a negative integer or None.


Context-manager that changes the current device to that of given object.
You can use both tensors and storages as arguments. If a given object is not allocated on a GPU, this is a no-op.
obj (Tensor or Storage) – object allocated on the selected device.


Releases all unoccupied cached memory currently held by the caching allocator so that those can be used in other GPU application and visible in nvidia-smi.
Note
empty_cache() doesn’t increase the amount of GPU memory available for PyTorch. However, it may help reduce fragmentation of GPU memory in certain cases. See Memory management for more details about GPU memory management.


Wrapper around a CUDA event.
CUDA events are synchronization markers that can be used to monitor the device’s progress, to accurately measure timing, and to synchronize CUDA streams.
The underlying CUDA events are lazily initialized when the event is first recorded or exported to another process. After creation, only streams on the same device may record the event. However, streams on any device can wait on the event.
Returns the time elapsed in milliseconds after the event was recorded and before the end_event was recorded.
Reconstruct an event from an IPC handle on the given device.
Returns an IPC handle of this event. If not recorded yet, the event will use the current device.
Checks if all work currently captured by event has completed.
A boolean indicating if all work currently captured by event has completed.
Records the event in a given stream.
Uses torch.cuda.current_stream() if no stream is specified. The stream’s device must match the event’s device.
Waits for the event to complete.
Waits until the completion of all work currently captured in this event. This prevents the CPU thread from proceeding until the event completes.
Note
This is a wrapper around cudaEventSynchronize(): see CUDA Event documentation for more info.
Makes all future work submitted to the given stream wait for this event.
Use torch.cuda.current_stream() if no stream is specified.


Returns the time elapsed in milliseconds after the event was recorded and before the end_event was recorded.


Reconstruct an event from an IPC handle on the given device.


Returns an IPC handle of this event. If not recorded yet, the event will use the current device.


Checks if all work currently captured by event has completed.
A boolean indicating if all work currently captured by event has completed.


Records the event in a given stream.
Uses torch.cuda.current_stream() if no stream is specified. The stream’s device must match the event’s device.


Waits for the event to complete.
Waits until the completion of all work currently captured in this event. This prevents the CPU thread from proceeding until the event completes.
Note
This is a wrapper around cudaEventSynchronize(): see CUDA Event documentation for more info.


Makes all future work submitted to the given stream wait for this event.
Use torch.cuda.current_stream() if no stream is specified.


Returns list CUDA architectures this library was compiled for.


Gets the cuda capability of a device.
device (torch.device or int, optional) – device for which to return the device capability. This function is a no-op if this argument is a negative integer. It uses the current device, given by current_device(), if device is None (default).
the major and minor cuda capability of the device
tuple(int, int)


Gets the name of a device.
device (torch.device or int, optional) – device for which to return the name. This function is a no-op if this argument is a negative integer. It uses the current device, given by current_device(), if device is None (default).
the name of the device
str


Gets the properties of a device.
device (torch.device or int or str) – device for which to return the properties of the device.
the properties of the device
_CudaDeviceProperties


Returns NVCC gencode flags this library were compiled with.


Returns the random number generator state of the specified GPU as a ByteTensor.
device (torch.device or int, optional) – The device to return the RNG state of. Default: 'cuda' (i.e., torch.device('cuda'), the current CUDA device).
Warning
This function eagerly initializes CUDA.


Returns a list of ByteTensor representing the random number states of all devices.


Initialize PyTorch’s CUDA state. You may need to call this explicitly if you are interacting with PyTorch via its C API, as Python bindings for CUDA functionality will not be available until this initialization takes place. Ordinary users should not need this, as all of PyTorch’s CUDA methods automatically initialize CUDA state on-demand.
Does nothing if the CUDA state is already initialized.


Returns the current random seed of the current GPU.
Warning
This function eagerly initializes CUDA.


Force collects GPU memory after it has been released by CUDA IPC.
Note
Checks if any sent CUDA tensors could be cleaned from the memory. Force closes shared memory file used for reference counting if there is no active counters. Useful when the producer process stopped actively sending tensors and want to release unused memory.


Returns a bool indicating if CUDA is currently available.


Returns whether PyTorch’s CUDA state has been initialized.


Returns a human-readable printout of the running processes and their GPU memory use for a given device.
This can be useful to display periodically during training, or when handling out-of-memory exceptions.
device (torch.device or int, optional) – selected device. Returns printout for the current device, given by current_device(), if device is None (default).


Sets the seed for generating random numbers for the current GPU. It’s safe to call this function if CUDA is not available; in that case, it is silently ignored.
seed (int) – The desired seed.
Warning
If you are working with a multi-GPU model, this function is insufficient to get determinism. To seed all GPUs, use manual_seed_all().


Sets the seed for generating random numbers on all GPUs. It’s safe to call this function if CUDA is not available; in that case, it is silently ignored.
seed (int) – The desired seed.


Returns the maximum GPU memory occupied by tensors in bytes for a given device.
By default, this returns the peak allocated memory since the beginning of this program. reset_peak_stats() can be used to reset the starting point in tracking this metric. For example, these two functions can measure the peak allocated memory usage of each iteration in a training loop.
device (torch.device or int, optional) – selected device. Returns statistic for the current device, given by current_device(), if device is None (default).
Note
See Memory management for more details about GPU memory management.


Returns the maximum GPU memory managed by the caching allocator in bytes for a given device.
By default, this returns the peak cached memory since the beginning of this program. reset_peak_stats() can be used to reset the starting point in tracking this metric. For example, these two functions can measure the peak cached memory amount of each iteration in a training loop.
device (torch.device or int, optional) – selected device. Returns statistic for the current device, given by current_device(), if device is None (default).
Note
See Memory management for more details about GPU memory management.


Returns the current GPU memory occupied by tensors in bytes for a given device.
device (torch.device or int, optional) – selected device. Returns statistic for the current device, given by current_device(), if device is None (default).
Note
This is likely less than the amount shown in nvidia-smi since some unused memory can be held by the caching allocator and some context needs to be created on GPU. See Memory management for more details about GPU memory management.


Returns the current GPU memory managed by the caching allocator in bytes for a given device.
device (torch.device or int, optional) – selected device. Returns statistic for the current device, given by current_device(), if device is None (default).
Note
See Memory management for more details about GPU memory management.


Returns a snapshot of the CUDA memory allocator state across all devices.
Interpreting the output of this function requires familiarity with the memory allocator internals.
Note
See Memory management for more details about GPU memory management.


Returns a dictionary of CUDA memory allocator statistics for a given device.
The return value of this function is a dictionary of statistics, each of which is a non-negative integer.
Core statistics:
For these core statistics, values are broken down as follows.
Pool type:
Metric type:
In addition to the core statistics, we also provide some simple event counters:
device (torch.device or int, optional) – selected device. Returns statistics for the current device, given by current_device(), if device is None (default).
Note
See Memory management for more details about GPU memory management.


Returns a human-readable printout of the current memory allocator statistics for a given device.
This can be useful to display periodically during training, or when handling out-of-memory exceptions.
Note
See Memory management for more details about GPU memory management.


Describe an instantaneous event that occurred at some point.
msg (string) – ASCII message to associate with the event.


Pops a range off of a stack of nested range spans. Returns the zero-based depth of the range that is ended.


Pushes a range onto a stack of nested range span. Returns zero-based depth of the range that is started.
msg (string) – ASCII message to associate with range


Resets the starting point in tracking maximum GPU memory occupied by tensors for a given device.
See max_memory_allocated() for details.
device (torch.device or int, optional) – selected device. Returns statistic for the current device, given by current_device(), if device is None (default).
Warning
This function now calls reset_peak_memory_stats(), which resets /all/ peak memory stats.
Note
See Memory management for more details about GPU memory management.


Resets the starting point in tracking maximum GPU memory managed by the caching allocator for a given device.
See max_memory_cached() for details.
device (torch.device or int, optional) – selected device. Returns statistic for the current device, given by current_device(), if device is None (default).
Warning
This function now calls reset_peak_memory_stats(), which resets /all/ peak memory stats.
Note
See Memory management for more details about GPU memory management.


Sets the seed for generating random numbers to a random number for the current GPU. It’s safe to call this function if CUDA is not available; in that case, it is silently ignored.
Warning
If you are working with a multi-GPU model, this function will only initialize the seed on one GPU. To initialize all GPUs, use seed_all().


Sets the seed for generating random numbers to a random number on all GPUs. It’s safe to call this function if CUDA is not available; in that case, it is silently ignored.


Sets the current device.
Usage of this function is discouraged in favor of device. In most cases it’s better to use CUDA_VISIBLE_DEVICES environmental variable.
device (torch.device or int) – selected device. This function is a no-op if this argument is negative.


Set memory fraction for a process. The fraction is used to limit an caching allocator to allocated memory on a CUDA device. The allowed value equals the total visible memory multiplied fraction. If trying to allocate more than the allowed value in a process, will raise an out of memory error in allocator.
Note
In general, the total available free memory is less than the total capacity.


Sets the random number generator state of the specified GPU.


Sets the random number generator state of all devices.
new_states (Iterable of torch.ByteTensor) – The desired state for each device


Wrapper around a CUDA stream.
A CUDA stream is a linear sequence of execution that belongs to a specific device, independent from other streams. See CUDA semantics for details.
Note
Although CUDA versions >= 11 support more than two levels of priorities, in PyTorch, we only support two levels of priorities.
Checks if all the work submitted has been completed.
A boolean indicating if all kernels in this stream are completed.
Records an event.
event (Event, optional) – event to record. If not given, a new one will be allocated.
Recorded event.
Wait for all the kernels in this stream to complete.
Note
This is a wrapper around cudaStreamSynchronize(): see CUDA Stream documentation for more info.
Makes all future work submitted to the stream wait for an event.
event (Event) – an event to wait for.
Note
This is a wrapper around cudaStreamWaitEvent(): see CUDA Stream documentation for more info.
This function returns without waiting for event: only future operations are affected.
Synchronizes with another stream.
All future work submitted to this stream will wait until all kernels submitted to a given stream at the time of call complete.
stream (Stream) – a stream to synchronize.
Note
This function returns without waiting for currently enqueued kernels in stream: only future operations are affected.


Context-manager that selects a given stream.
All CUDA kernels queued within its context will be enqueued on a selected stream.
stream (Stream) – selected stream. This manager is a no-op if it’s None.
Note
Streams are per-device. If the selected stream is not on the current device, this function will also change the current device to match the stream.


Checks if all the work submitted has been completed.
A boolean indicating if all kernels in this stream are completed.


Records an event.
event (Event, optional) – event to record. If not given, a new one will be allocated.
Recorded event.


Wait for all the kernels in this stream to complete.
Note
This is a wrapper around cudaStreamSynchronize(): see CUDA Stream documentation for more info.


Makes all future work submitted to the stream wait for an event.
event (Event) – an event to wait for.
Note
This is a wrapper around cudaStreamWaitEvent(): see CUDA Stream documentation for more info.
This function returns without waiting for event: only future operations are affected.


Synchronizes with another stream.
All future work submitted to this stream will wait until all kernels submitted to a given stream at the time of call complete.
stream (Stream) – a stream to synchronize.
Note
This function returns without waiting for currently enqueued kernels in stream: only future operations are affected.


Waits for all kernels in all streams on a CUDA device to complete.
device (torch.device or int, optional) – device for which to synchronize. It uses the current device, given by current_device(), if device is None (default).


Returns a namedtuple (values, indices) where values is the cumulative maximum of elements of input in the dimension dim. And indices is the index location of each maximum value found in the dimension dim.
out (tuple, optional) – the result tuple of two output tensors (values, indices)
Example:


Returns a namedtuple (values, indices) where values is the cumulative minimum of elements of input in the dimension dim. And indices is the index location of each maximum value found in the dimension dim.
out (tuple, optional) – the result tuple of two output tensors (values, indices)
Example:


Returns the cumulative product of elements of input in the dimension dim.
For example, if input is a vector of size N, the result will also be a vector of size N, with elements.
Example:


Returns the cumulative sum of elements of input in the dimension dim.
For example, if input is a vector of size N, the result will also be a vector of size N, with elements.
Example:


Returns a new tensor with each of the elements of input converted from angles in degrees to radians.
input (Tensor) – the input tensor.
out (Tensor, optional) – the output tensor.
Example:


Returns an fp32 Tensor by dequantizing a quantized Tensor
tensor (Tensor) – A quantized Tensor
Given a list of quantized Tensors, dequantize them and return a list of fp32 Tensors
tensors (sequence of Tensors) – A list of quantized Tensors


Calculates determinant of a square matrix or batches of square matrices.
Note
torch.det() is deprecated. Please use torch.linalg.det() instead.
Note
Backward through detdet  internally uses SVD results when input is not invertible. In this case, double backward through detdet  will be unstable when input doesn’t have distinct singular values. See torch.svd~torch.svd  for details.
input (Tensor) – the input tensor of size (*, n, n) where * is zero or more batch dimensions.
Example:


The argument diagonal controls which diagonal to consider:
out (Tensor, optional) – the output tensor.
See also
torch.diagonal() always returns the diagonal of its input.
torch.diagflat() always constructs a tensor with diagonal elements specified by the input.
Examples:
Get the square matrix where the input vector is the diagonal:
Get the k-th diagonal of a given matrix:


The argument offset controls which diagonal to consider:
Examples:


Returns a partial view of input with the its diagonal elements with respect to dim1 and dim2 appended as a dimension at the end of the shape.
The argument offset controls which diagonal to consider:
Applying torch.diag_embed() to the output of this function with the same arguments yields a diagonal matrix with the diagonal entries of the input. However, torch.diag_embed() has different default dimensions, so those need to be explicitly specified.
Note
To take a batch diagonal, pass in dim1=-2, dim2=-1.
Examples:


Creates a tensor whose diagonals of certain 2D planes (specified by dim1 and dim2) are filled by input. To facilitate creating batched diagonal matrices, the 2D planes formed by the last two dimensions of the returned tensor are chosen by default.
The argument offset controls which diagonal to consider:
The size of the new matrix will be calculated to make the specified diagonal of the size of the last input dimension. Note that for offset other than 00 , the order of dim1 and dim2 matters. Exchanging them is equivalent to changing the sign of offset.
Applying torch.diagonal() to the output of this function with the same arguments yields a matrix identical to input. However, torch.diagonal() has different default dimensions, so those need to be explicitly specified.
Example:


Computes the n-th forward difference along the given dimension.
The first-order differences are given by out[i] = input[i + 1] - input[i]. Higher-order differences are calculated by using torch.diff() recursively.
Note
Only n = 1 is currently supported
out (Tensor, optional) – the output tensor.
Example:


Computes the logarithmic derivative of the gamma function on input.
input (Tensor) – the tensor to compute the digamma function on
out (Tensor, optional) – the output tensor.
Note
This function is similar to SciPy’s scipy.special.digamma.
Note
From PyTorch 1.8 onwards, the digamma function returns -Inf for 0. Previously it returned NaN for 0.
Example:


Returns the p-norm of (input - other)
The shapes of input and other must be broadcastable.
Example:


Note
Please refer to PyTorch Distributed Overview for a brief introduction to all features related to distributed training.
torch.distributed supports three built-in backends, each with different capabilities. The table below shows which functions are available for use with CPU / CUDA tensors. MPI supports CUDA only if the implementation used to build PyTorch supports it.
Backend
gloo
mpi
nccl
Device
CPU
GPU
CPU
GPU
CPU
GPU
send
✓
✘
✓
?
✘
✘
recv
✓
✘
✓
?
✘
✘
broadcast
✓
✓
✓
?
✘
✓
all_reduce
✓
✓
✓
?
✘
✓
reduce
✓
✘
✓
?
✘
✓
all_gather
✓
✘
✓
?
✘
✓
gather
✓
✘
✓
?
✘
✘
scatter
✓
✘
✓
?
✘
✘
reduce_scatter
✘
✘
✘
✘
✘
✓
all_to_all
✘
✘
✓
?
✘
✘
barrier
✓
✘
✓
?
✘
✓
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 backend but NCCL, If the init_method argument of init_process_group() points to a file it must adhere to the following schema:
Same as on Linux platform, you can enable TcpStore by setting environment variables, MASTER_ADDR and MASTER_PORT.
In the past, we were often asked: “which backend should I use?”.
Rule of thumb
GPU hosts with InfiniBand interconnect
GPU hosts with Ethernet interconnect
CPU hosts with InfiniBand interconnect
CPU hosts with Ethernet interconnect
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):
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.
NCCL has also provided a number of environment variables for fine-tuning purposes.
Commonly used ones include the following for debugging purposes:
For the full list of NCCL environment variables, please refer to NVIDIA NCCL’s official documentation
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() 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 and 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() wrapper may still have advantages over other approaches to data-parallelism, including torch.nn.DataParallel():
The package needs to be initialized using the torch.distributed.init_process_group() function before calling any other methods. This blocks until all processes have joined.
Returns 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.
Initializes the default distributed process group, and this will also initialize the distributed package.
If neither is specified, init_method is assumed to be “env://”.
To enable backend == Backend.MPI, PyTorch needs to be built from source on a system that supports MPI.
An enum-like class of available backends: GLOO, NCCL, MPI, 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.
Returns the backend of the given process group.
group (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.
The backend of the given process group as a lower case string.
Returns the rank of current process group
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.
group (ProcessGroup, optional) – The process group to work on. If None, the default process group will be used.
The rank of the process group -1, if not part of the group
Returns the number of processes in the current process group
group (ProcessGroup, optional) – The process group to work on. If None, the default process group will be used.
The world size of the process group -1, if not part of the group
Checking if the default process group has been initialized
Checks if the MPI backend is available.
Checks if the NCCL backend is available.
Currently three initialization methods are supported:
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.
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.
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:
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://).
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 pacakge 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.
Base class for all store implementations, such as the 3 provided by PyTorch distributed: (TCPStore, FileStore, and HashStore).
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.
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.
A store implementation that uses a file to store the underlying key-value pairs.
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.
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.
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.
key (str) – The function will return the value associated with this key.
Value associated with key if key is in the store.
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.
Overloaded function.
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.
keys (list) – List of keys on which to wait until they are set in the store.
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.
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.
The number of keys present in the store.
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.
key (str) – The key to be deleted from the store
True if key was deleted, otherwise False.
Sets the store’s default timeout. This timeout is used during initialization and in wait() and get().
timeout (timedelta) – timeout to be set in the store.
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).
Creates 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
Using multiple process groups with the NCCL backend concurrently is not safe and the user should perform explicit synchronization in their application to ensure only one process group is used at a time. This means collectives from one process group should have completed execution on the device (not just enqueued since CUDA execution is async) before collectives from another process group are enqueued. See Using multiple NCCL communicators concurrently for more details.
A handle of distributed group that can be given to collective calls.
Sends a tensor synchronously.
Receives a tensor synchronously.
Sender rank -1, if not part of the group
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:
Sends a tensor asynchronously.
A distributed request object. None, if not part of the group
Receives a tensor asynchronously.
A distributed request object. None, if not part of the group
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. 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:
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:
Broadcasts the tensor to the whole group.
tensor must have the same number of elements in all processes participating in the collective.
Async work handle, if async_op is set to True. None, if not async_op or if not part of the group
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.
None. If rank is part of the group, object_list will contain the broadcasted objects from src rank.
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 responsiblity 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 all_gather() collective since it does not provide an async_op handle and thus will be a blocking call.
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.
Reduces the tensor data across all machines in such 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.
Async work handle, if async_op is set to True. None, if not async_op or if not part of the group
Reduces the tensor data across all machines.
Only the process with rank dst is going to receive the final result.
Async work handle, if async_op is set to True. None, if not async_op or if not part of the group
Gathers tensors from the whole group in a list.
Complex tensors are supported.
Async work handle, if async_op is set to True. None, if not async_op or if not part of the group
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.
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 responsiblity to ensure that this is set so that each rank has an individual GPU, via torch.cuda.set_device().
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.
Gathers a list of tensors in a single process.
Async work handle, if async_op is set to True. None, if not async_op or if not part of the group
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.
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
Note that this API is not supported when using the NCCL backend.
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.
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.
Async work handle, if async_op is set to True. None, if not async_op or if not part of the group
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.
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
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.
Reduces, then scatters a list of tensors to all processes in a group.
Async work handle, if async_op is set to True. None, if not async_op or if not part of the group.
Each process scatters list of input tensors to all processes in a group and return gathered list of tensors in output list.
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.
Synchronizes 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().
Async work handle, if async_op is set to True. None, if not async_op or if not part of the group
An enum-like class for available reduction operations: SUM, PRODUCT, MIN, MAX, BAND, BOR, and BXOR.
Note that BAND, BOR, and BXOR reductions are not available when using the NCCL backend.
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(), all_reduce_multigpu(), etc.
Members:
SUM
PRODUCT
MIN
MAX
BAND
BOR
BXOR
Deprecated enum-like class for reduction operations: SUM, PRODUCT, MIN, and MAX.
ReduceOp is recommended to use instead.
If you want to use collective communication functions supporting autograd you can find an implementation of those in the torch.distributed.nn.* module.
Functions here are synchronous and will be inserted in the autograd graph, so you need to ensure that all the processes that participated in the collective operation will do the backward pass for the backward communication to effectively happen and don’t cause a deadlock.
Please notice that currently the only backend where all the functions are guaranteed to work is gloo. .. autofunction:: torch.distributed.nn.broadcast .. autofunction:: torch.distributed.nn.gather .. autofunction:: torch.distributed.nn.scatter .. autofunction:: torch.distributed.nn.reduce .. autofunction:: torch.distributed.nn.all_gather .. autofunction:: torch.distributed.nn.all_to_all .. autofunction:: torch.distributed.nn.all_reduce
If you have more than one GPU on each node, when using the NCCL and Gloo backend, broadcast_multigpu() all_reduce_multigpu() reduce_multigpu() all_gather_multigpu() and reduce_scatter_multigpu() support distributed collective operations among multiple GPUs within each node. These functions can potentially improve the overall distributed training performance and be easily used by passing a list of tensors. Each Tensor in the passed tensor list needs to be on a separate GPU device of the host where the function is called. Note that the length of the tensor list needs to be identical among all the distributed processes. Also note that currently the multi-GPU collective functions are only supported by the NCCL backend.
For example, if the system we use for distributed training has 2 nodes, each of which has 8 GPUs. On each of the 16 GPUs, there is a tensor that we would like to all-reduce. The following code can serve as a reference:
Code running on Node 0
Code running on Node 1
After the call, all 16 tensors on the two nodes will have the all-reduced value of 16
Broadcasts the tensor to the whole group with multiple GPU tensors per node.
tensor must have the same number of elements in all the GPUs from all processes participating in the collective. each tensor in the list must be on a different GPU
Only nccl and gloo backend are currently supported tensors should only be GPU tensors
Async work handle, if async_op is set to True. None, if not async_op or if not part of the group
Reduces the tensor data across all machines in such a way that all get the final result. This function reduces a number of tensors on every node, while each tensor resides on different GPUs. Therefore, the input tensor in the tensor list needs to be GPU tensors. Also, each tensor in the tensor list needs to reside on a different GPU.
After the call, all tensor in tensor_list is going to be bitwise identical in all processes.
Complex tensors are supported.
Only nccl and gloo backend is currently supported tensors should only be GPU tensors
Async work handle, if async_op is set to True. None, if not async_op or if not part of the group
Reduces the tensor data on multiple GPUs across all machines. Each tensor in tensor_list should reside on a separate GPU
Only the GPU of tensor_list[dst_tensor] on the process with rank dst is going to receive the final result.
Only nccl backend is currently supported tensors should only be GPU tensors
Async work handle, if async_op is set to True. None, otherwise
Gathers tensors from the whole group in a list. Each tensor in tensor_list should reside on a separate GPU
Only nccl backend is currently supported tensors should only be GPU tensors
Complex tensors are supported.
output_tensor_lists (List[List[Tensor]]) –
Output lists. It should contain correctly-sized tensors on each GPU to be used for output of the collective, e.g. output_tensor_lists[i] contains the all_gather result that resides on the GPU of input_tensor_list[i].
Note that each element of output_tensor_lists has the size of world_size * len(input_tensor_list), since the function all gathers the result from every single GPU in the group. To interpret each element of output_tensor_lists[i], note that input_tensor_list[j] of rank k will be appear in output_tensor_lists[i][k * world_size + j]
Also note that len(output_tensor_lists), and the size of each element in output_tensor_lists (each element is a list, therefore len(output_tensor_lists[i])) need to be the same for all the distributed processes calling this function.
Async work handle, if async_op is set to True. None, if not async_op or if not part of the group
Reduce and scatter a list of tensors to the whole group. Only nccl backend is currently supported.
Each tensor in output_tensor_list should reside on a separate GPU, as should each list of tensors in input_tensor_lists.
output_tensor_list (List[Tensor]) –
Output tensors (on different GPUs) to receive the result of the operation.
Note that len(output_tensor_list) needs to be the same for all the distributed processes calling this function.
input_tensor_lists (List[List[Tensor]]) –
Input lists. It should contain correctly-sized tensors on each GPU to be used for input of the collective, e.g. input_tensor_lists[i] contains the reduce_scatter input that resides on the GPU of output_tensor_list[i].
Note that each element of input_tensor_lists has the size of world_size * len(output_tensor_list), since the function scatters the result from every single GPU in the group. To interpret each element of input_tensor_lists[i], note that output_tensor_list[j] of rank k receives the reduce-scattered result from input_tensor_lists[i][k * world_size + j]
Also note that len(input_tensor_lists), and the size of each element in input_tensor_lists (each element is a list, therefore len(input_tensor_lists[i])) need to be the same for all the distributed processes calling this function.
Async work handle, if async_op is set to True. None, if not async_op or if not part of the group.
Besides the 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 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.
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.
torch.distributed.launch is a module that spawns up multiple distributed training processes on each of the training nodes.
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 benefitial 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:
Node 1: (IP: 192.168.1.1, and has a free port: 1234)
Node 2:
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
Set your device to local rank using either
or
3. In your training program, you are supposed to call the following function at the beginning to start the distributed backend. You need to make sure that the init_method uses env://, which is the only supported init_method by this module.
4. In your training program, you can either use regular distributed functions or use torch.nn.parallel.DistributedDataParallel() module. If your training program uses GPUs for training and you would like to use torch.nn.parallel.DistributedDataParallel() module, here is how to configure it.
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 https://github.com/pytorch/pytorch/issues/12042 for an example of how things can go wrong if you don’t do this correctly.
The Multiprocessing package - torch.multiprocessing package also provides a spawn function in 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
Note that this function requires Python 3.4 or higher.


This DDP communication hook just calls allreduce using GradBucket tensors. Once gradient tensors are aggregated across all workers, its then callback takes the mean and returns the result. If user registers this hook, DDP results is expected to be same as the case where no hook was registered. Hence, this won’t change behavior of DDP and user can use this as a reference or modify this hook to log useful information or any other purposes while unaffecting DDP behavior.


This DDP communication hook implements a simple gradient compression approach that converts GradBucket tensors whose type is assumed to be torch.float32 to half-precision floating point format (torch.float16). It allreduces those float16 gradient tensors. Once compressed gradient tensors are allreduced, its then callback called decompress converts the aggregated result back to float32 and takes the mean.


This DDP communication hook implements a simplified PowerSGD gradient compression algorithm described in the paper. This variant does not compress the gradients layer by layer, but instead compresses the flattened input tensor that batches all the gradients. Therefore, it is faster than powerSGD_hook(), but usually results in a much lower accuracy, unless matrix_approximation_rank is 1.
Warning
Increasing matrix_approximation_rank here may not necessarily increase the accuracy, because batching per-parameter tensors without column/row alignment can destroy low-rank structure. Therefore, the user should always consider powerSGD_hook() first, and only consider this variant when a satisfactory accuracy can be achieved when matrix_approximation_rank is 1.
Once gradient tensors are aggregated across all workers, this hook applies compression as follows:
Note that this communication hook enforces vanilla allreduce for the first state.start_powerSGD_iter iterations. This not only gives the user more control over the tradeoff between speedup and accuracy, but also helps abstract away some complexity of the internal optimization of DDP for future communication hook developers.
Future handler of the communication, which updates the gradients in place.


Stores both the algorithm’s hyperparameters and the internal state for all the gradients during the training. Particularly, matrix_approximation_rank and start_powerSGD_iter are the main hyperparameters that should be tuned by the user. For performance, we suggest to keep binary hyperparameters use_error_feedback and warm_start on.
matrix_approximation_rank controls the size of compressed low-rank tensors, which determines the compression rate. The lower the rank, the stronger the compression.
1.1. If matrix_approximation_rank is too low, the full model quality will need more training steps to reach or will never reach and yield loss in accuracy.
1.2. The increase of matrix_approximation_rank can substantially increase the computation costs of the compression, and the accuracy may not be futher improved beyond a certain matrix_approximation_rank threshold.
To tune matrix_approximation_rank, we suggest to start from 1 and increase by factors of 2 (like an expoential grid search, 1, 2, 4, …), until a satisfactory accuracy is reached. Typically only a small value 1-4 is used. For some NLP tasks (as shown in Appendix D of the original paper), this value has been increased to 32.
To tune start_powerSGD_iter, we suggest to start with 10% of total training steps, and increase it until a satisfactory accuracy is reached.
Warning
If error feedback or warm-up is enabled, the minimum value of start_powerSGD_iter allowed in DDP is 2. This is because there is another internal optimization that rebuilds buckets at iteration 1 in DDP, and this can conflict with any tensor memorized before the rebuild process.


This DDP communication hook implements PowerSGD gradient compression algorithm described in the paper. Once gradient tensors are aggregated across all workers, this hook applies compression as follows:
Handles rank-1 tensors by allreducing them without compression:
2.1. Allocate contiguous memory for those rank-1 tensors, and allreduces all the rank-1 tensors as a batch, without compression;
2.2. Copies the individual rank-1 tensors from the contiguous memory back to the input tensor.
Handles high-rank tensors by PowerSGD compression:
3.1. For each high-rank tensor M, creates two low-rank tensors P and Q for decomposing M, such that M = PQ^T, where Q is initialized from a standard normal distribution and orthogonalized;
3.2. Computes each P in Ps, which is equal to MQ;
3.3. Allreduces Ps as a batch;
3.4. Orthogonalizes each P in Ps;
3.5. Computes each Q in Qs, which is approximately equal to M^TP;
3.6. Allreduces Qs as a batch;
3.7. Computes each M among all the high-rank tensors, which is approximately equal to PQ^T.
Note that this communication hook enforces vanilla allreduce for the first state.start_powerSGD_iter iterations. This not only gives the user more control over the tradeoff between speedup and accuracy, but also helps abstract away some complexity of the internal optimization of DDP for future communication hook developers.
Future handler of the communication, which updates the gradients in place.


Gathers tensors from the whole group in a list.
Complex tensors are supported.
Async work handle, if async_op is set to True. None, if not async_op or if not part of the group


Gathers tensors from the whole group in a list. Each tensor in tensor_list should reside on a separate GPU
Only nccl backend is currently supported tensors should only be GPU tensors
Complex tensors are supported.
output_tensor_lists (List[List[Tensor]]) –
Output lists. It should contain correctly-sized tensors on each GPU to be used for output of the collective, e.g. output_tensor_lists[i] contains the all_gather result that resides on the GPU of input_tensor_list[i].
Note that each element of output_tensor_lists has the size of world_size * len(input_tensor_list), since the function all gathers the result from every single GPU in the group. To interpret each element of output_tensor_lists[i], note that input_tensor_list[j] of rank k will be appear in output_tensor_lists[i][k * world_size + j]
Also note that len(output_tensor_lists), and the size of each element in output_tensor_lists (each element is a list, therefore len(output_tensor_lists[i])) need to be the same for all the distributed processes calling this function.
Async work handle, if async_op is set to True. None, if not async_op or if not part of the group


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.
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 responsiblity to ensure that this is set so that each rank has an individual GPU, via torch.cuda.set_device().
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.


Reduces the tensor data across all machines in such 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.
Async work handle, if async_op is set to True. None, if not async_op or if not part of the group


Reduces the tensor data across all machines in such a way that all get the final result. This function reduces a number of tensors on every node, while each tensor resides on different GPUs. Therefore, the input tensor in the tensor list needs to be GPU tensors. Also, each tensor in the tensor list needs to reside on a different GPU.
After the call, all tensor in tensor_list is going to be bitwise identical in all processes.
Complex tensors are supported.
Only nccl and gloo backend is currently supported tensors should only be GPU tensors
Async work handle, if async_op is set to True. None, if not async_op or if not part of the group


Each process scatters list of input tensors to all processes in a group and return gathered list of tensors in output list.
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.


Kicks off the distributed backward pass using the provided roots. This currently implements the FAST mode algorithm which assumes all RPC messages sent in the same distributed autograd context across workers would be part of the autograd graph during the backward pass.
We use the provided roots to discover the autograd graph and compute appropriate dependencies. This method blocks until the entire autograd computation is done.
We accumulate the gradients in the appropriate torch.distributed.autograd.context on each of the nodes. The autograd context to be used is looked up given the context_id that is passed in when torch.distributed.autograd.backward() is called. If there is no valid autograd context corresponding to the given ID, we throw an error. You can retrieve the accumulated gradients using the get_gradients() API.


Context object to wrap forward and backward passes when using distributed autograd. The context_id generated in the with statement is required to uniquely identify a distributed backward pass on all workers. Each worker stores metadata associated with this context_id, which is required to correctly execute a distributed autograd pass.


Retrieves a map from Tensor to the appropriate gradient for that Tensor accumulated in the provided context corresponding to the given context_id as part of the distributed autograd backward pass.
context_id (int) – The autograd context id for which we should retrieve the gradients.
A map where the key is the Tensor and the value is the associated gradient for that Tensor.


An enum-like class of available backends: GLOO, NCCL, MPI, 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.


Synchronizes 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().
Async work handle, if async_op is set to True. None, if not async_op or if not part of the group


Broadcasts the tensor to the whole group.
tensor must have the same number of elements in all processes participating in the collective.
Async work handle, if async_op is set to True. None, if not async_op or if not part of the group


Broadcasts the tensor to the whole group with multiple GPU tensors per node.
tensor must have the same number of elements in all the GPUs from all processes participating in the collective. each tensor in the list must be on a different GPU
Only nccl and gloo backend are currently supported tensors should only be GPU tensors
Async work handle, if async_op is set to True. None, if not async_op or if not part of the group


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.
None. If rank is part of the group, object_list will contain the broadcasted objects from src rank.
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 responsiblity 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 all_gather() collective since it does not provide an async_op handle and thus will be a blocking call.
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.


A store implementation that uses a file to store the underlying key-value pairs.


Gathers a list of tensors in a single process.
Async work handle, if async_op is set to True. None, if not async_op or if not part of the group


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.
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
Note that this API is not supported when using the NCCL backend.
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.


Returns the backend of the given process group.
group (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.
The backend of the given process group as a lower case string.


Returns the rank of current process group
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.
group (ProcessGroup, optional) – The process group to work on. If None, the default process group will be used.
The rank of the process group -1, if not part of the group


Returns the number of processes in the current process group
group (ProcessGroup, optional) – The process group to work on. If None, the default process group will be used.
The world size of the process group -1, if not part of the group


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.


Initializes the default distributed process group, and this will also initialize the distributed package.
If neither is specified, init_method is assumed to be “env://”.
To enable backend == Backend.MPI, PyTorch needs to be built from source on a system that supports MPI.


Receives a tensor asynchronously.
A distributed request object. None, if not part of the group


Sends a tensor asynchronously.
A distributed request object. None, if not part of the group


Returns 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.


Checking if the default process group has been initialized


Checks if the NCCL backend is available.


Creates 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
Using multiple process groups with the NCCL backend concurrently is not safe and the user should perform explicit synchronization in their application to ensure only one process group is used at a time. This means collectives from one process group should have completed execution on the device (not just enqueued since CUDA execution is async) before collectives from another process group are enqueued. See Using multiple NCCL communicators concurrently for more details.
A handle of distributed group that can be given to collective calls.


DistributedOptimizer takes remote references to parameters scattered across workers and applies the given optimizer locally for each parameter.
This class uses get_gradients() in order to retrieve the gradients for specific parameters.
Concurrent calls to step(), either from the same or different clients, will be serialized on each worker – as each worker’s optimizer can only work on one set of gradients at a time. However, there is no guarantee that the full forward-backward-optimizer sequence will execute for one client at a time. This means that the gradients being applied may not correspond to the latest forward pass executed on a given worker. Also, there is no guaranteed ordering across workers.
DistributedOptimizer creates the local optimizer with TorchScript enabled by default, so that optimizer updates are not blocked by the Python Global Interpreter Lock (GIL) during multithreaded training (e.g. Distributed Model Parallel). This feature is currently in beta stage, enabled for optimizers including Adagrad, Adam, SGD, RMSprop, AdamW and Adadelta. We are increasing the coverage to all optimizers in future releases.
Performs a single optimization step.
This will call torch.optim.Optimizer.step() on each worker containing parameters to be optimized, and will block until all workers return. The provided context_id will be used to retrieve the corresponding context that contains the gradients that should be applied to the parameters.
context_id – the autograd context id for which we should run the optimizer step.


Performs a single optimization step.
This will call torch.optim.Optimizer.step() on each worker containing parameters to be optimized, and will block until all workers return. The provided context_id will be used to retrieve the corresponding context that contains the gradients that should be applied to the parameters.
context_id – the autograd context id for which we should run the optimizer step.


Wraps an arbitrary nn.Sequential module to train on using synchronous pipeline parallelism. If the module requires lots of memory and doesn’t fit on a single GPU, pipeline parallelism is a useful technique to employ for training.
The implementation is based on the torchgpipe paper.
Pipe combines pipeline parallelism with checkpointing to reduce peak memory required to train while minimizing device under-utilization.
You should place all the modules on the appropriate devices and wrap them into an nn.Sequential module defining the desired order of execution.
Pipeline of two FC layers across GPUs 0 and 1.
Note
You can wrap a Pipe model with torch.nn.parallel.DistributedDataParallel only when the checkpoint parameter of Pipe is 'never'.
Note
Pipe only supports intra-node pipelining currently, but will be expanded to support inter-node pipelining in the future. The forward function returns an RRef to allow for inter-node pipelining in the future, where the output might be on a remote host. For intra-node pipelinining you can use local_value() to retrieve the output locally.
Warning
Pipe is experimental and subject to change.
Processes a single input mini-batch through the pipe and returns an RRef pointing to the output. Pipe is a fairly transparent module wrapper. It doesn’t modify the input and output signature of the underlying module. But there’s type restriction. Input and output have to be a Tensor or a sequence of tensors. This restriction is applied at partition boundaries too.
The input tensor is split into multiple micro-batches based on the chunks parameter used to initialize Pipe. The batch size is assumed to be the first dimension of the tensor and if the batch size is less than chunks, the number of micro-batches is equal to the batch size.
input (torch.Tensor or sequence of Tensor) – input mini-batch
RRef to the output of the mini-batch
TypeError – input is not a tensor or sequence of tensors.


Processes a single input mini-batch through the pipe and returns an RRef pointing to the output. Pipe is a fairly transparent module wrapper. It doesn’t modify the input and output signature of the underlying module. But there’s type restriction. Input and output have to be a Tensor or a sequence of tensors. This restriction is applied at partition boundaries too.
The input tensor is split into multiple micro-batches based on the chunks parameter used to initialize Pipe. The batch size is assumed to be the first dimension of the tensor and if the batch size is less than chunks, the number of micro-batches is equal to the batch size.
input (torch.Tensor or sequence of Tensor) – input mini-batch
RRef to the output of the mini-batch
TypeError – input is not a tensor or sequence of tensors.


The command to pop a skip tensor.
name (str) – name of skip tensor
the skip tensor previously stashed by another layer under the same name


The decorator to define a nn.Module with skip connections. Decorated modules are called “skippable”. This functionality works perfectly fine even when the module is not wrapped by Pipe.
Each skip tensor is managed by its name. Before manipulating skip tensors, a skippable module must statically declare the names for skip tensors by stash and/or pop parameters. Skip tensors with pre-declared name can be stashed by yield stash(name, tensor) or popped by tensor = yield
pop(name).
Here is an example with three layers. A skip tensor named “1to3” is stashed and popped at the first and last layer, respectively:
One skippable module can stash or pop multiple skip tensors:
Every skip tensor must be associated with exactly one pair of stash and pop. Pipe checks this restriction automatically when wrapping a module. You can also check the restriction by verify_skippables() without Pipe.


Verifies if the underlying skippable modules satisfy integrity.
Every skip tensor must have only one pair of stash and pop. If there are one or more unmatched pairs, it will raise TypeError with the detailed messages.
Here are a few failure cases. verify_skippables() will report failure for these cases:
To use the same name for multiple skip tensors, they must be isolated by different namespaces. See isolate().
TypeError – one or more pairs of stash and pop are not matched.


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.


Receives a tensor synchronously.
Sender rank -1, if not part of the group


Reduces the tensor data across all machines.
Only the process with rank dst is going to receive the final result.
Async work handle, if async_op is set to True. None, if not async_op or if not part of the group


An enum-like class for available reduction operations: SUM, PRODUCT, MIN, MAX, BAND, BOR, and BXOR.
Note that BAND, BOR, and BXOR reductions are not available when using the NCCL backend.
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(), all_reduce_multigpu(), etc.
Members:
SUM
PRODUCT
MIN
MAX
BAND
BOR
BXOR


Reduces the tensor data on multiple GPUs across all machines. Each tensor in tensor_list should reside on a separate GPU
Only the GPU of tensor_list[dst_tensor] on the process with rank dst is going to receive the final result.
Only nccl backend is currently supported tensors should only be GPU tensors
Async work handle, if async_op is set to True. None, otherwise


Deprecated enum-like class for reduction operations: SUM, PRODUCT, MIN, and MAX.
ReduceOp is recommended to use instead.


Reduces, then scatters a list of tensors to all processes in a group.
Async work handle, if async_op is set to True. None, if not async_op or if not part of the group.


Reduce and scatter a list of tensors to the whole group. Only nccl backend is currently supported.
Each tensor in output_tensor_list should reside on a separate GPU, as should each list of tensors in input_tensor_lists.
output_tensor_list (List[Tensor]) –
Output tensors (on different GPUs) to receive the result of the operation.
Note that len(output_tensor_list) needs to be the same for all the distributed processes calling this function.
input_tensor_lists (List[List[Tensor]]) –
Input lists. It should contain correctly-sized tensors on each GPU to be used for input of the collective, e.g. input_tensor_lists[i] contains the reduce_scatter input that resides on the GPU of output_tensor_list[i].
Note that each element of input_tensor_lists has the size of world_size * len(output_tensor_list), since the function scatters the result from every single GPU in the group. To interpret each element of input_tensor_lists[i], note that output_tensor_list[j] of rank k receives the reduce-scattered result from input_tensor_lists[i][k * world_size + j]
Also note that len(input_tensor_lists), and the size of each element in input_tensor_lists (each element is a list, therefore len(input_tensor_lists[i])) need to be the same for all the distributed processes calling this function.
Async work handle, if async_op is set to True. None, if not async_op or if not part of the group.


An enum class of available backends.
PyTorch ships with two builtin backends: BackendType.TENSORPIPE and BackendType.PROCESS_GROUP. Additional ones can be registered using the register_backend() function.


A decorator for a function indicating that the return value of the function is guaranteed to be a Future object and this function can run asynchronously on the RPC callee. More specifically, the callee extracts the Future returned by the wrapped function and installs subsequent processing steps as a callback to that Future. The installed callback will read the value from the Future when completed and send the value back as the RPC response. That also means the returned Future only exists on the callee side and is never sent through RPC. This decorator is useful when the wrapped function’s (fn) execution needs to pause and resume due to, e.g., containing rpc_async() or waiting for other signals.
Note
To enable asynchronous execution, applications must pass the function object returned by this decorator to RPC APIs. If RPC detected attributes installed by this decorator, it knows that this function returns a Future object and will handle that accordingly. However, this does not mean this decorator has to be outmost one when defining a function. For example, when combined with @staticmethod or @classmethod, @rpc.functions.async_execution needs to be the inner decorator to allow the target function be recognized as a static or class function. This target function can still execute asynchronously because, when accessed, the static or class method preserves attributes installed by @rpc.functions.async_execution.
The returned Future object can come from rpc_async(), then(), or Future constructor. The example below shows directly using the Future returned by then().
When combined with TorchScript decorators, this decorator must be the outmost one.
When combined with static or class method, this decorator must be the inner one.
This decorator also works with RRef helpers, i.e., . torch.distributed.rpc.RRef.rpc_sync(), torch.distributed.rpc.RRef.rpc_async(), and torch.distributed.rpc.RRef.remote().


Get WorkerInfo of a given worker name. Use this WorkerInfo to avoid passing an expensive string on every invocation.
worker_name (str) – the string name of a worker. If None, return the the id of the current worker. (default None)
WorkerInfo instance for the given worker_name or WorkerInfo of the current worker if worker_name is None.


Initializes RPC primitives such as the local RPC agent and distributed autograd, which immediately makes the current process ready to send and receive RPCs.


The backend options class for ProcessGroupAgent, which is derived from RpcBackendOptions.
URL specifying how to initialize the process group. Default is env://
The number of threads in the thread-pool used by ProcessGroupAgent.
A float indicating the timeout to use for all RPCs. If an RPC does not complete in this timeframe, it will complete with an exception indicating that it has timed out.


URL specifying how to initialize the process group. Default is env://


The number of threads in the thread-pool used by ProcessGroupAgent.


A float indicating the timeout to use for all RPCs. If an RPC does not complete in this timeframe, it will complete with an exception indicating that it has timed out.


Make a remote call to run func on worker to and return an RRef to the result value immediately. Worker to will be the owner of the returned RRef, and the worker calling remote is a user. The owner manages the global reference count of its RRef, and the owner RRef is only destructed when globally there are no living references to it.
A user RRef instance to the result value. Use the blocking API torch.distributed.rpc.RRef.to_here() to retrieve the result value locally.
Warning
Using GPU tensors as arguments or return values of func is not supported since we don’t support sending GPU tensors over the wire. You need to explicitly copy GPU tensors to CPU before using them as arguments or return values of func.
Warning
The remote API does not copy storages of argument tensors until sending them over the wire, which could be done by a different thread depending on the RPC backend type. The caller should make sure that the contents of those tensors stay intact until the returned RRef is confirmed by the owner, which can be checked using the torch.distributed.rpc.RRef.confirmed_by_owner() API.
Warning
Errors such as timeouts for the remote API are handled on a best-effort basis. This means that when remote calls initiated by remote fail, such as with a timeout error, we take a best-effort approach to error handling. This means that errors are handled and set on the resulting RRef on an asynchronous basis. If the RRef has not been used by the application before this handling (such as to_here or fork call), then future uses of the RRef will appropriately raise errors. However, it is possible that the user application will use the RRef before the errors are handled. In this case, errors may not be raised as they have not yet been handled.
Make sure that MASTER_ADDR and MASTER_PORT are set properly on both workers. Refer to init_process_group() API for more details. For example,
Then run the following code in two different processes:
Below is an example of running a TorchScript function using RPC.


An abstract structure encapsulating the options passed into the RPC backend. An instance of this class can be passed in to init_rpc() in order to initialize RPC with specific configurations, such as the RPC timeout and init_method to be used.
URL specifying how to initialize the process group. Default is env://
A float indicating the timeout to use for all RPCs. If an RPC does not complete in this timeframe, it will complete with an exception indicating that it has timed out.


Make a non-blocking RPC call to run function func on worker to. RPC messages are sent and received in parallel to execution of Python code. This method is thread-safe. This method will immediately return a Future that can be awaited on.
Returns a Future object that can be waited on. When completed, the return value of func on args and kwargs can be retrieved from the Future object.
Warning
Using GPU tensors as arguments or return values of func is not supported since we don’t support sending GPU tensors over the wire. You need to explicitly copy GPU tensors to CPU before using them as arguments or return values of func.
Warning
The rpc_async API does not copy storages of argument tensors until sending them over the wire, which could be done by a different thread depending on the RPC backend type. The caller should make sure that the contents of those tensors stay intact until the returned Future completes.
Make sure that MASTER_ADDR and MASTER_PORT are set properly on both workers. Refer to init_process_group() API for more details. For example,
Then run the following code in two different processes:
Below is an example of running a TorchScript function using RPC.


Make a blocking RPC call to run function func on worker to. RPC messages are sent and received in parallel to execution of Python code. This method is thread-safe.
Returns the result of running func with args and kwargs.
Warning
Using GPU tensors as arguments or return values of func is not supported since we don’t support sending GPU tensors over the wire. You need to explicitly copy GPU tensors to CPU before using them as arguments or return values of func.
Make sure that MASTER_ADDR and MASTER_PORT are set properly on both workers. Refer to init_process_group() API for more details. For example,
Then run the following code in two different processes:
Below is an example of running a TorchScript function using RPC.


Runs the backward pass using the RRef as the root of the backward pass. If dist_autograd_ctx_id is provided, we perform a distributed backward pass using the provided ctx_id starting from the owner of the RRef. In this case, get_gradients() should be used to retrieve the gradients. If dist_autograd_ctx_id is None, it is assumed that this is a local autograd graph and we only perform a local backward pass. In the local case, the node calling this API has to be the owner of the RRef. The value of the RRef is expected to be a scalar Tensor.
Returns whether this RRef has been confirmed by the owner. OwnerRRef always returns true, while UserRRef only returns true when the owner knowns about this UserRRef.
Returns whether or not the current node is the owner of this RRef.
If the current node is the owner, returns a reference to the local value. Otherwise, throws an exception.
Returns worker information of the node that owns this RRef.
Returns worker name of the node that owns this RRef.
Create a helper proxy to easily launch a remote using the owner of the RRef as the destination to run functions on the object referenced by this RRef. More specifically, rref.remote().func_name(*args, **kwargs) is the same as the following:
timeout (float, optional) – Timeout for rref.remote(). If the creation of this RRef is not successfully completed within the timeout, then the next time there is an attempt to use the RRef (such as to_here), a timeout will be raised. If not provided, the default RPC timeout will be used. Please see rpc.remote() for specific timeout semantics for RRef.
Create a helper proxy to easily launch an rpc_async using the owner of the RRef as the destination to run functions on the object referenced by this RRef. More specifically, rref.rpc_async().func_name(*args, **kwargs) is the same as the following:
timeout (float, optional) – Timeout for rref.rpc_async(). If the call does not complete within this timeframe, an exception indicating so will be raised. If this argument is not provided, the default RPC timeout will be used.
Create a helper proxy to easily launch an rpc_sync using the owner of the RRef as the destination to run functions on the object referenced by this RRef. More specifically, rref.rpc_sync().func_name(*args, **kwargs) is the same as the following:
timeout (float, optional) – Timeout for rref.rpc_sync(). If the call does not complete within this timeframe, an exception indicating so will be raised. If this argument is not provided, the default RPC timeout will be used.
Blocking call that copies the value of the RRef from the owner to the local node and returns it. If the current node is the owner, returns a reference to the local value.
timeout (float, optional) – Timeout for to_here. If the call does not complete within this timeframe, an exception indicating so will be raised. If this argument is not provided, the default RPC timeout (60s) will be used.


Runs the backward pass using the RRef as the root of the backward pass. If dist_autograd_ctx_id is provided, we perform a distributed backward pass using the provided ctx_id starting from the owner of the RRef. In this case, get_gradients() should be used to retrieve the gradients. If dist_autograd_ctx_id is None, it is assumed that this is a local autograd graph and we only perform a local backward pass. In the local case, the node calling this API has to be the owner of the RRef. The value of the RRef is expected to be a scalar Tensor.


Returns whether this RRef has been confirmed by the owner. OwnerRRef always returns true, while UserRRef only returns true when the owner knowns about this UserRRef.


Returns whether or not the current node is the owner of this RRef.


If the current node is the owner, returns a reference to the local value. Otherwise, throws an exception.


Returns worker information of the node that owns this RRef.


Returns worker name of the node that owns this RRef.


Create a helper proxy to easily launch a remote using the owner of the RRef as the destination to run functions on the object referenced by this RRef. More specifically, rref.remote().func_name(*args, **kwargs) is the same as the following:
timeout (float, optional) – Timeout for rref.remote(). If the creation of this RRef is not successfully completed within the timeout, then the next time there is an attempt to use the RRef (such as to_here), a timeout will be raised. If not provided, the default RPC timeout will be used. Please see rpc.remote() for specific timeout semantics for RRef.


Create a helper proxy to easily launch an rpc_async using the owner of the RRef as the destination to run functions on the object referenced by this RRef. More specifically, rref.rpc_async().func_name(*args, **kwargs) is the same as the following:
timeout (float, optional) – Timeout for rref.rpc_async(). If the call does not complete within this timeframe, an exception indicating so will be raised. If this argument is not provided, the default RPC timeout will be used.


Create a helper proxy to easily launch an rpc_sync using the owner of the RRef as the destination to run functions on the object referenced by this RRef. More specifically, rref.rpc_sync().func_name(*args, **kwargs) is the same as the following:
timeout (float, optional) – Timeout for rref.rpc_sync(). If the call does not complete within this timeframe, an exception indicating so will be raised. If this argument is not provided, the default RPC timeout will be used.


Blocking call that copies the value of the RRef from the owner to the local node and returns it. If the current node is the owner, returns a reference to the local value.
timeout (float, optional) – Timeout for to_here. If the call does not complete within this timeframe, an exception indicating so will be raised. If this argument is not provided, the default RPC timeout (60s) will be used.


Perform a shutdown of the RPC agent, and then destroy the RPC agent. This stops the local agent from accepting outstanding requests, and shuts down the RPC framework by terminating all RPC threads. If graceful=True, this will block until all local and remote RPC processes reach this method and wait for all outstanding work to complete. Otherwise, if graceful=False, this is a local shutdown, and it does not wait for other RPC processes to reach this method.
Warning
For Future objects returned by rpc_async(), future.wait() should not be called after shutdown().
graceful (bool) – Whether to do a graceful shutdown or not. If True, this will 1) wait until there is no pending system messages for UserRRefs and delete them; 2) block until all local and remote RPC processes have reached this method and wait for all outstanding work to complete.
Make sure that MASTER_ADDR and MASTER_PORT are set properly on both workers. Refer to init_process_group() API for more details. For example,
Then run the following code in two different processes:


The backend options for TensorPipeAgent, derived from RpcBackendOptions.
The device map locations.
URL specifying how to initialize the process group. Default is env://
The number of threads in the thread-pool used by TensorPipeAgent to execute requests.
A float indicating the timeout to use for all RPCs. If an RPC does not complete in this timeframe, it will complete with an exception indicating that it has timed out.
Set device mapping between each RPC caller and callee pair. This function can be called multiple times to incrementally add device placement configurations.


URL specifying how to initialize the process group. Default is env://


The number of threads in the thread-pool used by TensorPipeAgent to execute requests.


A float indicating the timeout to use for all RPCs. If an RPC does not complete in this timeframe, it will complete with an exception indicating that it has timed out.


Set device mapping between each RPC caller and callee pair. This function can be called multiple times to incrementally add device placement configurations.


A structure that encapsulates information of a worker in the system. Contains the name and ID of the worker. This class is not meant to be constructed directly, rather, an instance can be retrieved through get_worker_info() and the result can be passed in to functions such as rpc_sync(), rpc_async(), remote() to avoid copying a string on every invocation.
Globally unique id to identify the worker.
The name of the worker.


Globally unique id to identify the worker.


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.
Async work handle, if async_op is set to True. None, if not async_op or if not part of the group


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.
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
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.


Base class for all store implementations, such as the 3 provided by PyTorch distributed: (TCPStore, FileStore, and HashStore).


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.


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.
key (str) – The key to be deleted from the store
True if key was deleted, otherwise False.


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.
key (str) – The function will return the value associated with this key.
Value associated with key if key is in the store.


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.
The number of keys present in the store.


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.


Sets the store’s default timeout. This timeout is used during initialization and in wait() and get().
timeout (timedelta) – timeout to be set in the store.


Overloaded function.
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.
keys (list) – List of keys on which to wait until they are set in the store.
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.


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.