Garden | PyTorch 原生并行：从 Tensor 到 DTensor

Quick Start

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33


import torch
import torch.distributed as dist
from torch.distributed.tensor import DTensor, Shard, Replicate, Partial
from torch.distributed.tensor import distribute_tensor, distribute_module
from torch.distributed.device_mesh import DeviceMesh, init_device_mesh

import os
import torch.multiprocessing as mp

def main(rank, world_size):
    os.environ["NCCL_DEBUG"] = "WARN"
    os.environ["MASTER_ADDR"] = "127.0.0.1"
    os.environ["MASTER_PORT"] = "39500"
    os.environ["RANK"] = str(rank)
    os.environ["WORLD_SIZE"] = str(world_size)
    torch.cuda.set_device(rank)

    mesh = init_device_mesh("cuda", (world_size,), mesh_dim_names=("dp",))
    x = torch.randn(8192, 8192, device=f"cuda:{rank}")
    y = distribute_tensor(x, mesh, [Shard(0)])

    z = torch.sin(y)

    if rank == 0:
        print(f"Input tensor shape: {x.shape}")
        print(f"Distributed tensor y: placements={y.placements}, shape={y.shape}, _local_tensor={y._local_tensor.shape}")
        print(f"Result z: placements={z.placements}, shape={z.shape}, _local_tensor={z._local_tensor.shape}")

    dist.destroy_process_group()

if __name__ == "__main__":
    world_size = torch.cuda.device_count()
    mp.spawn(main, args=(world_size,), nprocs=world_size, join=True)

1
2
3
4


$ python3 dtensor.py
Input tensor shape: torch.Size([8192, 8192])
Distributed tensor y: placements=(Shard(dim=0),), shape=torch.Size([8192, 8192]), _local_tensor=torch.Size([1024, 8192])
Result z: placements=(Shard(dim=0),), shape=torch.Size([8192, 8192]), _local_tensor=torch.Size([1024, 8192])

从 Tensor 到 DTensor

Tensor

普通 torch.Tensor 可以先理解成「一段数据 + 一组元信息」：

数据本体存放在某个 device 上，例如 CPU 或一张 GPU。
元信息描述这段数据应该如何被解释，例如 shape、dtype、stride、requires_grad。
算子看到的是一个单设备张量，所有输入、输出和中间结果默认也都落在这个单设备语义里。

PyTorch 内部会用类似 TensorMetadata 的结构描述一个 Tensor 的逻辑属性：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14


class TensorMetadata(NamedTuple):  
    # TensorMetadata is a structure containing pertinent information  
    # about a tensor within a PyTorch program.
    
    # General Tensor metadata
    shape: torch.Size  
    dtype: torch.dtype  
    requires_grad: bool  
    stride: Tuple[int, ...]  
    memory_format: Optional[torch.memory_format]  
  
    # Quantization metadata  
    is_quantized: bool  
    qparams: Dict[str, Any]

当 Tensor 扩展到多设备时，仅有这些元信息是不够的。系统还必须回答几个分布式问题：

这组设备如何组织？例如是一维数据并行 mesh，还是二维 dp x tp mesh。
全局 Tensor 如何映射到每个 rank？例如按第 0 维切分、每张卡保存一份副本，或每张卡只保存部分归约结果。
算子执行后新的分布式布局是什么？哪些算子可以纯本地计算，哪些算子需要插入 all-gather、all-reduce、reduce-scatter 等通信。
用户代码应该看到全局 Tensor 语义，还是每个 rank 的局部 Tensor 语义。

DTensor

DTensor 的核心目标是把「多设备分布式张量」包装成一个仍然像 torch.Tensor 一样使用的对象：

1

DTensor = global Tensor metadata + local Tensor per rank + distributed layout

其中最重要的是区分两层视角：

Global view：用户代码看到的逻辑 Tensor，shape、dtype、算子语义都按完整 Tensor 理解。
Local view：每个 rank 实际持有的 _local_tensor，它可能只是 global Tensor 的一个 shard，也可能是完整副本，或者是尚未归约完成的 partial 结果。

因此，一个 DTensor 不只是多了几张卡上的存储，它还显式记录了分布式布局：

概念	作用
`DeviceMesh`	描述参与这个 DTensor 的设备拓扑，以及每个 mesh 维度对应的通信组
`Placement`	描述 Tensor 在每个 mesh 维度上的放置方式，例如 `Shard`、`Replicate`、`Partial`
`DTensorSpec`	把 `DeviceMesh`、`Placement` 和 Tensor 元信息组合起来，完整描述 DTensor 的布局

从编程模型看，DTensor 提供的是 SPMD 语义：所有 rank 运行同一份 Python 代码，但每个 rank 根据相同的 DTensorSpec 操作自己的 local shard。算子分发时，DTensor 会基于输入布局推导输出布局，并在必要时自动插入集合通信，让用户尽量按照单机 Tensor 的方式写代码。

DeviceMesh - Describing Device Topology

DeviceMesh 提供表达一组 device 布局的抽象，可以用一个多维数组表达，同时也提供 Mesh 内 device 通信的支持。

可以通过 init_device_mesh 来初始化一个 DeviceMesh：

1
2


from torch.distributed.device_mesh import init_device_mesh
mesh_1d = init_device_mesh("cuda", (8,))

对应可视化

1
2
3


┌───────────────────────────────────────────────────────────────┐
│ GPU 0 │ GPU 1 │ GPU 2 │ GPU 3 │ GPU 4 │ GPU 5 │ GPU 6 │ GPU 7 │
└───────────────────────────────────────────────────────────────┘

1
2


from torch.distributed.device_mesh import init_device_mesh
mesh_2d = init_device_mesh("cuda", (2, 4), mesh_dim_names=("dp", "tp"))

对应可视化：

1
2
3
4
5
6


              tp dimension (4 devices)
         ┌─────────────────────────────────┐
         │ GPU 0 │ GPU 1 │ GPU 2 │ GPU 3   │  dp=0
dp       ├─────────────────────────────────┤
dimension│ GPU 4 │ GPU 5 │ GPU 6 │ GPU 7   │  dp=1
         └─────────────────────────────────┘

访问 sub-meshes:

1
2
3
4
5
6
7


# Users can access the underlying process group thru `get_group` API.
dp_group = mesh_2d.get_group(mesh_dim="dp")
tp_group = mesh_2d.get_group(mesh_dim="tp")

