Backends

tensorplay.distributed supports four built-in backends, each with different capabilities. The table below shows which functions are available for use with a CPU or GPU for each backend. For NCCL, GPU refers to CUDA GPU while for XCCL to XPU GPU. MPI supports CUDA only if the implementation used to build TensorPlay supports it.

Basics

The tensorplay.distributed package provides TensorPlay support and communication primitives for multiprocess parallelism across several computation nodes running on one or more machines. The class tensorplay.nn.parallel.DistributedDataParallel() builds on this functionality to provide synchronous distributed training as a wrapper around any TensorPlay model. This differs from the kinds of parallelism provided by Multiprocessing package - tensorplay.multiprocessing and tensorplay.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, tensorplay.distributed or the tensorplay.nn.parallel.DistributedDataParallel() wrapper may still have advantages over other approaches to data-parallelism, including tensorplay.nn.DataParallel():

• Each process maintains its own optimizer and performs a complete optimization step with each iteration. While this may appear redundant, since the gradients have already been gathered together and averaged across processes and are thus the same for every process, this means that no parameter broadcast step is needed, reducing time spent transferring tensors between nodes.

• Each process contains an independent Python interpreter, eliminating the extra interpreter overhead and “GIL-thrashing” that comes from driving several execution threads, model replicas, or GPUs from a single Python process. This is especially important for models that make heavy use of the Python runtime, including models with recurrent layers or many small components.

Initialization

The package needs to be initialized using the tensorplay.distributed.init_process_group() or tensorplay.distributed.device_mesh.init_device_mesh() function before calling any other methods. Both block until all processes have joined.

Currently three initialization methods are supported:

import tensorplay.distributed as dist

# Use address of one of the machines
dist.init_process_group(backend, init_method='tcp://10.1.1.20:23456',
                        rank=args.rank, world_size=4)

import tensorplay.distributed as dist

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

Post-Initialization

Once tensorplay.distributed.init_process_group() was run, the following functions can be used. To check whether the process group has already been initialized use tensorplay.distributed.is_initialized().

Shutdown

It is important to clean up resources on exit by calling destroy_process_group(). The simplest pattern to follow is to destroy every process group and backend by calling destroy_process_group() with the default value of None for the group argument, at a point in the training script where communications are no longer needed, usually near the end of main(). The call should be made once per trainer-process, not at the outer process-launcher level. if destroy_process_group() is not called by all ranks in a pg within the timeout duration, especially when there are multiple process-groups in the application e.g. for N-D parallelism, hangs on exit are possible. This is because the destructor for ProcessGroupNCCL calls ncclCommAbort, which must be called collectively, but the order of calling ProcessGroupNCCL’s destructor if called by python’s GC is not deterministic. Calling destroy_process_group() helps by ensuring ncclCommAbort is called in a consistent order across ranks, and avoids calling ncclCommAbort during ProcessGroupNCCL’s destructor.

Groups

By default collectives operate on the default group (also called the world) and require all processes to enter the distributed function call. However, some workloads can benefit from more fine-grained communication. This is where distributed groups come into play. new_group() function can be used to create new groups, with arbitrary subsets of all processes. It returns an opaque group handle that can be given as a group argument to all collectives (collectives are distributed functions to exchange information in certain well-known programming patterns).

tensorplay.distributed.distributed_core.GroupMember = <class 'tensorplay.distributed.distributed_core.GroupMember'>[source]

DeviceMesh

DeviceMesh is a higher level abstraction that manages process groups (or NCCL communicators). It allows user to easily create inter node and intra node process groups without worrying about how to set up the ranks correctly for different sub process groups, and it helps manage those distributed process group easily. init_device_mesh() function can be used to create new DeviceMesh, with a mesh shape describing the device topology.

Point-to-point communication

isend() and irecv() return distributed request objects when used. In general, the type of this object is unspecified as they should never be created manually, but they are guaranteed to support two methods:

