torchcomms

https://meta-pytorch.org/torchcomms/
Admin
torchcomms is an experimental, lightweight communication API for PyTorch Distributed (PTD) that...
Tokens:24,356
Snippets:189
Trust Score:6.4
Update:2 months ago
Show doc for...
Tokens:
Raw
### Install torchcomms from source

Source: https://meta-pytorch.org/torchcomms/main/getting_started.html

Install required dependencies and the torchcomms package.

```bash
# Install PyTorch (if not already installed)
pip install -r requirements.txt
pip install -v --no-build-isolation .
```

--------------------------------

### Build RCCL backend

Source: https://meta-pytorch.org/torchcomms/main/getting_started.html

Install prerequisites and set environment variables before building the RCCL backend.

```bash
conda install conda-forge::glog=0.4.0 conda-forge::gflags conda-forge::fmt -y
```

```bash
export ROCM_HOME=/opt/rocm
export RCCL_INCLUDE=$ROCM_HOME/include/rccl
```

```bash
./build_rccl.sh
```

--------------------------------

### Configure build backends

Source: https://meta-pytorch.org/torchcomms/main/getting_started.html

Set environment variables to enable or disable specific backends before installation.

```bash
# Enable/disable specific backends (ON/OFF or 1/0)
export USE_NCCL=ON    # Default: ON
export USE_NCCLX=ON   # Default: ON
export USE_GLOO=ON    # Default: ON
export USE_RCCL=OFF   # Default: OFF
```

--------------------------------

### Execute Distributed Training Script

Source: https://meta-pytorch.org/torchcomms/main/getting_started.html

Commands to launch the distributed example across multiple processes using torchrun or the legacy torch.distributed.launch module.

```bash
# Using torchrun (recommended)
torchrun --nproc_per_node=2 example.py

# Or using python -m torch.distributed.launch
python -m torch.distributed.launch --nproc_per_node=2 example.py
```

--------------------------------

### Broadcast Object List Example

Source: https://meta-pytorch.org/torchcomms/main/_modules/torchcomms/objcol.html

Demonstrates how to broadcast a list of picklable objects from a root rank to all other ranks in the communication group.

```python
>>> # xdoctest: +SKIP("need comm init")
>>> # Note: comm initialization omitted on each rank.
>>> from torchcomms import objcol
>>> if comm.get_rank() == 0:
>>>     # Assumes world_size of 3.
>>>     objects = ["foo", 12, {1: 2}] # any picklable object
>>> else:
>>>     objects = [None, None, None]
>>> # Assumes backend is not NCCL
>>> objcol.broadcast_object_list(comm, objects, src=0, device=device)
>>> objects
['foo', 12, {1: 2}]
```

--------------------------------

### Install torchcomms via pip

Source: https://meta-pytorch.org/torchcomms/main/getting_started.html

Install torchcomms and PyTorch nightly builds for specific CUDA versions.

```bash
# Cuda 12.6
pip install --pre torch torchcomms --index-url https://download.pytorch.org/whl/nightly/cu126

# Cuda 13.0
pip install --pre torch torchcomms --index-url https://download.pytorch.org/whl/nightly/cu130

# Cuda 13.2
pip install --pre torch torchcomms --index-url https://download.pytorch.org/whl/nightly/cu132
```

--------------------------------

### Send and Receive Object Lists with torchcomms

Source: https://meta-pytorch.org/torchcomms/main/api.html

Demonstrates the usage of send_object_list and recv_object_list to exchange picklable objects between ranks. Requires a pre-initialized communicator and assumes a non-NCCL backend for this specific example.

```python
>>> # xdoctest: +SKIP("need comm init")
>>> # Note: comm initialization omitted on each rank.
>>> from torchcomms import objcol
>>> # Assumes backend is not NCCL
>>> if comm.get_rank() == 0:
>>>     # Assumes world_size of 2.
>>>     objects = ["foo", 12, {1: 2}] # any picklable object
>>>     objcol.send_object_list(comm, objects, dst=1)
>>> else:
>>>     objects = [None, None, None]
>>>     objcol.recv_object_list(comm, objects, src=0)
>>> objects
['foo', 12, {1: 2}]
```

```python
>>> # xdoctest: +SKIP("need comm init")
>>> # Note: comm initialization omitted on each rank.
>>> from torchcomms import objcol
>>> # Assumes backend is not NCCL
>>> if comm.get_rank() == 0:
>>>     # Assumes world_size of 2.
>>>     objects = ["foo", 12, {1: 2}] # any picklable object
>>>     objcol.send_object_list(comm, objects, dst=1)
>>> else:
>>>     objects = [None, None, None]
>>>     objcol.recv_object_list(comm, objects, src=0)
>>> objects
['foo', 12, {1: 2}]
```

--------------------------------

### Clone and prepare environment for source build

Source: https://meta-pytorch.org/torchcomms/main/getting_started.html

Set up a conda environment and clone the repository to prepare for building from source.

```bash
# Create a conda environment
conda create -n torchcomms python=3.10
conda activate torchcomms
# Clone the repository
git clone git@github.com:meta-pytorch/torchcomms.git
cd torchcomms
```

--------------------------------

### Register and initialize a backend

Source: https://meta-pytorch.org/torchcomms/main/api.html

Register the custom backend class and create a new communicator instance.

```python
torchcomms.register_backend("my_backend" , MyBackend)

comm = torchcomms.new_comm("my_backend", device, "my_comm")
```

--------------------------------

### Manage RMA Windows with TorchComm

Source: https://meta-pytorch.org/torchcomms/main/api.html

Demonstrates creating windows, registering tensors, and performing one-sided communication operations. Requires the ncclx backend.