# 或者直接访问
dp_mesh = mesh_2d["dp"]
tp_mesh = mesh_2d["tp"]

Placement - Describing Tensor Distribution

Placements describe how a tensor is distributed across each dimension of the DeviceMesh.

一个 DTensor 的 placements 数量必须和 DeviceMesh 的维度数一致。例如一维 mesh 只需要一个 placement，二维 mesh 则需要两个 placement：

1
2
3
4
5


mesh_1d = init_device_mesh("cuda", (4,), mesh_dim_names=("tp",))
dtensor_1d = distribute_tensor(x, mesh_1d, [Shard(0)])

mesh_2d = init_device_mesh("cuda", (2, 4), mesh_dim_names=("dp", "tp"))
dtensor_2d = distribute_tensor(x, mesh_2d, [Replicate(), Shard(1)])

可以把 placements 理解成对每个 mesh 维度的逐项说明：

Mesh 维度	Placement	含义
`dp`	`Replicate()`	沿数据并行维度复制完整 Tensor
`tp`	`Shard(1)`	沿张量并行维度切分 Tensor 的第 1 维

DTensor 支持三类基础 placement：Shard、Replicate 和 Partial。

Shard(dim)

Shard(dim) 表示沿 Tensor 的第 dim 个维度切分数据，并把切分后的 shards 分发到 mesh 维度上的不同 rank。

1
2
3
4


from torch.distributed.tensor import Shard, distribute_tensor

x = torch.randn(8, 4)
sharded = distribute_tensor(x, mesh_1d_4, [Shard(0)])

逻辑上，用户仍然看到一个完整的 [8, 4] global Tensor：

1
2
3
4
5
6


Before: global tensor, shape = [8, 4]

┌─────────────┐
│   8 rows    │
│   4 cols    │
└─────────────┘

如果在 4 张 GPU 上沿第 0 维切分，每个 rank 只持有 [2, 4] 的 local shard：

1
2
3
4
5
6


After: Shard(0) across 4 GPUs

GPU 0: [2, 4]     GPU 1: [2, 4]     GPU 2: [2, 4]     GPU 3: [2, 4]
┌──────────┐      ┌──────────┐      ┌──────────┐      ┌──────────┐
│ rows 0-1 │      │ rows 2-3 │      │ rows 4-5 │      │ rows 6-7 │
└──────────┘      └──────────┘      └──────────┘      └──────────┘

Shard 是节省显存的主要手段：global Tensor 的总数据量被拆到多个 rank 上。但它也意味着某些算子需要通信。例如，当一个算子要求完整 Tensor 时，DTensor 可能需要先 all-gather；当输出仍可保持分片时，则可以只在 local shard 上计算。

Replicate()

Replicate() 表示 mesh 维度上的每个 rank 都保存一份完整 Tensor 副本。

1
2
3
4


from torch.distributed.tensor import Replicate, distribute_tensor

x = torch.randn(8, 4)
replicated = distribute_tensor(x, mesh_1d_4, [Replicate()])

复制后的每个 rank 都持有完整 [8, 4]：

1
2
3
4
5
6
7


Result: Replicate() across 4 GPUs

GPU 0: [8, 4]     GPU 1: [8, 4]     GPU 2: [8, 4]     GPU 3: [8, 4]
┌─────────┐       ┌─────────┐       ┌─────────┐       ┌─────────┐
│  full   │       │  full   │       │  full   │       │  full   │
│  copy   │       │  copy   │       │  copy   │       │  copy   │
└─────────┘       └─────────┘       └─────────┘       └─────────┘

Replicate 的优点是读访问简单，很多 elementwise 算子可以直接在每个 rank 上独立执行；缺点是显存占用随副本数线性增加。数据并行里的参数副本、或者需要广播到所有 rank 的小张量，通常适合用 Replicate 表达。

Partial(reduce_op)

Partial(reduce_op) 表示每个 rank 持有的是某个 global Tensor 的「部分结果」，这些 partial values 需要通过归约操作才能变成完整语义的 DTensor。

1
2
3


from torch.distributed.tensor import DTensor, Partial

partial = DTensor.from_local(local_tensor, mesh, [Partial("sum")])

常见的 reduce_op 包括 "sum"、"avg"、"product"、"max"、"min"。其中最常见的是 "sum"，例如矩阵乘法或线性层在某个维度被切分后，每个 rank 只计算了输出的一部分累加项：

1
2
3
4
5
6


Each rank holds partial output with the same logical shape

GPU 0: partial result ┐
GPU 1: partial result ├── all-reduce(sum) ──> replicated full result
GPU 2: partial result │
GPU 3: partial result ┘

Partial 通常是算子传播过程中的中间布局，而不是用户手动构造数据时最常用的布局。它的核心意义是把「尚未完成归约」这件事显式记录在 DTensor 的 layout 中，这样后续算子可以决定是继续保留 partial 状态，还是在需要完整值时触发 all-reduce。

Multi-Dimensional Placements

对于多维 DeviceMesh，placements 是一个 tuple/list，每一项对应一个 mesh 维度。下面的例子使用二维 mesh：

1
2
3
4


mesh_2d = init_device_mesh("cuda", (2, 4), mesh_dim_names=("dp", "tp"))

x = torch.randn(16, 8)
distributed_x = distribute_tensor(x, mesh_2d, [Replicate(), Shard(1)])

含义是：

沿 dp 维度使用 Replicate()：两组数据并行 rank 拥有相同的逻辑数据。
沿 tp 维度使用 Shard(1)：每个 dp 组内部再按 Tensor 的第 1 维做列切分。

可视化如下：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10


Global tensor shape = [16, 8]
placements = [Replicate(), Shard(1)]

                    tp dimension: Shard(1)
              ┌──────────┬──────────┬──────────┬──────────┐
dp = 0        │ cols 0-1 │ cols 2-3 │ cols 4-5 │ cols 6-7 │
              ├──────────┼──────────┼──────────┼──────────┤
dp = 1        │ cols 0-1 │ cols 2-3 │ cols 4-5 │ cols 6-7 │
              └──────────┴──────────┴──────────┴──────────┘
              replicated across dp dimension

这里同一列位置上的两个 dp rank 互为 replica，而同一行里的四个 tp rank 共同组成一个按列切分的 global Tensor。