• is_completed() - returns True if the operation has finished

• wait() - will block the process until the operation is finished. is_completed() is guaranteed to return True once it returns.

Synchronous and asynchronous collective operations

Every collective operation function supports the following two kinds of operations, depending on the setting of the async_op flag passed into the collective: Synchronous operation - the default mode, when async_op is set to False. When the function returns, it is guaranteed that the collective operation is performed. In the case of CUDA operations, it is not guaranteed that the CUDA operation is completed, since CUDA operations are asynchronous. For CPU collectives, any further function calls utilizing the output of the collective call will behave as expected. For CUDA collectives, function calls utilizing the output on the same CUDA stream will behave as expected. Users must take care of synchronization under the scenario of running under different streams. For details on CUDA semantics such as stream synchronization, see CUDA Semantics. See the below script to see examples of differences in these semantics for CPU and CUDA operations. Asynchronous operation - when async_op is set to True. The collective operation function returns a distributed request object. In general, you don’t need to create it manually and it is guaranteed to support two methods:

• is_completed() - in the case of CPU collectives, returns True if completed. In the case of CUDA operations, returns True if the operation has been successfully enqueued onto a CUDA stream and the output can be utilized on the default stream without further synchronization.

• wait() - in the case of CPU collectives, will block the process until the operation is completed. In the case of CUDA collectives, will block the currently active CUDA stream until the operation is completed (but will not block the CPU).

• get_future() - returns tensorplay._C.Future object. Supported for NCCL, also supported for most operations on GLOO and MPI, except for peer to peer operations. Note: as we continue adopting Futures and merging APIs, get_future() call might become redundant. Example The following code can serve as a reference regarding semantics for CUDA operations when using distributed collectives. It shows the explicit need to synchronize when using collective outputs on different CUDA streams:

# Code runs on each rank.
dist.init_process_group("nccl", rank=rank, world_size=2)
output = tensorplay.tensor([rank]).cuda(rank)
s = tensorplay.cuda.Stream()
handle = dist.all_reduce(output, async_op=True)
# Wait ensures the operation is enqueued, but not necessarily complete.
handle.wait()
# Using result on non-default stream.
with tensorplay.cuda.stream(s):
    s.wait_stream(tensorplay.cuda.default_stream())
    output.add_(100)
if rank == 0:
    # if the explicit call to wait_stream was omitted, the output below will be
    # non-deterministically 1 or 101, depending on whether the allreduce overwrote
    # the value after the add completed.
    print(output)

Collective functions

class reduce_op

Deprecated enum-like class for reduction operations: SUM, PRODUCT, MIN, and MAX.

ReduceOp is recommended to use instead.

Distributed Key-Value Store

The distributed package comes with a distributed key-value store, which can be used to share information between processes in the group as well as to initialize the distributed package in tensorplay.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.

Profiling Collective Communication

Note that you can use tensorplay.profiler (recommended, only available after 1.8.1) or tensorplay.autograd.profiler to profile collective communication and point-to-point communication APIs mentioned here. All out-of-the-box backends (gloo, nccl, mpi) are supported and collective communication usage will be rendered as expected in profiling output/traces. Profiling your code is the same as any regular tensorplay operator:

import tensorplay
import tensorplay.distributed as dist
with tensorplay.profiler():
    tensor = tensorplay.randn(20, 10)
    dist.all_reduce(tensor)

Please refer to the profiler documentation for a full overview of profiler features.

Multi-GPU collective functions

Object collectives

Object collectives are a set of collective-like operations that work on arbitrary Python objects, as long as they can be pickled. There are various collective patterns implemented (e.g. broadcast, all_gather, …) but they each roughly follow this pattern:

1. convert the input object into a pickle (raw bytes), then shove it into a byte tensor

2. communicate the size of this byte tensor to peers (first collective operation)

3. allocate appropriately sized tensor to perform the real collective

4. communicate the object data (second collective operation)