```python
import torchcomms

allocator = torchcomms.get_mem_allocator(comm.get_backend())
pool = torch.cuda.MemPool(allocator)
with torch.cuda.use_mem_pool(pool):
    buffer = torch.ones([size], dtype=dtype, device=device)

# Option 1: Create window with tensor registration in one step
window = comm.new_window(buffer)

# Option 2: Create window and register buffer separately
window = comm.new_window()
window.tensor_register(buffer)

# Sender
window.put(data, dst_rank=1, target_offset_nelems=0, async_op=False)
window.signal(peer_rank=1, async_op=False)

# Receiver
window.wait_signal(peer_rank=0, async_op=False)
received = buffer[:size]

# Cleanup
window.tensor_deregister()
```

--------------------------------

### Initialize Torchcomms Module

Source: https://meta-pytorch.org/torchcomms/main/_modules/torchcomms.html

Handles library loading, metaclass registration for PyTorch compatibility, and backend discovery.

```python
# Copyright (c) Meta Platforms, Inc. and affiliates.

# pyre-strict
# patternlint-disable fbcode-nonempty-init-py
import ctypes
import os
import sys
from datetime import timedelta
from importlib.metadata import entry_points

# We need to load this upfront since libtorchcomms depend on libtorch
import torch  # noqa: F401
from torchcomms.functional import (
    is_torch_compile_supported,
    is_torch_compile_supported_and_enabled,
)

if is_torch_compile_supported():
    try:
        # PyTorch 2.14+ renamed torch._opaque_base to torch._custom_class_base
        # (pytorch/pytorch#188456)
        from torch._custom_class_base import OpaqueBaseMeta
    except ImportError:
        from torch._opaque_base import OpaqueBaseMeta

    # make the metaclass available to the pybind module
    sys.modules["torchcomms._opaque_meta"] = type(
        "module", (), {"OpaqueBaseMeta": OpaqueBaseMeta}
    )()

    # to support opaque registration for time delta.
    class Timeout(timedelta, metaclass=OpaqueBaseMeta):
        pass
else:
    # When compile support is disabled, define Timeout without the metaclass


[docs]
    class Timeout(timedelta):
        pass




def _load_libtorchcomms() -> None:
    libtorchcomms_path = os.path.join(os.path.dirname(__file__), "libtorchcomms.so")
    # OSS build, buck native linking links everything together so this is not needed
    if os.path.exists(libtorchcomms_path):
        # load this using RTLD_LOCAL so that we don't pollute the global namespace
        # We need to load this upfront since _comms and _comms_* depend on it
        # and won't be able to find it themselves.
        ctypes.CDLL(libtorchcomms_path, mode=ctypes.RTLD_LOCAL)


_load_libtorchcomms()
from torchcomms._comms import *  # noqa: E402, F401, F403
import torchcomms.hooks as hooks  # noqa: E402, F401
import torchcomms.objcol as objcol  # noqa: E402, F401, F403

if is_torch_compile_supported_and_enabled():
    # Import collectives first to ensure all operations are registered
    # This must happen before patch_torchcomm() so that window operations
    # and other collectives are registered and can be patched
    from torchcomms.functional import collectives  # noqa: F401


# The documentation uses __all__ to determine what is documented and in what
# order.
__all__ = [  # noqa: F405
    "new_comm",
    "TorchComm",
    "ReduceOp",
    "TorchWork",
    "Timeout",
    "BatchP2POptions",
    "BatchSendRecv",
    "P2POp",
    "CommOptions",
    "TorchCommWindow",
    "register_backend",
    "TorchCommBackend",
]

for name in __all__:
    cls = globals()[name]
    cls.__module__ = "torchcomms"


def _load_backend(backend: str) -> None:
    """Used to load backends lazily from C++

    C++ calls this only when the backend is not already registered via
    register_backend.
    """
    found = entry_points(group="torchcomms.backends", name=backend)
    if not found:
        raise ModuleNotFoundError(
            f"failed to find backend {backend}, is it registered via entry_points.txt?"
        )
    wheel = next(iter(found))
    wheel.load()


def is_backend_built(backend: str) -> bool:
    """True if the wheel was built with this backend's extension."""
    return bool(entry_points(group="torchcomms.backends", name=backend))


def built_backends() -> list[str]:
    """Names of all backends the wheel was built with."""
    return [ep.name for ep in entry_points(group="torchcomms.backends")]
```

--------------------------------

### Reconfigure Communicator Peers

Source: https://meta-pytorch.org/torchcomms/main/api.html

Shows how to initialize a communicator in a dynamic regime by exchanging handles and calling reconfigure.

```python
>>> comm = torchcomms.new_comm("mccl", device, "my_comm", init_dynamic_regime=True)
>>> my_handle = comm.get_init_handle()
>>> # Exchange handles with all ranks via store/coordinator
>>> all_handles = collect_handles_from_all_ranks(my_handle)
>>> work = comm.reconfigure(
...     uuid=0,
...     init_handles=all_handles,
...     timeout=timedelta(milliseconds=5000)
... )
>>> work.wait_blocking()
>>> # Now collective operations are permitted
>>> comm.all_reduce(tensor, op=torchcomms.ReduceOp.SUM, async_op=False)
```

--------------------------------

### Initialize and execute all_gather_p

Source: https://meta-pytorch.org/torchcomms/main/api.html

Use this pattern for persistent all-gather operations requiring the rcclx backend. The handle must be explicitly freed after use.

```python
# Initialize once
handle = comm.all_gather_p_init(output_tensor)

# Execute many times
for input_tensor in inputs:
    work = comm.all_gather_p_exec(handle, input_tensor, async_op=True)
    work.wait()

# Free when done
comm.all_gather_p_free(handle)
```

--------------------------------

### Initialize DeviceMesh with torchcomm