作为对比， PyTorch 的 DTensor 和 OneFlow¹ 以及 GSPMD² 定义的区别：

PT-D DistributedTensor	OneFlow’s SBP	GSPMD’s tensor sharding
Shard	Split	Tiled
Replicate	Broadcast	Replicated
Partial	Partial	Partially tiled = tiled + replicated

DTensorSpec

DTensorSpec 完全表达了一个 DTensor 的元信息，分别由以下三个部分组成：

DeviceMesh 对象来表达 DTensor 的 mesh 信息，
Tuple[Placement] 来表达 placements 方法
TensorMetadata 对象来表达 global tensor 的 meta 信息

1
2
3
4
5
6


class DTensorSpec:  
    mesh: DeviceMesh  
    placements: Tuple[Placement, ...]  
  
    # tensor meta will only be set during sharding propagation  
    tensor_meta: Optional[TensorMeta] = None

实际举例：

1

DTensorSpec(mesh=DeviceMesh:([0, 1]), placements=[Shard(dim=0)], tensor_meta=TensorMetadata(shape=torch.Size([6, 3]), dtype=torch.int64, requires_grad=False, stride=(3, 1), memory_format=None, is_quantized=False, qparams={}))

DTensor

Torch 的 DTensor 在 torch.Tensor 类型上进行了简单的封装：

1
2
3
4
5
6
7


class DTensor(torch.Tensor):  
    _local_tensor: torch.Tensor  
    _spec: DTensorSpec  
    __slots__ = ["_local_tensor", "_spec"]  
  
    # _op_dispatcher instance as a class attribute to handle runtime dispatching logic  
    _op_dispatcher: op_dispatch.OpDispatcher = op_dispatch.OpDispatcher()

主要包括 _local_tensor 和 _spec：

_local_tensor 是实际存储的 torch.Tensor 变量 (per rank)。
_spec 中存储了 DTensor 的全部元信息，对应于 DTensorSpec 字段
- 包括 DeviceMesh、切分策略（Placements）以及传统 tensor 的属性信息，例如 shape、dtype 等

Creating DTensor

创建 DTensor 最常见有两条路径：

distribute_tensor()：从一个 global logical tensor 出发，由 DTensor 负责 scatter / broadcast 到各个 rank。
DTensor.from_local()：从每个 rank 已经存在的 local tensor 出发，告诉 DTensor 这些 local tensor 共同组成什么 global layout。

两者的区别在于数据来源不同：

API	输入 Tensor 语义	典型场景
`distribute_tensor()`	输入被当成 global tensor，rank 0 通常作为 source of truth	初始化参数、分发输入、把普通 Tensor 转成 DTensor
`DTensor.from_local()`	输入就是当前 rank 的 local shard / replica / partial	算子中间结果、已有分片权重、手动构造 local shard

Method 1: `distribute_tensor()`

distribute_tensor() 适合从一个普通 torch.Tensor 创建 DTensor。它会根据 placements 把 global tensor 分发到 DeviceMesh 上：

1
2
3
4


from torch.distributed.tensor import Shard, distribute_tensor

global_tensor = torch.arange(32, device=device).reshape(8, 4)
dtensor = distribute_tensor(global_tensor, mesh, [Shard(0)])

如果 mesh 里有 4 个 rank，Shard(0) 会把 [8, 4] 的 global tensor 按第 0 维切成 4 份，每个 rank 的 local tensor 形状是 [2, 4]。

需要注意的是，distribute_tensor() 保持的是 single-device semantic：逻辑上它从一个完整 Tensor 出发，然后由 DTensor runtime 负责在 mesh 内 scatter / broadcast。实践中应把 rank 0 上的输入视为 source of truth，其他 rank 上的输入值不应该承载额外语义。

Method 2: `DTensor.from_local()`

DTensor.from_local() 适合每个 rank 已经有自己的 local tensor 的情况：

1
2
3
4
5
6
7
8
9


from torch.distributed.tensor import DTensor, Shard

local_tensor = torch.full((2, 4), fill_value=rank, device=device)
dtensor = DTensor.from_local(
    local_tensor,
    device_mesh=mesh,
    placements=[Shard(0)],
    run_check=True,
)

这里每个 rank 都提供一个 [2, 4] 的 local shard。如果 world_size = 4，那么 DTensor 的 global shape 会被解释为 [8, 4]。

run_check 的含义：

run_check=True：DTensor 会做一致性检查，例如 local tensor 的 shape / stride 是否能组成合法的 global tensor。更安全，但会引入额外通信。
run_check=False：不做检查，直接相信调用方提供的 local tensor 和 placement 是正确的。更快，但调用方必须自己保证每个 rank 的 local shard 合法。

单脚本可运行示例

下面这个脚本不依赖 torchrun，直接用 torch.multiprocessing.spawn 启动多进程。机器有 GPU 时使用 cuda + nccl，否则退化到 cpu + gloo。

保存为 dtensor_create.py 后直接运行：

1

python3 dtensor_create.py

完整脚本：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85


import os
import socket

import torch
import torch.distributed as dist
import torch.multiprocessing as mp
from torch.distributed.device_mesh import init_device_mesh
from torch.distributed.tensor import DTensor, Replicate, Shard, distribute_tensor


def find_free_port() -> str:
    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
        sock.bind(("127.0.0.1", 0))
        return str(sock.getsockname()[1])


def describe_dtensor(name: str, dtensor: DTensor) -> str:
    local = dtensor.to_local()
    return (
        f"{name}: global_shape={tuple(dtensor.shape)}, "
        f"placements={dtensor.placements}, "
        f"local_shape={tuple(local.shape)}, "
        f"local_value=\n{local.cpu()}"
    )


def setup_dist(rank: int, world_size: int, backend: str, port: str) -> None:
    os.environ["MASTER_ADDR"] = "127.0.0.1"
    os.environ["MASTER_PORT"] = port
    os.environ["RANK"] = str(rank)
    os.environ["WORLD_SIZE"] = str(world_size)
    dist.init_process_group(backend=backend, rank=rank, world_size=world_size)