5. convert raw data back into Python (unpickle) Object collectives sometimes have surprising performance or memory characteristics that lead to long runtimes or OOMs, and thus they should be used with caution. Here are some common issues. Asymmetric pickle/unpickle time - Pickling objects can be slow, depending on the number, type and size of the objects. When the collective has a fan-in (e.g. gather_object), the receiving rank(s) must unpickle N times more objects than the sending rank(s) had to pickle, which can cause other ranks to time out on their next collective. Inefficient tensor communication - Tensors should be sent via regular collective APIs, not object collective APIs. It is possible to send Tensors via object collective APIs, but they will be serialized and deserialized (including a CPU-sync and device-to-host copy in the case of non-CPU tensors), and in almost every case other than debugging or troubleshooting code, it would be worth the trouble to refactor the code to use non-object collectives instead. Unexpected tensor devices - If you still want to send tensors via object collectives, there is another aspect specific to cuda (and possibly other accelerators) tensors. If you pickle a tensor that is currently on cuda:3, and then unpickle it, you will get another tensor on cuda:3 regardless of which process you are on, or which CUDA device is the ‘default’ device for that process. With regular tensor collective APIs, ‘output tensors’ will always be on the same, local device, which is generally what you’d expect. Unpickling a tensor will implicitly activate a CUDA context if it is the first time a GPU is used by the process, which can waste significant amounts of GPU memory. This issue can be avoided by moving tensors to CPU before passing them as inputs to an object collective.

Third-party backends

Besides the builtin GLOO/MPI/NCCL backends, TensorPlay 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 the samples under test/cpp_extension/. The capability of third-party backends are decided by their own implementations. The new backend derives from tensorplay.distributed.ProcessGroup and registers the backend name and the instantiating interface through tensorplay.distributed.Backend.register_backend() when imported. When manually importing this backend and invoking tensorplay.distributed.init_process_group() with the corresponding backend name, the tensorplay.distributed package runs on the new backend.

TorchComms backend

TorchComms is an optional communication backend for tensorplay.distributed. When enabled, it overrides the normal backend instantiation in init_process_group() so that all process groups are created through TorchComms instead of the built-in ProcessGroup implementations.

export TP_DISTRIBUTED_USE_TORCHCOMMS=1

import tensorplay.distributed.config as dist_config

dist_config.use_torchcomms = True

dist.init_process_group(backend="nccl", device_id=tensorplay.device("cuda", local_rank))

Launch utility

The tensorplay.distributed package also provides a launch utility in tensorplay.distributed.launch. This helper utility can be used to launch multiple processes per node for distributed training.

Spawn utility

The Multiprocessing package - tensorplay.multiprocessing package also provides a spawn function in tensorplay.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 TensorPlay example - ImageNet implementation Note that this function requires Python 3.4 or higher.

Debugging tensorplay.distributed applications

Debugging distributed applications can be challenging due to hard to understand hangs, crashes, or inconsistent behavior across ranks. tensorplay.distributed provides a suite of tools to help debug training applications in a self-serve fashion:

import os
from datetime import timedelta

import tensorplay
import tensorplay.distributed as dist
import tensorplay.multiprocessing as mp


def worker(rank):
    dist.init_process_group("nccl", rank=rank, world_size=2)
    # monitored barrier requires gloo process group to perform host-side sync.
    group_gloo = dist.new_group(backend="gloo")
    if rank not in [1]:
        dist.monitored_barrier(group=group_gloo, timeout=timedelta(seconds=2))


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

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

┌───────────────────────────────────┐
│     Frontend Server (rank 0)      │
│  HTTP server on a fixed port      │
│  Renders HTML dashboards          │
│  Fans out requests to all ranks   │
└──────────┬────────────────────────┘
           │  HTTP POST /handler/<endpoint>?<args>
           ▼