Source: https://meta-pytorch.org/torchcomms/main/_modules/torchcomms/device_mesh.html

Logic for creating a DeviceMesh by mapping global ranks to torchcomm process groups.

```python
global_pg = None
if _global_comm is not None:
    global_pg = _create_torchcomm_process_group(
        comm=_global_comm,
        group_name=_global_comm.get_name(),
        global_ranks_mapping=None,  # Will use default mapping
    )
elif len(mesh_dim_comms) != 1:
    raise RuntimeError(
        "More than one torch comm objects are passed but no global comm(_global_comm) is provided. "
        "Please provide a global comm object via _global_comm."
    )

group_names = []
idx = 0
for comm, name in zip(mesh_dim_comms, mesh_dim_names):
    group_name = name

    # Calculate global ranks mapping for this mesh dimension
    global_ranks = mesh.transpose(idx, -1).reshape(-1, mesh.size(idx))
    # Find the row containing the global rank
    row_idx = int(torch.where(global_ranks == global_rank)[0].item())
    list_rank = global_ranks[row_idx].tolist()
    global_ranks_mapping = {x: j for j, x in enumerate(list_rank)}

    # Use helper function to create the process group
    pg = _create_torchcomm_process_group(
        comm=comm,
        group_name=group_name,
        backend_str=backend_str,
        global_ranks_mapping=global_ranks_mapping,
    )
    if _global_comm is None and idx == 0:
        global_pg = pg

    group_names.append(group_name)
    idx += 1

# Set as the default world process group
dist.distributed_c10d.GroupMember.WORLD = global_pg

device_mesh = dist.DeviceMesh(
    device_type=device.type,
    mesh=mesh,
    mesh_dim_names=mesh_dim_names,
    _init_backend=False,
    _rank=global_rank,
)
device_mesh._dim_group_names = group_names

return device_mesh
```

--------------------------------

### Build NCCLX backend

Source: https://meta-pytorch.org/torchcomms/main/getting_started.html

Commands to build the NCCLX backend using either system libraries or source dependencies.

```bash
USE_SYSTEM_LIBS=1 ./build_ncclx.sh
```

```bash
./build_ncclx.sh
```

--------------------------------

### Perform Asynchronous AllReduce

Source: https://meta-pytorch.org/torchcomms/main/getting_started.html

Demonstrates initiating an asynchronous AllReduce operation and waiting for its completion while performing other tasks.

```python
import torch
from torchcomms import new_comm, ReduceOp

device = torch.device("cuda")
torchcomm = new_comm("ncclx", device, name="main_comm")

rank = torchcomm.get_rank()
device_id = rank % torch.cuda.device_count()
target_device = torch.device(f"cuda:{device_id}")

# Create tensor
tensor = torch.full((1024,), float(rank + 1), dtype=torch.float32, device=target_device)

# Start async AllReduce
work = torchcomm.all_reduce(tensor, ReduceOp.SUM, async_op=True)

# Do other work while communication happens
print(f"Rank {rank}: Doing other work while AllReduce is in progress...")

# Wait for completion
work.wait()
print(f"Rank {rank}: AllReduce completed")

torchcomm.finalize()
```

--------------------------------

### Implement Serialization Helper

Source: https://meta-pytorch.org/torchcomms/main/_modules/torchcomms/objcol.html

The _Serialization class handles object conversion using either pickle or torch.save based on the TORCHCOMMS_SERIALIZATION environment variable.

```python
class _Serialization:
    """Serialization helper with serialize and deserialize methods."""

    def __init__(self) -> None:
        self.use_pickle: bool = os.getenv("TORCHCOMMS_SERIALIZATION") == "pickle"

    def serialize(self, f: io.BytesIO, obj: object) -> None:
        if self.use_pickle:
            pickle.Pickler(f).dump(obj)
        else:
            torch.save(obj, f)

    def deserialize(self, f: io.BytesIO, weights_only: bool) -> object:
        if self.use_pickle:
            return pickle.Unpickler(f).load()
        else:
            return torch.load(f, weights_only=weights_only)
```

--------------------------------

### TorchComm.all_gather_p_init

Source: https://meta-pytorch.org/torchcomms/main/api.html

Initializes a persistent AllGather operation.

```APIDOC
## TorchComm.all_gather_p_init(output, hints=None, timeout=None)

### Description
Initialize a persistent AllGather operation. This is a SM free collective operation where memory is pre-registered.

### Parameters
- **output** (torch.Tensor) - Required - Pre-allocated output tensor of size (world_size * input_size).
- **hints** (Mapping[str, str]) - Optional - Dictionary of string hints for backend-specific options.
- **timeout** (timedelta) - Optional - Timeout for the operation.
```

--------------------------------

### Initialize DeviceMesh with TorchComm

Source: https://meta-pytorch.org/torchcomms/main/_modules/torchcomms/device_mesh.html

Initializes a PyTorch DeviceMesh using a collection of TorchComm instances for mesh dimensions.

```python
def init_device_mesh(
    mesh_dim_comms: tuple[TorchComm, ...],  # noqa: F405
    mesh_dim_names: tuple[str, ...],
    _global_comm: Optional[TorchComm] = None,  # noqa: F405
) -> dist.DeviceMesh:
    """
    Initializes a `DeviceMesh` from the list of provided `TorchComm` instances.

    See `DeviceMesh` for more details.
    """

    device = mesh_dim_comms[0].get_device()
    mesh_shape = tuple(comm.get_size() for comm in mesh_dim_comms)
    world_size = math.prod(mesh_shape)

    mesh = torch.arange(world_size, dtype=torch.int, device="cpu").view(mesh_shape)

    local_ranks = [comm.get_rank() for comm in mesh_dim_comms]
    global_rank = cast(int, mesh[tuple(local_ranks)].item())
    backend_str = "torchcomm"
    # Register the backend
    dist.Backend.register_backend(backend_str, new_comm)
```