def run(rank: int, world_size: int, backend: str, device_type: str, port: str) -> None:
    setup_dist(rank, world_size, backend, port)

    if device_type == "cuda":
        torch.cuda.set_device(rank)
        device = torch.device("cuda", rank)
    else:
        device = torch.device("cpu")

    mesh = init_device_mesh(device_type, (world_size,), mesh_dim_names=("tp",))

    # Method 1: distribute a logical global tensor.
    # rank 0 is the source of truth for the data to scatter / broadcast.
    global_tensor = torch.arange(world_size * 8, device=device).reshape(world_size * 2, 4)
    sharded = distribute_tensor(global_tensor, mesh, [Shard(0)])
    replicated = distribute_tensor(global_tensor, mesh, [Replicate()])

    # Method 2: wrap pre-existing local shards.
    local_shard = torch.full((2, 4), fill_value=rank, dtype=torch.float32, device=device)
    from_local = DTensor.from_local(
        local_shard,
        device_mesh=mesh,
        placements=[Shard(0)],
        run_check=True,
    )

    # Print rank by rank to avoid interleaved output.
    for print_rank in range(world_size):
        dist.barrier()
        if rank == print_rank:
            print(f"\n===== rank {rank} =====")
            print(describe_dtensor("distribute_tensor + Shard(0)", sharded))
            print(describe_dtensor("distribute_tensor + Replicate()", replicated))
            print(describe_dtensor("DTensor.from_local + Shard(0)", from_local))
        dist.barrier()

    dist.destroy_process_group()


if __name__ == "__main__":
    if torch.cuda.is_available():
        device_type = "cuda"
        backend = "nccl"
        world_size = torch.cuda.device_count()
    else:
        device_type = "cpu"
        backend = "gloo"
        world_size = 4

    port = find_free_port()
    mp.spawn(run, args=(world_size, backend, device_type, port), nprocs=world_size, join=True)

输出里重点看三件事：

global_shape：用户看到的逻辑 Tensor shape，所有 rank 一致。
placements：DTensor 当前的分布式布局。
local_shape / local_value：当前 rank 实际持有的数据。

在 4 个 rank 上，distribute_tensor + Shard(0) 的输出大致是：

1
2
3
4
5
6
7
8
9


===== rank 0 =====
distribute_tensor + Shard(0): global_shape=(8, 4), placements=(Shard(dim=0),), local_shape=(2, 4)
tensor([[0, 1, 2, 3],
        [4, 5, 6, 7]])

===== rank 1 =====
distribute_tensor + Shard(0): global_shape=(8, 4), placements=(Shard(dim=0),), local_shape=(2, 4)
tensor([[ 8,  9, 10, 11],
        [12, 13, 14, 15]])

而 DTensor.from_local + Shard(0) 会保留每个 rank 预先构造的 local shard：

1
2
3
4
5
6
7
8
9


===== rank 0 =====
DTensor.from_local + Shard(0): global_shape=(8, 4), placements=(Shard(dim=0),), local_shape=(2, 4)
tensor([[0., 0., 0., 0.],
        [0., 0., 0., 0.]])

===== rank 1 =====
DTensor.from_local + Shard(0): global_shape=(8, 4), placements=(Shard(dim=0),), local_shape=(2, 4)
tensor([[1., 1., 1., 1.],
        [1., 1., 1., 1.]])

这正是两个 API 的核心差异：distribute_tensor() 负责从 global tensor 分发数据；DTensor.from_local() 只是把已有 local tensors 标注成同一个 global DTensor 的不同分片。

Working with DTensor

拿到 DTensor 之后，用户代码通常仍然按普通 torch.Tensor 的方式写：

直接调用 PyTorch 算子，例如 +、relu、matmul、sum。
用 to_local() 查看当前 rank 的 local shard / replica。
用 full_tensor() 收集完整 global tensor。
用 redistribute() 显式改变 DTensor 的 layout。

DTensor 的关键价值在于：算子看到的是 global Tensor 语义，但实际执行时会尽量在 local tensor 上完成计算，并在必要时自动插入 collective communication。

Automatic Operator Parallelization

大多数 PyTorch 算子可以直接作用在 DTensor 上：

1
2
3
4
5
6


a = distribute_tensor(torch.randn(1024, 1024), mesh, [Shard(0)])
b = distribute_tensor(torch.randn(1024, 1024), mesh, [Shard(0)])

c = a + b
d = torch.relu(c)
e = c.sum(dim=0)

DTensor 会根据输入 placements 推导输出 placements：

对 elementwise 算子，如果输入都是相同的 Shard(0)，输出通常仍然是 Shard(0)，每个 rank 只算自己的 local shard。
对 reduction 算子，如果 reduction 维度正好是 sharded 维度，输出可能变成 Partial("sum")，表示每个 rank 只有部分归约结果。
当后续算子需要完整值时，DTensor 会通过 all-reduce、all-gather 等通信把 layout 转成需要的形式。

Local Tensor and Full Tensor

to_local() 返回当前 rank 实际持有的 local tensor：

1

local_tensor = dtensor.to_local()

如果 DTensor 是 Shard(0)，那么 local_tensor 是当前 rank 的分片；如果 DTensor 是 Replicate()，那么 local_tensor 是当前 rank 的完整副本。

full_tensor() 会把所有 shards 收集起来，返回完整的 logical tensor：

1

full_tensor = dtensor.full_tensor()

这一步通常会触发通信。例如 Shard(0) -> Replicate() 本质上需要 all-gather，所以它适合调试、保存、验证，不适合放在训练 hot path 里频繁调用。

Redistributing DTensor

redistribute() 用来显式改变 DTensor 的 placements：

1
2
3
4


sharded = distribute_tensor(x, mesh, [Shard(0)])

replicated = sharded.redistribute(mesh, [Replicate()])
resharded = replicated.redistribute(mesh, [Shard(1)])

常见 layout 转换和通信关系：

From	To	可能触发的通信
`Shard(dim)`	`Replicate()`	`all-gather`
`Replicate()`	`Shard(dim)`	local chunk / narrow
`Shard(src_dim)`	`Shard(dst_dim)`	`all-to-all` 或组合通信
`Partial()`	`Replicate()`	`all-reduce`
`Partial()`	`Shard(dim)`	`reduce-scatter`

从使用习惯上看，redistribute() 是 DTensor 里非常重要的显式同步点：当自动 layout propagation 推导出的 layout 不是后续计算想要的形式时，就需要手动指定目标 placements。

单脚本测试代码

下面的脚本可以直接保存为 dtensor_working.py 运行：