┌──────────────────┐  ┌──────────────────┐  ┌──────────────────┐
│ WorkerServer (0) │  │ WorkerServer (1) │  │ WorkerServer (N) │
│ Per-rank C++     │  │ Per-rank C++     │  │ Per-rank C++     │
│ HTTP server      │  │ HTTP server      │  │ HTTP server      │
│ Exposes handlers │  │ Exposes handlers │  │ Exposes handlers │
└──────────────────┘  └──────────────────┘  └──────────────────┘

pip install jinja2 aiohttp tabulate

import tensorplay
import tensorplay.distributed as dist
from tensorplay.distributed.debug import start_debug_server, stop_debug_server

dist.init_process_group("nccl")

# Start the debug server on all ranks.
# The frontend (browser UI) is served on rank 0 at port 25999.
start_debug_server(port=25999)

# ... your training loop ...

stop_debug_server()
dist.destroy_process_group()

# train.py
import tensorplay
import tensorplay.distributed as dist
from tensorplay.distributed.debug import start_debug_server

def main():
    dist.init_process_group("nccl")
    start_debug_server()

    rank = dist.get_rank()
    device = tensorplay.device(f"cuda:{rank}")
    model = tensorplay.nn.Linear(10, 10).to(device)
    model = tensorplay.nn.parallel.DistributedDataParallel(model, device_ids=[rank])

    for _ in range(1000):
        x = tensorplay.randn(32, 10, device=device)
        loss = model(x).sum()
        loss.backward()

    dist.destroy_process_group()

if __name__ == "__main__":
    main()

python -m tensorplay.distributed.run --nproc-per-node=2 train.py
# Open http://localhost:25999

POST http://<worker-host>:<worker-port>/handler/<handler_name>?<params>

curl -X POST \
  http://worker-host:port/handler/ping  # @lint-ignore

curl -X POST \
  http://worker-host:port/handler/dump_traceback  # @lint-ignore

curl -X POST \
  "http://worker-host:port/handler/pyspy_dump?nonblocking=1&native=1"  # @lint-ignore

curl -X POST \
  http://worker-host:port/handler/fr_trace_json  # @lint-ignore

curl -X POST \
  "http://worker-host:port/handler/torchcomms_fr_trace_json?onlyactive=true"  # @lint-ignore

curl -X POST \
  --output trace.json \
  "http://worker-host:port/handler/debug_profile?duration=5"  # @lint-ignore
# Open trace.json in chrome://tracing or https://ui.perfetto.dev

curl -X POST \
  http://worker-host:port/handler/wait_counter_values  # @lint-ignore

start_debug_server(
    dump_dir="/shared/nfs/debug_dumps",
    dump_interval=120.0,  # dump every 2 minutes
    enabled_dumps={"stacks", "fr_trace", "pyspy_dump", "wait_counters", "tcpstore"},
)

class DebugHandler(ABC):
    fetch_timeout: float = _DEFAULT_FETCH_TIMEOUT

    @abstractmethod
    def routes(self) -> list[Route]: ...

    @abstractmethod
    def nav_links(self) -> list[NavLink]: ...

    def templates(self) -> dict[str, str]:
        return {}

    def dump(self) -> str | None:
        return None

    def dump_filename(self) -> str:
        return type(self).__name__.lower()

from tensorplay.distributed.debug._debug_handlers import default_handlers

handlers = default_handlers() + [MyCustomHandler()]
start_debug_server(handlers=handlers)

Logging

The underlying C++ library of tensorplay.distributed outputs log messages at various levels. These messages can be helpful to understand the execution state of a distributed training job and to troubleshoot problems such as network connection failures.

Watchdog (Experimental)

The tensorplay.distributed._watchdog module provides a pure-Python watchdog for detecting hung operations in Python-based distributed backends and related distributed primitives. It monitors both CPU-side hangs (e.g., a stuck rendezvous) and GPU-side hangs (e.g., a kernel that never completes) by running an asyncio event loop on a background daemon thread. By default, timeouts dump all thread stack traces and abort the process.