tvm.relax.distributed

The infrastructure for distributed inference on Relax.

class tvm.relax.distributed.DeviceMesh(shape: Shape | list[int] | tuple[int], device_ids: list[int] | Range)
Device mesh express a view of topology of devices,

represented by an n-d matrix of device ids.

Parameters:

  • shape (Union[ShapeTuple, List[int], Tuple[int]]) – Logical shape of device mesh

  • device_ids (Union[List[int], Range]) – Represents the device id in the mesh

tvm.relax.distributed.device_mesh(shape: Shape, device_ids: list[int] | Range) → DeviceMesh

Create a device mesh expression. :param shape: The shape of the device mesh. :type shape: ShapeTuple :param device_ids: Represents the device id in the mesh :type device_ids: Union[List[int], Range]

Returns:

res – The device mesh.

Return type:

DeviceMesh

class tvm.relax.distributed.Placement(dim_specs: list[PlacementSpec])

Describes how data is distributed in each dimension of the device mesh

Parameters:

dim_specs (List[PlacementSpec]) – The placement spec for each dimension of the device mesh.

static from_text(text: str) → Placement

Create a placement from a text string.

Parameters:

text (str) – The text string.

Returns:

placement – The placement.

Return type:

Placement

class tvm.relax.distributed.DTensorStructInfo(tensor_sinfo: TensorStructInfo, device_mesh: DeviceMesh, placement: Placement, span: Span | None = None)

StructInfo of a Distributed Tensor value.

Parameters:

  • tensor_sinfo (TensorStructInfo) – The struct info inherited from TensorStructInfo

  • device_mesh (DeviceMesh) – The device mesh of the tensor.

  • placement (Placement) – The placement of the tensor among the device mesh

class tvm.relax.distributed.PlacementSpec(*args, **kwargs)

Describes how data is distributed in one dimension of the device mesh

Parameters:

  • axis (int) – If the kind is sharding, this value represents the tensor dimension to shard. otherwise, axis is -1

  • kind (PlacementSpecKind) – The kind of placement spec. Possible values: kSharding and kReplica.

static sharding(axis: int) → PlacementSpec

Create a sharding placement spec

Parameters:

axis (int) – The tensor dimension to shard.

Returns:

placement_spec – The placement spec.

Return type:

PlacementSpec

static replica() → PlacementSpec

Create a replica placement spec

Returns:

placement_spec – The placement spec.

Return type:

PlacementSpec