tensorplay.distributed.tensor

Creating distributed tensors

• distribute_tensor() wraps an existing tensor into a DTensor with the given placements, redistributing the data if the current layout on the mesh differs from the requested one.

• distribute_module() converts an entire module into its distributed form: a partition function decides, per parameter, whether it is sharded or left replicated, and optional input and output functions transform the tensors that enter and leave the module.

• from_local() builds a DTensor out of the local shards already present on each rank, without any cross-rank movement of data (optionally checking that the local shards match the requested global shape).

Factory functions

The factory functions create DTensor values directly on a mesh. They accept the usual construction arguments (dtype, layout, requires_grad) plus device_mesh and placements:

For example, creating an 8-element replicated vector on a 4-rank mesh keeps a full copy on every rank, while Shard(0) distributes it one element per pair of ranks:

from tensorplay.distributed.tensor import Shard, ones

# requires an initialized process group, e.g. via dist.init_process_group()
mesh = ...   # a DeviceMesh, see distributed.device_mesh
local = ones(8, device_mesh=mesh, placements=[Shard(0)])

Placements

A placement tells the runtime how one dimension of the mesh participates in storing the tensor:

• Shard splits the tensor along a logical dimension across the ranks of the corresponding mesh dimension.

• Replicate stores a full copy of the tensor on every rank of the corresponding mesh dimension.

• Partial marks a tensor that has been only partially reduced on each rank; the reduce_op names the operation (e.g. "sum") that completes the reduction when a full value is needed.

The DTensor class

The class itself exposes the shard on the current rank through to_local() (moving the tensor across devices if the shard is on a different accelerator than the local one), the fully materialized value through full_tensor(), and a re-layout operation through redistribute() that converts between placements, possibly asynchronously. Its metadata follows a plain tensor: shape, size, stride, ndim, numel, dtype, device, plus the device_mesh and placements that describe its layout.

Where to go next

• device mesh — the mesh abstraction that DTensor runs on.

• tensor parallelism — sharding entire modules with parallelize_module.

• the distributed package — process groups, collectives, and the key-value stores underneath.