1

python3 dtensor_working.py

它会测试四件事：

DTensor elementwise 算子结果和普通 Tensor 一致。
sum(dim=0) 这类跨 shard 维度的 reduction 可以得到正确 global 结果。
to_local() 和 full_tensor() 分别对应 local view 和 global view。
redistribute() 可以在 Shard(0)、Replicate()、Shard(1) 之间转换，并保持 global tensor 不变。

完整脚本：

  1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113


import os
import socket

import torch
import torch.distributed as dist
import torch.multiprocessing as mp
from torch.distributed.device_mesh import init_device_mesh
from torch.distributed.tensor import Replicate, Shard, distribute_tensor


def find_free_port() -> str:
    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
        sock.bind(("127.0.0.1", 0))
        return str(sock.getsockname()[1])


def setup_dist(rank: int, world_size: int, backend: str, port: str) -> None:
    os.environ["MASTER_ADDR"] = "127.0.0.1"
    os.environ["MASTER_PORT"] = port
    os.environ["RANK"] = str(rank)
    os.environ["WORLD_SIZE"] = str(world_size)
    dist.init_process_group(backend=backend, rank=rank, world_size=world_size)


def describe(name: str, dtensor) -> str:
    local = dtensor.to_local()
    return (
        f"{name}: global_shape={tuple(dtensor.shape)}, "
        f"placements={dtensor.placements}, "
        f"local_shape={tuple(local.shape)}, "
        f"local_value=\n{local.cpu()}"
    )


def assert_same_global(name: str, dtensor, expected: torch.Tensor) -> None:
    full = dtensor.full_tensor()
    torch.testing.assert_close(full.cpu(), expected.cpu(), msg=f"{name} mismatch")


def run(rank: int, world_size: int, backend: str, device_type: str, port: str) -> None:
    setup_dist(rank, world_size, backend, port)

    if device_type == "cuda":
        torch.cuda.set_device(rank)
        device = torch.device("cuda", rank)
    else:
        device = torch.device("cpu")

    mesh = init_device_mesh(device_type, (world_size,), mesh_dim_names=("tp",))

    rows = world_size * 2
    cols = world_size * 2
    x = torch.arange(rows * cols, dtype=torch.float32, device=device).reshape(rows, cols)

    # 1. Start from a Shard(0) DTensor.
    sharded0 = distribute_tensor(x, mesh, [Shard(0)])
    assert_same_global("sharded0", sharded0, x)

    # 2. Elementwise ops stay sharded and can run locally.
    elementwise = torch.relu(sharded0 * 2 - 10)
    expected_elementwise = torch.relu(x * 2 - 10)
    assert_same_global("elementwise", elementwise, expected_elementwise)

    # 3. Reduction over the sharded dimension may create Partial internally.
    # Redistribute to Replicate() when we need the complete reduced value.
    reduced = sharded0.sum(dim=0)
    reduced_replicated = reduced.redistribute(mesh, [Replicate()])
    expected_reduced = x.sum(dim=0)
    assert_same_global("reduced", reduced_replicated, expected_reduced)

    # 4. Convert between local and global views.
    local_shard = sharded0.to_local()
    full_tensor = sharded0.full_tensor()
    torch.testing.assert_close(full_tensor.cpu(), x.cpu(), msg="full_tensor mismatch")

    # 5. Explicitly change placements.
    replicated = sharded0.redistribute(mesh, [Replicate()])
    sharded1 = replicated.redistribute(mesh, [Shard(1)])
    assert_same_global("replicated", replicated, x)
    assert_same_global("sharded1", sharded1, x)

    # Print rank by rank to keep output readable.
    for print_rank in range(world_size):
        dist.barrier()
        if rank == print_rank:
            print(f"\n===== rank {rank} =====")
            print(describe("sharded0", sharded0))
            print(describe("elementwise", elementwise))
            print(describe("reduced_replicated", reduced_replicated))
            print(describe("replicated", replicated))
            print(describe("sharded1", sharded1))
            print(f"to_local() shape: {tuple(local_shard.shape)}")
            print(f"full_tensor() shape: {tuple(full_tensor.shape)}")
        dist.barrier()

    if rank == 0:
        print("\nAll DTensor working tests passed.")

    dist.destroy_process_group()


if __name__ == "__main__":
    if torch.cuda.is_available():
        device_type = "cuda"
        backend = "nccl"
        world_size = min(torch.cuda.device_count(), 4)
    else:
        device_type = "cpu"
        backend = "gloo"
        world_size = 4

    port = find_free_port()
    mp.spawn(run, args=(world_size, backend, device_type, port), nprocs=world_size, join=True)

预期输出里可以重点观察 placements 的变化：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10


===== rank 0 =====
sharded0: global_shape=(8, 8), placements=(Shard(dim=0),), local_shape=(2, 8)
elementwise: global_shape=(8, 8), placements=(Shard(dim=0),), local_shape=(2, 8)
reduced_replicated: global_shape=(8,), placements=(Replicate(),), local_shape=(8,)
replicated: global_shape=(8, 8), placements=(Replicate(),), local_shape=(8, 8)
sharded1: global_shape=(8, 8), placements=(Shard(dim=1),), local_shape=(8, 2)
to_local() shape: (2, 8)
full_tensor() shape: (8, 8)

All DTensor working tests passed.

这个例子里，sharded0 和 elementwise 都保持 Shard(0)，说明 elementwise 计算可以直接在 local shard 上完成；reduced_replicated 变成 Replicate()，说明跨 shard 维度做 reduction 后需要归约到完整结果；sharded1 则展示了显式 reshard 的结果。

Automatic Parallelization in Action

前面几节介绍了 DTensor 的 layout 表达方式。本节关注另一个问题：当我们真的执行 PyTorch 算子时，DTensor 如何自动决定输出 placement，以及什么时候需要插入通信。

核心流程可以概括为：

用户调用普通 PyTorch 算子，例如 torch.sin(x)、torch.matmul(a, b)、x.sum()。
DTensor 拦截算子调用，读取输入 DTensor 的 DTensorSpec。
根据算子的 sharding rule 推导合法的输出 placement。
如果当前输入 placement 不满足某个合法策略，就先对输入做 redistribute()。
在 local tensor 上执行真实算子，并把结果重新包装成 DTensor。

