λLLM Wiki
3 min read#pytorch#distributed-training#dtensor#fsdp#tensor-parallelism

DTensor

M
memset0
Apr 9, 2026

DTensor(Distributed Tensor) 是 PyTorch 的分布式张量抽象。它将一个逻辑上完整的张量分布到多个设备上，并封装了分布式元数据，使得上层代码可以像操作普通张量一样操作分布式张量。DTensor 是 FSDP2 和 PyTorch 分布式张量并行的底层基础。

核心结构

一个 DTensor 包装了本地张量（local tensor）并附带三项分布式元数据：

DTensor
├── _local_tensor : torch.Tensor    — 当前设备上实际存储的数据分片
├── device_mesh   : DeviceMesh      — 参与分布的设备拓扑
└── placements    : tuple[Placement] — 每个 mesh 维度上的分布方式

DeviceMesh：设备拓扑

DeviceMesh(设备网格) 描述了参与分布式计算的设备集合及其拓扑关系。

from torch.distributed.device_mesh import init_device_mesh

# 1D mesh：4 个 GPU 组成一行（用于纯 FSDP 或纯 TP）
mesh_1d = init_device_mesh("cuda", (4,), mesh_dim_names=("dp",))

# 2D mesh：8 个 GPU 排成 4×2（外层 FSDP，内层 TP）
#   DP group: {GPU0,GPU1,GPU2,GPU3}, {GPU4,GPU5,GPU6,GPU7}  — 不对，是按列/行分
mesh_2d = init_device_mesh("cuda", (4, 2), mesh_dim_names=("dp", "tp"))
# dp 维度 size=4: 4 个数据并行组
# tp 维度 size=2: 每组内 2 个 GPU 做张量并行

2D mesh 是 FSDP + TP 组合并行的基础：一个参数可以在 dp 维度上用 Shard 做 FSDP 分片，同时在 tp 维度上用 Shard 做张量并行切分。

Placement：分布方式

Placement 描述了张量在 mesh 的某个维度上如何分布。三种基本类型：

Placement含义用途
Shard(dim)张量沿 dim 维切分，每个设备持有 1/N1/NFSDP 参数分片、TP 权重切分
Replicate()张量在每个设备上完整复制DDP 参数复制、TP 中不切分的参数
Partial(reduce_op)每个设备持有部分结果，需要归约才能得到完整值反向传播中的梯度（需 reduce-scatter）

示例——一个 [512, 512] 的权重在 4 个 GPU 上：

import torch
from torch.distributed.tensor import DTensor, Shard, Replicate

# Shard(0)：按行切分 → 每个 GPU 持有 [128, 512]
# 这就是 FSDP2 对参数的分片方式
dtensor_shard0 = DTensor.from_local(local_tensor, mesh, placements=(Shard(0),))

# Shard(1)：按列切分 → 每个 GPU 持有 [512, 128]
# 这是 TP 列并行的权重切分方式
dtensor_shard1 = DTensor.from_local(local_tensor, mesh, placements=(Shard(1),))

# Replicate()：每个 GPU 持有完整的 [512, 512]
# 这是 DDP 的参数存储方式
dtensor_rep = DTensor.from_local(local_tensor, mesh, placements=(Replicate(),))

在 FSDP2 中的角色

FSDP2 用 DTensor 替换了 FSDP1 的 FlatParameter。调用 fully_shard(module) 后，模块的每个参数被就地转换为 DTensor(Shard(0))：

from torch.distributed.fsdp import fully_shard

model = Transformer()
for layer in model.layers:
    fully_shard(layer)
fully_shard(model)

# 包装后，每个参数变成 DTensor
param = model.layers[0].self_attn.q_proj.weight
print(type(param))        # → DTensor
print(param.placements)   # → (Shard(0),)
print(param.shape)        # → torch.Size([512, 512])  ← 逻辑形状（完整）
print(param._local_tensor.shape)  # → torch.Size([128, 512])  ← 实际存储（1/4）

Forward/Backward 的状态转换

FSDP2 的 hook 在前向/反向过程中动态转换参数的分布状态：

平时(idle):     DTensor(Shard(0))    — 只存本地分片，省显存
     ↓ pre-forward hook: All-Gather
计算中(compute): plain torch.Tensor   — 完整参数，正常计算
     ↓ post-forward hook: 释放
计算后(idle):    DTensor(Shard(0))    — 回到分片状态

反向传播时相同流程再走一遍，加上反向后对梯度做 Reduce-Scatter。

对比 FSDP1 FlatParameter

维度FlatParameter (FSDP1)DTensor (FSDP2)
表示所有参数拼接成一个 1D 张量每个参数独立
元数据只有一份（整个 flat tensor 共享 requires_grad、dtype）每个参数独立持有
部分冻结❌ 报错：FlatParameter requires uniform requires_grad✅ 每个 DTensor 独立设置 requires_grad
参数名丢失（变成 _fsdp_wrapped_module.flat_param）保留原始名称
Checkpoint需要 All-Gather 通信还原完整参数直接保存分片，无通信
与 TP 组合困难原生支持（2D mesh + 不同维度的 Shard）

FSDP1 FlatParameter 在 LoRA 等部分冻结场景下的具体问题和 workaround，详见 raw/note/fsdp-partial-parameter-freezing.md。

在张量并行中的角色

DTensor 也是 PyTorch 实现 张量并行(TP) 的基础。TP 中，权重矩阵按行或列切分，每个 GPU 只计算矩阵乘法的一部分：

# 列并行 Linear：权重按列切分
# W: [512, 2048] → Shard(1) → 每个 GPU 持有 [512, 512]（4 GPU）
# 输入 x 是 Replicate()，每个 GPU 计算 x @ W_local → 输出是 Shard(1)

# 行并行 Linear：权重按行切分
# W: [2048, 512] → Shard(0) → 每个 GPU 持有 [512, 512]
# 输入按列切分（Shard(1)），每个 GPU 计算 x_local @ W_local → 输出是 Partial(SUM)
# 需要 AllReduce 得到完整输出

2D 并行：FSDP + TP 的组合

通过 2D DeviceMesh，DTensor 可以同时表示 FSDP 和 TP 的分布：

# 8 GPU: 4-way FSDP × 2-way TP
mesh = init_device_mesh("cuda", (4, 2), mesh_dim_names=("dp", "tp"))

# 一个参数的分布状态：
# placements = (Shard(0), Shard(0))
#               ↑ dp维度   ↑ tp维度
# → 先在 dp 维度按行切分（FSDP），再在 tp 维度按行切分（TP 列并行）

这种组合是训练超大模型的标准做法：节点内用 TP（利用 NVLink 高带宽），节点间用 FSDP（利用集群总显存）。

相关页面

Comments