--------------------------------

### Perform Synchronous AllReduce with TorchComms

Source: https://meta-pytorch.org/torchcomms/main/getting_started.html

Initializes a communicator with the NCCLX backend and performs a sum reduction across all ranks. Requires CUDA-enabled devices and proper environment configuration.

```python
#!/usr/bin/env python3
# example.py
import torch
from torchcomms import new_comm, ReduceOp

def main():
    # Initialize TorchComm with NCCLX backend
    device = torch.device("cuda")
    torchcomm = new_comm("ncclx", device, name="main_comm")

    # Get rank and world size
    rank = torchcomm.get_rank()
    world_size = torchcomm.get_size()

    # Calculate device ID
    num_devices = torch.cuda.device_count()
    device_id = rank % num_devices
    target_device = torch.device(f"cuda:{device_id}")

    print(f"Rank {rank}/{world_size}: Running on device {device_id}")

    # Create a tensor with rank-specific data
    tensor = torch.full(
        (1024,),
        float(rank + 1),
        dtype=torch.float32,
        device=target_device
    )

    print(f"Rank {rank}: Before AllReduce: {tensor[0].item()}")

    # Perform synchronous AllReduce (sum across all ranks)
    torchcomm.all_reduce(tensor, ReduceOp.SUM, async_op=False)

    # Synchronize CUDA stream
    torch.cuda.current_stream().synchronize()

    print(f"Rank {rank}: After AllReduce: {tensor[0].item()}")

    # Cleanup
    torchcomm.finalize()

if __name__ == "__main__":
    main()
```

--------------------------------

### init_device_mesh

Source: https://meta-pytorch.org/torchcomms/main/_modules/torchcomms/device_mesh.html

Initializes a PyTorch DeviceMesh from a list of provided TorchComm instances.

```APIDOC
## init_device_mesh(mesh_dim_comms, mesh_dim_names, _global_comm=None)

### Description
Initializes a `DeviceMesh` from the list of provided `TorchComm` instances. This allows seamless use of TorchComm with PyTorch's distributed tensor parallelism APIs.

### Parameters
- **mesh_dim_comms** (tuple[TorchComm, ...]) - Required - A tuple of TorchComm instances representing the dimensions of the mesh.
- **mesh_dim_names** (tuple[str, ...]) - Required - A tuple of names corresponding to each dimension in the mesh.
- **_global_comm** (Optional[TorchComm]) - Optional - An optional global TorchComm instance.

### Returns
- **dist.DeviceMesh** - The initialized PyTorch DeviceMesh instance.
```

--------------------------------

### torchcomms.TorchCommWindow

Source: https://meta-pytorch.org/torchcomms/main/api.html

Manages local buffers for communication windows.

```APIDOC
## torchcomms.TorchCommWindow

### Description
Interface for managing registered local buffers for communication.

### Methods
- **deregister_local_buffer(handle: SupportsInt) -> None**: Deregister a previously registered local buffer.
```

--------------------------------

### Send and Receive Object Lists

Source: https://meta-pytorch.org/torchcomms/main/_modules/torchcomms/objcol.html

Demonstrates the usage of send_object_list and recv_object_list to exchange a list of picklable objects between two ranks. Requires an initialized communication backend.

```python
>>> # xdoctest: +SKIP("need comm init")
>>> # Note: comm initialization omitted on each rank.
>>> from torchcomms import objcol
>>> # Assumes backend is not NCCL
>>> if comm.get_rank() == 0:
>>>     # Assumes world_size of 2.
>>>     objects = ["foo", 12, {1: 2}] # any picklable object
>>>     objcol.send_object_list(comm, objects, dst=1)
>>> else:
>>>     objects = [None, None, None]
>>>     objcol.recv_object_list(comm, objects, src=0)
>>> objects
['foo', 12, {1: 2}]
```

--------------------------------

### recv_object_list(comm, object_list, src, timeout=None, weights_only=True)

Source: https://meta-pytorch.org/torchcomms/main/_modules/torchcomms/objcol.html

Receives a list of picklable objects from a source rank synchronously.

```APIDOC
## recv_object_list(comm, object_list, src, timeout=None, weights_only=True)

### Description
Receives a list of picklable objects from a source rank. The provided object_list must have a length equal to the number of objects being sent.

### Parameters
- **comm** (TorchComm) - Required - The communication object to work on.
- **object_list** (list[Any]) - Required - List of objects to receive into.
- **src** (int) - Required - Source rank from which to receive.
- **timeout** (timedelta, optional) - Optional - Timeout for collective operations.
- **weights_only** (bool, optional) - Optional - If True, only safe objects such as weights are allowed to be deserialized.
```

--------------------------------

### Registering and managing custom hooks

Source: https://meta-pytorch.org/torchcomms/main/hooks.html

Demonstrates how to define, register, and remove pre, post, and abort hooks on a communicator. Hooks must not be modified while collective operations are in progress.

```python
import torch
import torchcomms
from torchcomms._comms import OpName, PreHookArgs, PostHookArgs

# Create a communicator
device = torch.device("cuda:0")
comm = torchcomms.new_comm("ncclx", device)

# Define hook callbacks
def my_pre_hook(args: PreHookArgs) -> None:
    print(f"Starting operation: {args.name} (op_id={args.op_id})")

def my_post_hook(args: PostHookArgs) -> None:
    print(f"Completed operation: {args.name} (op_id={args.op_id})")

def my_abort_hook() -> None:
    print("Process is about to abort, saving debug info...")

# Register hooks - each returns a RemovableHandle
pre_handle = comm.register_pre_hook(my_pre_hook)
post_handle = comm.register_post_hook(my_post_hook)
abort_handle = comm.register_abort_hook(my_abort_hook)

# Run collective operations - hooks will be called
tensor = torch.ones(10, device=device)
comm.all_reduce(tensor, torchcomms.ReduceOp.SUM, async_op=False)

# Unregister hooks when done
pre_handle.remove()
post_handle.remove()
abort_handle.remove()

comm.finalize()
```