Example: Elementwise Operations

Elementwise 算子通常是最简单的情况。如果输入是 Shard(0)，输出通常仍然是 Shard(0)，每个 rank 只处理自己的 local shard：

1
2
3
4
5
6
7


mesh = init_device_mesh("cuda", (4,))

x = distribute_tensor(torch.randn(8192, 8192), mesh, [Shard(0)])

y = torch.sin(x)
z = y + 1
w = x * 2

逻辑上，用户看到的是完整 Tensor：

1
2
3
4


x: global_shape = [8192, 8192], placement = Shard(0)
y: global_shape = [8192, 8192], placement = Shard(0)
z: global_shape = [8192, 8192], placement = Shard(0)
w: global_shape = [8192, 8192], placement = Shard(0)

实际执行时，每张 GPU 只计算自己的 rows：

1
2
3
4


GPU0: sin(x[0:2048, :])
GPU1: sin(x[2048:4096, :])
GPU2: sin(x[4096:6144, :])
GPU3: sin(x[6144:8192, :])

这类算子一般不需要通信，因为每个输出元素只依赖同位置的输入元素。

Example: Matrix Multiplication

矩阵乘法的 placement 决策取决于 shard 的维度是否落在输出维度或 contracting 维度上。

先看一个不需要通信的例子：

1
2
3
4
5
6


mesh = init_device_mesh("cuda", (4,))

A = distribute_tensor(torch.randn(1024, 512), mesh, [Shard(0)])
B = distribute_tensor(torch.randn(512, 256), mesh, [Replicate()])

C = torch.matmul(A, B)

A @ B 的数学语义是：

1
2
3


A: [M, K]
B: [K, N]
C: [M, N]

这里 A 沿 M 维度做 Shard(0)，而 B 在每个 rank 上都有完整副本。因此每个 rank 都能独立计算自己负责的输出 rows：

1
2
3
4
5
6
7
8


A: [1024, 512] with Shard(0)       B: [512, 256] with Replicate()

GPU0: A[0:256, :]     @ B  -> C[0:256, :]
GPU1: A[256:512, :]   @ B  -> C[256:512, :]
GPU2: A[512:768, :]   @ B  -> C[512:768, :]
GPU3: A[768:1024, :]  @ B  -> C[768:1024, :]

C: [1024, 256] with Shard(0)

输出 placement 可以自然保持 Shard(0)，不需要额外通信。

Example: Automatic Redistribution

再看一个需要自动重分布的例子：

1
2
3
4
5
6
7
8
9


mesh = init_device_mesh("cuda", (4,))

A = distribute_tensor(torch.randn(1024, 512), mesh, [Shard(0)])
B = distribute_tensor(torch.randn(512, 512), mesh, [Shard(0)])
C = distribute_tensor(torch.randn(512, 256), mesh, [Replicate()])

D = torch.matmul(A, B)
E = torch.matmul(D, C)
loss = E.sum()

初始状态：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17


A: [1024, 512] with Shard(0)
GPU0: [256, 512]
GPU1: [256, 512]
GPU2: [256, 512]
GPU3: [256, 512]

B: [512, 512] with Shard(0)
GPU0: [128, 512]
GPU1: [128, 512]
GPU2: [128, 512]
GPU3: [128, 512]

C: [512, 256] with Replicate()
GPU0: [512, 256]
GPU1: [512, 256]
GPU2: [512, 256]
GPU3: [512, 256]

Step 1: `D = torch.matmul(A, B)`

对 matmul(A, B) 来说：

A 的第 0 维是输出 rows，Shard(0) 是自然可保留的 output sharding。
B 的第 0 维是 contracting dimension K，但当前 B 也被 Shard(0) 切开。
每个 rank 只有一部分 K，无法直接算完整的 A_local @ B。

因此，DTensor 需要先把 B 从 Shard(0) 转成 Replicate()，本质上触发一次 all-gather：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13


Before:
B: Shard(0)
GPU0: [128, 512]
GPU1: [128, 512]
GPU2: [128, 512]
GPU3: [128, 512]

After all-gather:
B: Replicate()
GPU0: [512, 512]
GPU1: [512, 512]
GPU2: [512, 512]
GPU3: [512, 512]

之后每个 rank 可以独立计算自己的输出 rows：

1
2
3
4
5
6


D: [1024, 512] with Shard(0)

GPU0: D[0:256, :]     = A[0:256, :]     @ B
GPU1: D[256:512, :]   = A[256:512, :]   @ B
GPU2: D[512:768, :]   = A[512:768, :]   @ B
GPU3: D[768:1024, :]  = A[768:1024, :]  @ B

Step 2: `E = torch.matmul(D, C)`

此时：

D 是 Shard(0)，沿输出 rows 切分。
C 是 Replicate()，每个 rank 都有完整权重。

这又回到了理想情况，不需要通信：

1
2
3
4
5


D: Shard(0)          C: Replicate()       E: Shard(0)
GPU0: [256, 512]  @  [512, 256]       ->  GPU0: [256, 256]
GPU1: [256, 512]  @  [512, 256]       ->  GPU1: [256, 256]
GPU2: [256, 512]  @  [512, 256]       ->  GPU2: [256, 256]
GPU3: [256, 512]  @  [512, 256]       ->  GPU3: [256, 256]

输出 E 仍然是 Shard(0)。

Step 3: `loss = E.sum()`

E.sum() 是对整个 global Tensor 求和。由于 E 沿第 0 维切分，每个 rank 只能先算自己的局部和：

1
2
3
4


GPU0: sum(E[0:256, :])     = partial_0
GPU1: sum(E[256:512, :])   = partial_1
GPU2: sum(E[512:768, :])   = partial_2
GPU3: sum(E[768:1024, :])  = partial_3

为了得到 global scalar，需要一次 all-reduce(sum)：

1
2
3
4


partial_0 ┐
partial_1 ├── all-reduce(sum) ──> loss: Replicate() scalar
partial_2 │
partial_3 ┘

因此，reduction 算子是最容易触发 Partial -> Replicate 或 Partial -> Shard 通信的场景。

Placement Decision Rules

下面是一些常见算子的直觉规则：

