tensorplay.distributed.nn

Autograd-aware collectives

The functions in tensorplay.distributed.nn.functional wrap the blocking collectives of tensorplay.distributed in autograd Functions, so they can appear anywhere in a forward pass and still backpropagate. Their gradients are the mathematical inverses of the collectives, applied to the gradient that flows in:

• broadcast() — backward reduces the gradients (SUM) back to the source rank and zeroes them elsewhere, so only the source accumulates.

• reduce() — backward broadcasts the reduced result’s gradient to every rank.

• gather() / scatter() are inverse to each other: gather’s backward scatters the per-rank gradient slices, and scatter’s backward gathers them.

• all_reduce() — backward applies the same all-reduce to the gradient, since every rank needs the sum.

• all_gather() — backward all-reduces the incoming output gradients and keeps the slice this rank contributed.

• reduce_scatter() — backward all-gathers the scattered gradient.

• all_to_all() (and its single-output variant) route gradient slices back along the exchange pattern of the forward.

import tensorplay as tp
import tensorplay.distributed as dist
import tensorplay.distributed.nn.functional as df

dist.init_process_group("nccl")

x = tp.randn(4, 8, device="cuda", requires_grad=True)
y = df.all_reduce(x)          # differentiable all-reduce
y.sum().backward()            # gradient flows back through the collective

Remote modules

RemoteModule constructs a module of your choosing on another worker and hands you a local handle: calling forward (or any generated method) transparently issues an RPC to the remote module, moves the arguments there, and returns the result. The remote_device string says where the module lives: "worker_name/device" or "rank:N/device" — for example "trainer1/cuda:0". The constructor takes the module class and the arguments its own constructor needs; the instantiation itself happens on the remote worker.

import tensorplay as tp
import tensorplay.distributed.rpc as rpc
from tensorplay.distributed.nn import RemoteModule

rpc.init_rpc(f"driver", rank=0, world_size=2)

# an nn.Linear now lives on the remote worker, on its cuda:0
remote_linear = RemoteModule(
    "worker1/cuda:0",
    tp.nn.Linear,
    args=(16, 8),
)

out = remote_linear(tp.randn(4, 16))          # runs on worker1

param_rrefs = remote_linear.remote_parameters()  # RRefs, not tensors
module_rref = remote_linear.get_module_rref()

rpc.shutdown()

interface() is the companion decorator: mark a class with it to define the exact set of remote methods a remote module exposes, validated at instantiation time.

RemoteModule implements the usual nn.Module surface, but the operations that would mutate local state (cuda, to, load_state_dict, register_buffer, …) raise, because there is no local module to mutate — the state lives on the remote worker. Use remote_parameters() to get RRefs to the parameters for a DistributedOptimizer, or get_module_rref() to pass the module itself into further RPC calls.

Where to go next

• RPC — the transport remote modules are built on.

• the distributed package — the blocking collectives these functions differentiate.

• distributed autograd — accumulating gradients across workers when the forward pass spans RPC calls.