--------------------------------

### Manage Asynchronous Operations with TorchWork

Source: https://meta-pytorch.org/torchcomms/main/api.html

Demonstrates how to use TorchWork to overlap communication with computation by enqueuing operations on a background stream and waiting for completion.

```python
tensor = ...
# run all_reduce on a background stream and return a TorchWork object
work = torchcomms.all_reduce(tensor, ReduceOp.SUM, async_op=True)

# Schedule some other work on the current stream
a = b * 2

# block the current stream until the all_reduce is complete
work.wait()

# safely use the tensor after the collective completes
tensor.sum()

# block CPU until stream is complete
torch.accelerator.current_stream().synchronize()
```

--------------------------------

### Gathering Objects Across Ranks

Source: https://meta-pytorch.org/torchcomms/main/_modules/torchcomms/objcol.html

Demonstrates gathering picklable objects from all ranks to a root rank using objcol.gather_object.

```python
>>> # xdoctest: +SKIP("need comm init")
>>> # Note: comm initialization omitted on each rank.
>>> from torchcomms import objcol
>>> # Assumes world_size of 3.
>>> gather_objects = ["foo", 12, {1: 2}] # any picklable object
>>> output = [None for _ in gather_objects]
>>> objcol.gather_object(
...     comm,
...     gather_objects[comm.get_rank()],
...     root=0,
...     object_gather_list=output,
... )
>>> # On rank 0
>>> output
['foo', 12, {1: 2}]
```

--------------------------------

### Implement a custom TorchComm backend

Source: https://meta-pytorch.org/torchcomms/main/api.html

Subclass TorchCommBackend and override required methods to define a custom communication backend.

```python
class MyBackend(torchcomms.TorchCommBackend)

def __init__(self):
    super().__init__() 

def init(self, device, name, options):
    self.rank = 0 self.size = 1 

def finalize(self):
    pass 

def get_rank(self):
    return self.rank 

def get_size(self):
    return self.size 

def get_backend_name(self):
    return “my_backend” 

def get_comm_name(self):
    return “my_comm” 

def all_reduce(self, tensor, op, async_op):
    return None # synchronous no-op
# … implement other methods … 
```

--------------------------------

### Gather objects across ranks with all_gather_object

Source: https://meta-pytorch.org/torchcomms/main/_modules/torchcomms/objcol.html

Demonstrates gathering picklable objects from all ranks into a list. Requires an initialized communicator and assumes a world size of 3.

```python
>>> # xdoctest: +SKIP("need comm init")
>>> # Note: comm initialization omitted on each rank.
>>> from torchcomms import objcol
>>> # Assumes world_size of 3.
>>> gather_objects = ["foo", 12, {1: 2}] # any picklable object
>>> output = [None for _ in gather_objects]
>>> objcol.all_gather_object(comm, output, gather_objects[comm.get_rank()])
>>> output
['foo', 12, {1: 2}]
```

--------------------------------

### send_object_list(comm, object_list, dst, timeout=None)

Source: https://meta-pytorch.org/torchcomms/main/_modules/torchcomms/objcol.html

Sends a list of picklable objects to a destination rank synchronously.

```APIDOC
## send_object_list(comm, object_list, dst, timeout=None)

### Description
Sends a list of picklable Python objects to a specified destination rank. This function serializes the objects into tensors before transmission.

### Parameters
- **comm** (TorchComm) - Required - The communication object to work on.
- **object_list** (list[Any]) - Required - A list of picklable objects to send.
- **dst** (int) - Required - The destination rank.
- **timeout** (timedelta, optional) - Optional - Timeout for collective operations.
```

--------------------------------

### Scatter Python objects across ranks

Source: https://meta-pytorch.org/torchcomms/main/_modules/torchcomms/objcol.html

Demonstrates how to use scatter_object_list to distribute picklable objects from a root rank to all other ranks in the communication group.

```python
>>> # xdoctest: +SKIP("need comm init")
>>> # Note: comm initialization omitted on each rank.
>>> from torchcomms import objcol
>>> if comm.get_rank() == 0:
>>>     # Assumes world_size of 3.
>>>     objects = ["foo", 12, {1: 2}] # any picklable object
>>> else:
>>>     # Can be any list on non-root ranks, elements are not used.
>>>     objects = [None, None, None]
>>> output_list = [None]
>>> objcol.scatter_object_list(comm, output_list, objects, root=0)
>>> # Rank i gets objects[i]. For example, on rank 2:
>>> output_list
[{1: 2}]
```

--------------------------------

### Create a TorchComm-backed ProcessGroup

Source: https://meta-pytorch.org/torchcomms/main/_modules/torchcomms/device_mesh.html

Helper function to register a TorchComm instance as a PyTorch ProcessGroup within the distributed runtime.