Operation	Input placements	Output placement	说明
elementwise	inputs placements match	same as input	local shard 独立计算
`matmul(Shard(0), Replicate())`	rows sharded + full weight	`Shard(0)`	无通信
`matmul(Shard(0), Shard(0))`	rows sharded + K sharded	usually `Shard(0)`	需要先 gather 右矩阵
`matmul(Shard(1), Shard(0))`	K 维被两边切分	`Partial("sum")`	每个 rank 得到部分累加结果
`sum()` on `Shard(dim)`	reduced sharded dimension	`Partial("sum")` then `Replicate()`	需要 all-reduce 得到完整值
`sum(dim=d)` where `d != shard_dim`	non-sharded dimension	`Shard(adjusted_dim)`	通常不需要跨 rank 通信

这些规则不是硬编码在用户代码里的，而是由 DTensor 的 sharding propagation 机制根据每个 operator 的 schema 决定。

How DTensor Makes Decisions

DTensor 的自动并行不是全图优化，而是 eager execution 下的逐算子决策：

每个算子有一组合法的 sharding strategies，描述输入 placements 和输出 placements 的组合。
当当前输入 placements 已经匹配某个 strategy 时，算子可以直接执行。
当不匹配时，DTensor 会估算不同重分布方案的 communication cost，并选择一个局部成本较低的策略。
输入会先被 redistribute() 到目标 layout，再执行 local operator。

相关源码入口可以从这些模块理解：

torch/distributed/tensor/_sharding_prop.py：负责 sharding propagation，推导 output DTensorSpec。
torch/distributed/tensor/_op_schema.py：定义 OpSchema、OpStrategy 等策略表达。
torch/distributed/tensor/_ops/：不同 operator 的 sharding rule。
torch/distributed/tensor/_dispatch.py：DTensor 的 dispatch 流程，在需要时调用 redistribution。

这也带来一个限制：DTensor 当前更接近「per-operator greedy decision」，不是像 XLA GSPMD 那样对整张计算图做全局最优 sharding 规划。因此，局部通信最少的选择不一定是整个模型端到端通信最少的选择。实际训练中仍然需要用户通过合适的初始 placement、redistribute()、以及更高层的 TP/FSDP 策略来引导 layout。

Comparison with JAX GSPMD

JAX 的 GSPMD / XLA 通常会在 compile-time trace 完整计算图，然后为全图 sharding 求解更全局的方案。PyTorch DTensor 则更贴近 PyTorch eager 模型，运行时逐算子做 placement propagation 和 redistribution。

Aspect	PyTorch DTensor	JAX GSPMD / XLA
Optimization scope	Per-operator, local greedy decision	Whole-graph optimization
Decision time	Runtime eager execution	Compile-time tracing
Cost model	Lightweight redistribution cost	Global cost model / solver
User control	`placements` + `redistribute()`	`with_sharding_constraint` / partition spec
Strength	更贴近 PyTorch eager，调试和增量迁移更直接	全图视角更容易得到全局更优 sharding
Trade-off	可能出现局部最优但全局非最优的通信	编译成本和约束表达更重

总结来说，DTensor 的自动并行能让很多算子在不改用户代码的情况下自动分布式执行，但它不是完全替代并行策略设计。更好的理解方式是：DTensor 提供了统一的 layout 表达、算子级 propagation 和必要通信插入；用户仍然需要在模型并行、数据并行、参数布局这些更高层做结构化设计。

Guiding Placement Decisions with Explicit Redistribution

用户不能直接覆盖某个 PyTorch operator 的 DTensor sharding rule，也不能在调用 torch.matmul() 时手动指定“这个 op 的输出一定要是什么 placement”。但实际使用中有一个很重要的控制手段：在 operator 之前插入显式的 redistribute()。

这相当于给 DTensor 一个 placement hint：虽然你不能改 op schema 的选择逻辑，但你可以控制这个 operator 实际看到的输入 layout。由于 DTensor 是逐算子做 placement propagation，这个技巧常常能避免后续不必要的 reshard。

The Problem: Greedy Per-Operator Decisions

考虑下面这个计算图：

1
2
3
4
5
6
7
8
9


mesh = init_device_mesh("cuda", (4,))

A = distribute_tensor(torch.randn(1024, 512), mesh, [Shard(1)])
B = distribute_tensor(torch.randn(512, 256), mesh, [Shard(0)])
F = distribute_tensor(torch.randn(1024, 256), mesh, [Shard(1)])

C = torch.matmul(A, B)  # Shard(1) @ Shard(0) -> Partial("sum")
D = torch.relu(C)       # nonlinear op must resolve Partial first
result = D + F          # F is Shard(1), but D's placement was already chosen

对 matmul(A, B) 来说：

A 是 Shard(1)，沿 columns 切分，也就是沿 contracting dimension K 切分。
B 是 Shard(0)，沿 rows 切分，同样也是沿 contracting dimension K 切分。
每个 rank 可以计算一部分 K 对输出的贡献，因此输出 C 的自然 layout 是 Partial("sum")。

可以把这一步理解成：

1
2
3
4
5
6


A[:, K_i] @ B[K_i, :] = partial_C_i

GPU0: partial_C_0 ┐
GPU1: partial_C_1 ├── needs sum over K partitions
GPU2: partial_C_2 │
GPU3: partial_C_3 ┘

接下来执行 torch.relu(C) 时，问题出现了：非线性算子不能直接作用在 Partial 上，因为：

1

relu(partial_0) + relu(partial_1) != relu(partial_0 + partial_1)

所以 relu() 之前必须把 C 从 Partial("sum") resolve 成一个 concrete layout。由于 DTensor 在 eager 模式下只看到当前这个 relu(C) operator，它并不知道后面马上会有一个 D + F，也不知道 F 已经是 Shard(1)。

如果 DTensor 在 relu() 前把 C resolve 成 Replicate()，那么通信路径可能变成：

1
2
3
4
5
6
7
8
9


C: Partial("sum")
  -- all-reduce -->
D: Replicate()

D + F:
  D: Replicate()
  F: Shard(1)
  -- another redistribution may be needed -->
result

这就是 per-operator greedy decision 的局限：单个 operator 看起来合理的 layout，不一定是整个后续计算图通信最少的 layout。

The Solution: Use `redistribute()` as a Placement Hint

用户通常知道更长的计算上下文。因此，可以在 relu() 前显式把 C resolve 成后续更需要的 layout：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13


mesh = init_device_mesh("cuda", (4,))

A = distribute_tensor(torch.randn(1024, 512), mesh, [Shard(1)])
B = distribute_tensor(torch.randn(512, 256), mesh, [Shard(0)])
F = distribute_tensor(torch.randn(1024, 256), mesh, [Shard(1)])

C = torch.matmul(A, B)  # Shard(1) @ Shard(0) -> Partial("sum")

# User knows F is Shard(1) downstream, so resolve C to Shard(1) now.
C = C.redistribute(mesh, [Shard(1)])  # Partial("sum") -> Shard(1)

D = torch.relu(C)       # Shard(1) -> Shard(1), no communication
result = D + F          # Shard(1) + Shard(1) -> Shard(1), no reshard

这里 C.redistribute(mesh, [Shard(1)]) 的语义是：把每个 rank 上的 partial output 沿第 1 维做 reduce-scatter，直接得到按 columns 切分的完整结果：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15


Before:
C: Partial("sum"), each rank has logical shape [1024, 256]

GPU0: partial_C_0 [1024, 256] ┐
GPU1: partial_C_1 [1024, 256] ├── reduce-scatter(sum, dim=1)
GPU2: partial_C_2 [1024, 256] │
GPU3: partial_C_3 [1024, 256] ┘

After:
C: Shard(1), global shape [1024, 256]

GPU0: C[:,   0: 64]
GPU1: C[:,  64:128]
GPU2: C[:, 128:192]
GPU3: C[:, 192:256]

这样后续两步都可以保持 Shard(1)：

torch.relu(Shard(1)) -> Shard(1)：elementwise 算子在 local shard 上执行，不需要通信。
Shard(1) + Shard(1) -> Shard(1)：两个输入 layout 匹配，也不需要额外 reshard。

Communication Trade-Off

显式 redistribute() 不是“免费优化”，它只是让用户选择通信发生在哪里、以及通信后的 layout 长什么样。

在上面的例子里，有两种可能路径：

Strategy	Communication	Downstream layout
Let DTensor resolve `Partial -> Replicate` before `relu`	`all-reduce` over full `[1024, 256]`	后续和 `F: Shard(1)` 计算时可能还要 reshard
User explicitly resolves `Partial -> Shard(1)`	`reduce-scatter` over `[1024, 256]`	后续 `relu` 和 `+ F` 都保持 `Shard(1)`

选择哪一种取决于后续计算和 tensor size。如果后面大量算子都希望 Replicate()，那么 all-reduce 可能是合理的；如果后面会和 Shard(1) 的张量继续计算，那么提前 reduce-scatter 到 Shard(1) 往往更好。

因此，redistribute() 的价值不只是“改变 layout”，而是把用户对后续计算图的全局理解显式写进 eager 程序里。

Comparison with JAX `with_sharding_constraint()`

JAX 提供了更 declarative 的方式：jax.lax.with_sharding_constraint() 可以告诉 XLA 某个中间结果期望采用什么 sharding。编译器在看到完整图之后，可以把这个 constraint 纳入全局优化。

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25


import jax
import jax.numpy as jnp
from jax.sharding import Mesh, NamedSharding, PartitionSpec as P

mesh = Mesh(jax.devices(), ("x",))


@jax.jit
def computation_with_hint(A, B, F):
    C = jnp.matmul(A, B)

    # User tells the compiler: keep C compatible with downstream F,
    # which is sharded on tensor dim 1.
    C = jax.lax.with_sharding_constraint(C, NamedSharding(mesh, P(None, "x")))

    D = jax.nn.relu(C)
    return D + F


# A is sharded on columns, B on rows, F on columns.
A = jax.device_put(A_data, NamedSharding(mesh, P(None, "x")))
B = jax.device_put(B_data, NamedSharding(mesh, P("x", None)))
F = jax.device_put(F_data, NamedSharding(mesh, P(None, "x")))

result = computation_with_hint(A, B, F)

两者的差异可以总结为：

Aspect	PyTorch DTensor `redistribute()`	JAX `with_sharding_constraint()`
Style	Imperative layout transformation	Declarative compiler constraint
Timing	立即执行，立即通信	编译期参与全图优化
Scope	影响下一个及后续 eager operators 看到的 input placement	影响 XLA 对整张图的 sharding 规划
User control	直接、显式、容易调试	更抽象，但可能获得更全局的优化

所以，在 PyTorch DTensor 里，“引导 placement 决策”的核心方式就是：在关键边界上显式插入 redistribute()，把中间结果转成后续计算最自然的 layout。

实现原理 `__torch_dispatch__`

PyTorch 算子下发流程：

OneFlow: Redesign the Distributed Deep Learning Framework from Scratch, https://arxiv.org/pdf/2110.15032 ↩︎
GSPMD: General and Scalable Parallelization for ML Computation Graphs, https://arxiv.org/pdf/2105.04663 ↩︎

Quick Start

从 Tensor 到 DTensor

Tensor

DTensor

DeviceMesh - Describing Device Topology

Placement - Describing Tensor Distribution

Shard(dim)

Replicate()

Partial(reduce_op)

Multi-Dimensional Placements

DTensorSpec

DTensor

Creating DTensor

Method 1: distribute_tensor()

Method 2: DTensor.from_local()

单脚本可运行示例

Working with DTensor

Automatic Operator Parallelization

Local Tensor and Full Tensor

Redistributing DTensor

单脚本测试代码

Automatic Parallelization in Action

Example: Elementwise Operations

Example: Matrix Multiplication

Example: Automatic Redistribution

Step 1: D = torch.matmul(A, B)

Step 2: E = torch.matmul(D, C)

Step 3: loss = E.sum()

Placement Decision Rules

How DTensor Makes Decisions

Comparison with JAX GSPMD

Guiding Placement Decisions with Explicit Redistribution

The Problem: Greedy Per-Operator Decisions

The Solution: Use redistribute() as a Placement Hint

Communication Trade-Off

Comparison with JAX with_sharding_constraint()

实现原理 __torch_dispatch__

Method 1: `distribute_tensor()`

Method 2: `DTensor.from_local()`

Step 1: `D = torch.matmul(A, B)`

Step 2: `E = torch.matmul(D, C)`

Step 3: `loss = E.sum()`

The Solution: Use `redistribute()` as a Placement Hint

Comparison with JAX `with_sharding_constraint()`

实现原理 `__torch_dispatch__`