```python
def _create_torchcomm_process_group(
    comm: TorchComm,
    group_name: str,
    backend_str: str = "torchcomm",
    global_ranks_mapping: Optional[dict[int, int]] = None,
) -> dist.ProcessGroup:
    """
    Helper function to create a ProcessGroup backed by TorchComm and register it
    with the distributed runtime.

    Args:
        comm: TorchComm instance to wrap
        group_name: Name for the process group
        backend_str: Backend string identifier
        global_ranks_mapping: Mapping from global rank to group rank

    Returns:
        The created and registered ProcessGroup instance
    """
    # Make the linter happy. GroupName is just an alias for str. The cost of
    # this conversion is negligible.
    group_name = GroupName(group_name)

    wrapper = _BackendWrapper(comm)  # noqa: F405
    backend_type = dist.ProcessGroup.BackendType.CUSTOM
    backend_config = dist.BackendConfig(dist.Backend(backend_str))

    # Create process group. TorchComm backends handle all communication
    # directly, so no real store is needed. We use a dummy HashStore to
    # satisfy the ProcessGroup interface (e.g., splitGroup calls
    # store_->clone()).
    dummy_store = dist.HashStore()
    pg = dist.ProcessGroup(dummy_store, comm.get_rank(), comm.get_size())

    # Register backend
    # pyre-fixme[6]: BackendWrapper implements dist.Backend but types isn't aware
    pg._register_backend(comm.get_device(), backend_type, wrapper)
    pg._set_group_name(group_name)

    # Update global state
    dist.distributed_c10d._world.pg_map[pg] = (backend_str, dummy_store)
    dist.distributed_c10d._world.pg_names[pg] = group_name
    dist.distributed_c10d._world.pg_backend_config[pg] = str(backend_config)
    dist.distributed_c10d._register_process_group(group_name, pg)

    # Set up rank mapping
    if global_ranks_mapping is not None:
        dist.distributed_c10d._world.pg_group_ranks[pg] = global_ranks_mapping
    else:
        # Default mapping for global process groups
        dist.distributed_c10d._world.pg_group_ranks[pg] = {
            i: i for i in range(comm.get_size())
        }

    # Set up process group tag
    pg_tag = f"ptd:{group_name}"
    dist.distributed_c10d._world.tags_to_pg.setdefault(pg_tag, []).append(pg)
    dist.distributed_c10d._world.pg_to_tag[pg] = pg_tag

    return pg
```

--------------------------------

### new_window

Source: https://meta-pytorch.org/torchcomms/main/api.html

Creates a new window object for Remote Memory Access (RMA) operations.

```APIDOC
## new_window(self: torchcomms.TorchComm, tensor: torch.Tensor | None = None) -> torchcomms.TorchCommWindow

### Description
Create a new window object for Remote Memory Access (RMA) operations. Windows enable one-sided communication where data can be written directly to a remote rank’s buffer without receiver-side participation.

### Parameters
- **tensor** (torch.Tensor) - Optional - Contiguous tensor to register with the window. Must be allocated within a memory pool created via torchcomms.get_mem_allocator().

### Returns
- **TorchCommWindow** - Window object, registered if tensor was provided.

### Raises
- **RuntimeError** - If tensor is provided and a buffer is already registered.
```

--------------------------------

### scatter_object_list(comm, root, scatter_object_output_list, scatter_object_input_list=None, timeout=None, weights_only=True)

Source: https://meta-pytorch.org/torchcomms/main/_modules/torchcomms/objcol.html

Scatters a list of picklable objects from the root rank to all ranks in the communication group. Each rank receives one object, which is stored in the first element of the provided output list.

```APIDOC
## scatter_object_list

### Description
Scatters picklable objects from a source rank to all ranks in the communication group. On each rank, the received object is stored as the first element of the `scatter_object_output_list`.

### Parameters
- **comm** (TorchComm) - Required - The communication group to work on.
- **root** (int) - Required - The source rank from which to scatter the objects.
- **scatter_object_output_list** (list[Any]) - Required - A non-empty list where the first element will store the object scattered to this rank.
- **scatter_object_input_list** (list[Any], optional) - Optional - List of input objects to scatter. Only required on the root rank.
- **timeout** (timedelta, optional) - Optional - Timeout for the collective operation.
- **weights_only** (bool, optional) - Optional - If True, only safe objects such as weights are allowed to be deserialized.

### Returns
- **None** - The function modifies `scatter_object_output_list` in-place.
```

--------------------------------

### torchcomms.register_backend(name, backend_class)

Source: https://meta-pytorch.org/torchcomms/main/api.html

Registers a Python class as a TorchComm backend. The class must inherit from TorchCommBackend.

```APIDOC
## torchcomms.register_backend(name, backend_class)

### Description
Registers a Python class as a TorchComm backend. The provided class must inherit from TorchCommBackend and will be instantiated when a new communicator is created.

### Parameters
- **name** (str) - Required - Backend name used in torchcomms.new_comm.
- **backend_class** (object) - Required - A Python class that inherits from TorchCommBackend.
```

--------------------------------

### broadcast_object_list(comm, object_list, src, timeout=None, weights_only=True)

Source: https://meta-pytorch.org/torchcomms/main/api.html

Broadcasts a list of picklable objects from the source rank to all other ranks in the communication group.

```APIDOC
## broadcast_object_list

### Description
Broadcasts a list of picklable objects from the source rank to all other ranks in the communication group. This is a blocking call.

### Parameters
- **comm** (TorchComm) - Required - The communication group to work on.
- **object_list** (List[object]) - Required - List of input objects to broadcast. Each object must be picklable.
- **src** (int) - Required - Source rank from which to broadcast.
- **timeout** (timedelta) - Optional - Timeout for collective operations.
- **weights_only** (bool) - Optional - If True, only safe objects such as weights are allowed to be deserialized.

### Returns
- **None** - The object_list will contain the broadcasted objects on all ranks.
```

--------------------------------

### batch_op_create

Source: https://meta-pytorch.org/torchcomms/main/api.html

Creates a batch operation object for batched P2P operations.

```APIDOC
## batch_op_create(self: torchcomms.TorchComm) -> torchcomms.BatchSendRecv

### Description
Create a batch operation object for batched P2P operations.
```

--------------------------------

### get_device_window(signal_count, counter_count, barrier_count)

Source: https://meta-pytorch.org/torchcomms/main/api.html

Obtains a device-side window handle for GPU-initiated operations, allowing interaction with Triton kernels.

```APIDOC
## get_device_window(signal_count: SupportsInt = -1, counter_count: SupportsInt = -1, barrier_count: SupportsInt = 1) -> int

### Description
Get a device-side window handle for GPU-initiated operations. Returns a pointer (as int64) that can be passed to Triton kernels.

### Parameters
- **signal_count** (int) - Optional - Number of signal slots to allocate.
- **counter_count** (int) - Optional - Number of counter slots to allocate.
- **barrier_count** (int) - Optional - Number of barrier slots to allocate.

### Returns
- **int** - Device window pointer as int64.
```

--------------------------------

### torchcomms.objcol.recv_object_list

Source: https://meta-pytorch.org/torchcomms/main/api.html

Receives a list of picklable objects from a specified source rank.

```APIDOC
## torchcomms.objcol.recv_object_list(comm, object_list, src, timeout=None, weights_only=True)

### Description
Synchronously receives picklable objects into the provided object_list from a source rank.

### Parameters
- **comm** (TorchComm) - Required - The comm to work on.
- **object_list** (List[Any]) - Required - List of objects to receive into.
- **src** (int) - Required - Source rank from which to receive.
- **timeout** (timedelta) - Optional - Timeout for collective operations.
- **weights_only** (bool) - Optional - If True, only safe objects such as weights are allowed to be deserialized.
```

--------------------------------

### reconfigure

Source: https://meta-pytorch.org/torchcomms/main/api.html

Reconfigures the communicator with a new set of peers in a dynamic regime.

```APIDOC
## reconfigure(self: torchcomms.TorchComm, uuid: SupportsInt, init_handles: collections.abc.Set[str] | collections.abc.Sequence[str], timeout: datetime.timedelta | None = None, hints: collections.abc.Mapping[str, str] | None = None) -> torchcomms.TorchWork

### Description
Reconfigure the communicator with a new set of peers. In dynamic regime, this method initializes the communicator with the provided set of peers. After a successful reconfigure call, the communicator is fully initialized and collective operations are permitted.

### Parameters
- **uuid** (SupportsInt) - Required - Uniquely identifies this instance of the communicator.
- **init_handles** (Set[str] | Sequence[str]) - Required - The init handles of all ranks that will participate in this communicator.
- **timeout** (datetime.timedelta) - Optional - How long to allow reconfiguration to take before failing.
- **hints** (Mapping[str, str]) - Optional - Additional configuration key-value pairs.

### Returns
- **TorchWork** - A TorchWork handle that can be used to wait for completion.

### Raises
- **RuntimeError** - If not implemented by the backend.
```

--------------------------------

### all_to_all

Source: https://meta-pytorch.org/torchcomms/main/api.html

Scatter the split list to all ranks.

```APIDOC
## all_to_all

### Description
Scatter the split list to all ranks.

### Parameters
- **output_tensor_list** (collections.abc.Sequence[torch.Tensor]) - Required - Output tensor list
- **input_tensor_list** (collections.abc.Sequence[torch.Tensor]) - Required - Input tensor list to scatter
- **async_op** (bool) - Required - Whether to perform the operation asynchronously
- **hints** (collections.abc.Mapping[str, str]) - Optional - Dictionary of string hints for backend-specific options
- **timeout** (datetime.timedelta) - Optional - Timeout for the operation

### Returns
- **torchcomms.TorchWork** - The work object for the operation
```

--------------------------------

### Register and use FlightRecorderHook

Source: https://meta-pytorch.org/torchcomms/main/hooks.html

Tracks collective operations for debugging by recording metadata and timing in a ring buffer.

```python
import torch
import torchcomms
from torchcomms.hooks import FlightRecorderHook

# Create a communicator
device = torch.device("cuda:0")
comm = torchcomms.new_comm("ncclx", device)

# Create and register a flight recorder hook
recorder = FlightRecorderHook(max_entries=1024)
recorder.register_with_comm(comm)

# Run some collective operations
tensor = torch.ones(10, device=device)
comm.all_reduce(tensor, torchcomms.ReduceOp.SUM, async_op=False)
comm.barrier(async_op=False)

# Dump the recorded trace as JSON
json_trace = recorder.dump_json()
print(json_trace)

# Optionally dump to a file
recorder.dump_file(rank=comm.get_rank())

# Unregister when done
recorder.unregister()

# Finalize the communicator
comm.finalize()
```

--------------------------------

### Log collective operations with clog

Source: https://meta-pytorch.org/torchcomms/main/hooks.html

Initialize a clog logger and register it with multiple communicators to track operation signatures and lifecycle events.

```python
>>> from torchcomms.hooks import clog
>>> logger = clog(output="/tmp/clog.log", events=["ALL"])
>>> logger.register_with_comm(comm_a)
>>> logger.register_with_comm(comm_b)
>>> # ... run collectives ...
```

--------------------------------

### send_object_list

Source: https://meta-pytorch.org/torchcomms/main/_modules/torchcomms/objcol.html

Sends a list of picklable objects synchronously to a destination rank.

```APIDOC
## send_object_list

### Description
Sends a list of picklable objects synchronously to a destination rank. All objects in the list must be picklable.

### Parameters
- **comm** (TorchComm) - The communication object to work on.
- **object_list** (list[Any]) - List of input objects to send.
- **dst** (int) - Destination rank to send the object list to.
- **timeout** (timedelta, optional) - Timeout for collective operations.
```

--------------------------------

### torchcomms.objcol.broadcast_object_list

Source: https://meta-pytorch.org/torchcomms/main/api.html

Broadcasts a list of picklable objects to all ranks in the comm.

```APIDOC
## torchcomms.objcol.broadcast_object_list(comm, object_list, root, timeout=None, weights_only=True)

### Description
Broadcasts picklable objects in object_list to the whole comm.

### Parameters
- **comm** (TorchComm) - Required - The comm to work on.
- **object_list** (List[Any]) - Required - List of objects to broadcast.
- **root** (int) - Required - The root rank.
- **timeout** (timedelta) - Optional - Timeout for collective operations.
- **weights_only** (bool) - Optional - If True, only safe objects such as weights are allowed to be deserialized.
```

--------------------------------

### Sending Object Lists Synchronously

Source: https://meta-pytorch.org/torchcomms/main/_modules/torchcomms/objcol.html

Sends a list of picklable objects to a destination rank. Requires the receiver to provide a list of equal size.

```python
def send_object_list(
    comm: TorchComm,
    object_list: list[Any],
    dst: int,
    timeout: timedelta | None = None,
) -> None:
    """
    Sends picklable objects in ``object_list`` synchronously.

    Similar to :func:`send`, but Python objects can be passed in.
    Note that all objects in ``object_list`` must be picklable in order to be
    sent.

    Args:
        comm: The comm to work on.
        object_list (List[object]): List of input objects to sent.
            Each object must be picklable. Receiver must provide lists of equal sizes.
        dst (int): Destination rank to send ``object_list`` to.
        timeout: (timedelta, optional): Timeout for collective operations. If
            ``None``, will use the default timeout for the backend.
    Returns:
        ``None``.

    .. note:: For NCCL-based comms, internal tensor representations
        of objects must be moved to the GPU device before communication takes
        place. In this case, the device used is given by
        ``torch.cuda.current_device()`` and it is the user's responsibility to
        ensure that this is set so that each rank has an individual GPU, via
        ``torch.cuda.set_device()``.

    .. warning::
        Object collectives have a number of serious performance and scalability
        limitations.

    .. warning::
        :func:`send_object_list` uses ``pickle`` module implicitly, which
        is known to be insecure. It is possible to construct malicious pickle
        data which will execute arbitrary code during unpickling. Only call this
        function with data you trust.

    .. warning::
        Calling :func:`send_object_list` with GPU tensors is not well supported
        and inefficient as it incurs GPU -> CPU transfer since tensors would be
        pickled. Please consider using :func:`send` instead.

    """
```

--------------------------------

### broadcast_object_list(comm, object_list, root, timeout=None, weights_only=True)

Source: https://meta-pytorch.org/torchcomms/main/_modules/torchcomms/objcol.html

Broadcasts a list of picklable Python objects from a root rank to all other ranks in the communication group. This is a blocking operation.

```APIDOC
## broadcast_object_list

### Description
Broadcasts picklable objects in `object_list` to the whole comm. This function is similar to standard broadcast collectives but supports Python objects by serializing them.

### Parameters
- **comm** (TorchComm) - Required - The communication group to work on.
- **object_list** (list[Any]) - Required - List of input objects to broadcast. Each object must be picklable. Each rank must provide lists of equal sizes.
- **root** (int) - Required - Source rank from which to broadcast `object_list`.
- **timeout** (timedelta, optional) - Optional - Timeout for collective operations. If None, uses the default backend timeout.
- **weights_only** (bool, optional) - Optional - If True, only safe objects such as weights are allowed to be deserialized.

### Returns
- **None** - The `object_list` on each rank will be updated with the broadcasted objects from the root rank.

### Example
```python
from torchcomms import objcol
# Assumes comm is initialized
if comm.get_rank() == 0:
    objects = ["foo", 12, {1: 2}]
else:
    objects = [None, None, None]
objcol.broadcast_object_list(comm, objects, root=0)
# objects now contains ['foo', 12, {1: 2}] on all ranks
```
```

--------------------------------

### gather_object(comm, obj, root, object_gather_list=None, timeout=None, weights_only=True)

Source: https://meta-pytorch.org/torchcomms/main/_modules/torchcomms/objcol.html

Gathers picklable objects from all ranks in the communication group to a single root process. This is a blocking call that serializes objects, performs the gather, and deserializes them on the root rank.

```APIDOC
## gather_object

### Description
Gathers picklable objects from the whole comm in a single process. Note that the object must be picklable in order to be gathered.

### Parameters
- **comm** (TorchComm) - Required - The comm to work on.
- **obj** (object) - Required - Input object. Must be picklable.
- **root** (int) - Required - Destination rank on comm.
- **object_gather_list** (list[object]) - Optional - Output list. On the root rank, it should be correctly sized as the size of the comm for this collective and will contain the output. Must be None on non-root ranks.
- **timeout** (timedelta) - Optional - Timeout for collective operations. If None, will use the default timeout for the backend.
- **weights_only** (bool) - Optional - If True, only safe objects such as weights are allowed to be deserialized.

### Returns
- **None** - On the root rank, object_gather_list will contain the output of the collective.
```

--------------------------------

### signal(peer_rank, async_op, hints=None, timeout=None)

Source: https://meta-pytorch.org/torchcomms/main/api.html

Atomically sends a signal to notify a remote peer that data is ready. This is used in conjunction with wait_signal to synchronize operations.

```APIDOC
## signal(peer_rank, async_op, hints=None, timeout=None)

### Description
Atomically send a signal to notify a remote peer that data is ready. Used with wait_signal to synchronize.

### Parameters
- **peer_rank** (int) - Required - Destination rank to signal.
- **async_op** (bool) - Required - If True, returns TorchWork for async completion.
- **hints** (Dict[str, str]) - Optional - Backend-specific hints.
- **timeout** (timedelta) - Optional - Operation timeout.

### Returns
- **TorchWork** - Work object for synchronization.
```
Tokens:
Raw