code_hyper_parallel_dtensor · Pull Request !121 · MindSpore/hyper-parallel

# DTensor API 文档

## 概述

DTensor（Distributed Tensor，分布式张量）是 Hyper-Parallel 框架中的核心组件，用于表示分布在多个设备上的张量。它提供了一种统一的抽象，使用户能够以声明式的方式定义张量的分布策略，而无需手动管理底层的通信和数据切分。

### 核心概念

| 概念           | 说明                                                         |
| -------------- | ------------------------------------------------------------ |
| **DeviceMesh** | 设备网格，将物理设备组织成多维逻辑网格，支持数据并行（DP）、张量并行（TP）、流水线并行（PP）等多种并行策略 |
| **Placement**  | 放置策略，定义张量在网格各维度上的分布方式，包括 `Shard`（切分）、`Replicate`（复制）、`Partial`（部分归约）等 |
| **Layout**     | 内部布局对象，用于管理张量的分布映射和重分布操作             |

### 设计目标

1. **简化分布式编程**：用户只需声明张量的分布策略，框架自动处理数据切分和通信
2. **支持混合并行**：通过多维设备网格，灵活组合不同的并行策略
3. **高效重分布**：提供优化的张量重分布机制，最小化通信开销

---

## 类定义

### DTensor

```python
class DTensor(DTensorBase):
    """
    DTensor - 分布式张量

DTensor 表示一个分布在多个设备上的张量，根据 DeviceMesh 和放置规范
    （Placement）进行分布。

Attributes:
        _local_tensor (Tensor): 当前设备上的本地张量分片
        _device_mesh (DeviceMesh): 描述设备拓扑的设备网格
        _placements (Tuple[Placement, ...]): 每个网格维度的放置策略
        _layout (Layout): 用于重分布操作的内部布局对象
    """
```

---

## 构造方法

### `__init_data__`

初始化 DTensor 的内部数据。

```python
def __init_data__(
    self,
    local_tensor: Tensor,
    device_mesh: DeviceMesh,
    placements: Sequence[Placement]
) -> None
```

**参数说明：**

| 参数           | 类型                  | 说明                                     |
| -------------- | --------------------- | ---------------------------------------- |
| `local_tensor` | `Tensor`              | 当前设备上的本地张量分片                 |
| `device_mesh`  | `DeviceMesh`          | 描述设备拓扑的设备网格                   |
| `placements`   | `Sequence[Placement]` | 放置策略序列，长度应与设备网格维度数一致 |

**内部行为：**

1. 存储本地张量、设备网格和放置策略
2. 调用 `_build_layout()` 构建内部布局对象，用于后续的重分布操作

---

## 静态方法

### `from_local`

从本地张量创建 DTensor。

```python
@staticmethod
def from_local(
    local_tensor: Tensor,
    device_mesh: DeviceMesh,
    placements: Sequence[Placement]
) -> 'DTensor'
```

**参数说明：**

| 参数           | 类型                  | 说明                                                         |
| -------------- | --------------------- | ------------------------------------------------------------ |
| `local_tensor` | `Tensor`              | 当前设备上的本地张量分片                                     |
| `device_mesh`  | `DeviceMesh`          | 描述设备拓扑的设备网格                                       |
| `placements`   | `Sequence[Placement]` | 放置策略序列，每个元素应为 `Shard`、`Replicate`、`Partial` 等 |

**返回值：** `DTensor` - 新创建的分布式张量实例

**使用场景：**

- 当每个 rank 已经拥有正确的本地分片时使用
- 通常用于从已存在的分布式数据创建 DTensor

**示例：**

```python
from hyper_parallel.core.placement_types import Shard, Replicate

# 创建 2x2 的设备网格，维度名称为 "dp" 和 "tp"
mesh = DeviceMesh((2, 2), ("dp", "tp"))

# 假设每个 rank 已有本地张量（shape 可能不同）
local_tensor = Tensor(np.ones((4, 4)))

# 创建 DTensor：沿第 0 维在 "dp" 维度切分，在 "tp" 维度复制
dtensor = DTensor.from_local(local_tensor, mesh, [Shard(0), Replicate()])
```

---

### `distribute_tensor`

将全局张量分发到设备网格。

```python
@staticmethod
def distribute_tensor(
    tensor: Tensor,
    device_mesh: DeviceMesh,
    placements: Sequence[Placement]
) -> 'DTensor'
```

**参数说明：**

| 参数          | 类型                  | 说明                                         |
| ------------- | --------------------- | -------------------------------------------- |
| `tensor`      | `Tensor`              | 要分发的全局张量（所有 rank 应持有相同数据） |
| `device_mesh` | `DeviceMesh`          | 描述设备拓扑的设备网格                       |
| `placements`  | `Sequence[Placement]` | 分发的放置策略                               |

**返回值：** `DTensor` - 包含本地分片的新 DTensor 实例

**使用场景：**

- 当所有 rank 拥有相同的全局张量，需要按策略切分时使用
- 通常用于模型初始化或数据加载阶段

**与 `from_local` 的区别：**

| 方法                | 输入要求                     | 行为                     |
| ------------------- | ---------------------------- | ------------------------ |
| `from_local`        | 每个 rank 已有正确的本地分片 | 直接封装为 DTensor       |
| `distribute_tensor` | 所有 rank 持有相同的全局张量 | 根据布局切片获取本地分片 |

**示例：**

```python
from hyper_parallel.core.placement_types import Shard, Replicate

mesh = DeviceMesh((2, 2), ("dp", "tp"))

# 所有 rank 持有相同的全局张量
global_tensor = Tensor(np.arange(16).reshape(4, 4))

# 分发：沿第 0 维在 "dp" 维度切分
dtensor = DTensor.distribute_tensor(global_tensor, mesh, [Shard(0), Replicate()])

# 结果：
# rank 0, 1 得到: [[0,1,2,3], [4,5,6,7]]
# rank 2, 3 得到: [[8,9,10,11], [12,13,14,15]]
```

---

## 实例方法

### `to_local`

获取本地张量分片。

```python
def to_local(self) -> Tensor
```

**返回值：** `Tensor` - 当前设备上的本地张量分片

**说明：**

- 直接返回存储的本地张量，不涉及任何通信
- 返回的张量形状取决于分布策略

**示例：**

```python
local_tensor = dtensor.to_local()
print(local_tensor.shape)  # 本地分片的形状
```

---

### `redistribute`

将 DTensor 重分布到新的设备网格和放置策略。

```python
def redistribute(
    self,
    device_mesh: DeviceMesh,
    placements: Sequence[Placement]
) -> 'DTensor'
```

**参数说明：**

| 参数          | 类型                  | 说明         |
| ------------- | --------------------- | ------------ |
| `device_mesh` | `DeviceMesh`          | 目标设备网格 |
| `placements`  | `Sequence[Placement]` | 目标放置策略 |

**返回值：** `DTensor` - 具有指定分布的新 DTensor

**内部行为：**

1. 根据目标设备网格和放置策略构建目标布局（`dst_layout`）
2. 调用 `_tensor_redistribution.redistribution()` 执行重分布
3. 自动处理必要的集合通信（AllGather、AllReduce、AllToAll 等）

**使用场景：**

- 在不同并行策略之间转换
- 前向传播和反向传播中调整张量分布
- 实现算子级别的并行策略

**示例：**

```python
from hyper_parallel.core.placement_types import Shard, Replicate

# 原始分布：沿第 0 维切分
dtensor = DTensor.from_local(local_tensor, mesh, [Shard(0), Replicate()])

# 重分布：改为沿第 1 维切分
new_dtensor = dtensor.redistribute(mesh, [Replicate(), Shard(1)])
```

---

### `reduce_partial`

归约 Partial 状态的 DTensor。

```python
def reduce_partial(self) -> 'DTensor'
```

**返回值：** `DTensor` - 归约后的新 DTensor

**说明：**

- 当 DTensor 处于 `Partial` 状态时（如矩阵乘法后的部分和），此方法执行归约操作
- 调用 `_tensor_redistribution.reduce_partial()` 完成 AllReduce 通信

**使用场景：**

- 矩阵乘法后的部分结果归约
- 梯度聚合

---

### `full_tensor`

收集所有分片，返回完整的全局张量。

```python
def full_tensor(self) -> Tensor
```

**返回值：** `Tensor` - 包含完整数据的张量

**内部行为：**

1. 创建全复制（Replicate）的目标布局
2. 将所有放置策略设为 `Replicate()`
3. 清除 Partial 状态
4. 执行重分布操作
5. 返回本地张量（此时已包含完整数据）

**性能警告：**

此操作涉及跨所有 rank 的通信，对于大型张量可能开销较大。请在性能关键路径中谨慎使用。

**使用场景：**

- 调试和验证
- 保存检查点
- 需要完整张量进行计算的场景

**示例：**

```python
# 获取本地分片（仅当前 rank 的数据）
local_tensor = dtensor.to_local()

# 获取完整张量（所有 rank 的数据合并）
full_tensor = dtensor.full_tensor()

print(f"Local shape: {local_tensor.shape}")
print(f"Full shape: {full_tensor.shape}")
```

---

## 属性

### `device_mesh`

```python
@property
def device_mesh(self) -> DeviceMesh
```

**返回值：** `DeviceMesh` - DTensor 关联的设备网格

---

### `placements`

```python
@property
def placements(self) -> Tuple[Placement, ...]
```

**返回值：** `Tuple[Placement, ...]` - DTensor 的放置策略元组

---

### `layout`

```python
@property
def layout(self) -> Layout
```

**返回值：** `Layout` - 用于重分布的内部布局对象；若未初始化则返回 `None`

---

### `shape`

```python
@property
def shape(self) -> Tuple[int, ...]
```

**返回值：** `Tuple[int, ...]` - 全局张量形状

**说明：**

通过布局对象根据本地形状计算全局形状。例如，若本地形状为 `(4, 8)`，在第 0 维切分为 2 份，则全局形状为 `(8, 8)`。

---

### `local_shape`

```python
@property
def local_shape(self) -> Tuple[int, ...]
```

**返回值：** `Tuple[int, ...]` - 本地张量形状

## 对比总结

| API                 | 输入       | 输出       | 通信              | 场景             |
| ------------------- | ---------- | ---------- | ----------------- | ---------------- |
| `from_local`        | 本地tensor | DTensor    | 无/可选broadcast  | 每个rank已有数据 |
| `to_local`          | DTensor    | 本地tensor | 无                | 获取本地数据     |
| `redistribute`      | DTensor    | DTensor    | 视情况            | 改变分布策略     |
| `full_tensor`       | DTensor    | 完整tensor | all_gather        | 收集完整数据     |
| `distribute_tensor` | 完整tensor | DTensor    | scatter/broadcast | 初始化时分发数据 |

## 数据流示意

![输入图片说明](https://foruda.gitee.com/images/1768874722433490033/3487738c_13767016.png "dtensor.png")

---

## 辅助类

### `SkipDTensorDispatch`

上下文管理器，用于临时禁用 DTensor 的自动分发机制。

```python
class SkipDTensorDispatch:
    def __enter__(self):
        """进入上下文时禁用 DTensor 分发"""
        from hyper_parallel.core.shard._op_dispatch import disable_dtensor_dispatch
        disable_dtensor_dispatch()

def __exit__(self, exc_type, exc_val, exc_tb):
        """退出上下文时重新启用 DTensor 分发"""
        from hyper_parallel.core.shard._op_dispatch import enable_dtensor_dispatch
        enable_dtensor_dispatch()
```

**使用场景：**

- 需要直接操作底层 Tensor 而不触发 DTensor 分发逻辑时
- 调试和性能分析
- 实现自定义的分布式算子

**示例：**

```python
with SkipDTensorDispatch():
    # 在此上下文中，张量操作不会触发 DTensor 的自动分发
    result = some_local_operation(tensor)
```

---

## 辅助函数

### `_build_layout`

从设备网格和放置策略构建布局对象。

```python
def _build_layout(
    device_mesh: DeviceMesh,
    placements: Sequence[Placement],
    tensor_dim: int
) -> Layout
```

**参数说明：**

| 参数          | 类型                  | 说明         |
| ------------- | --------------------- | ------------ |
| `device_mesh` | `DeviceMesh`          | 设备网格     |
| `placements`  | `Sequence[Placement]` | 放置策略序列 |
| `tensor_dim`  | `int`                 | 张量的维度数 |

**返回值：** `Layout` - 构建的布局对象

**内部行为：**

1. 从设备网格创建基础布局
2. 应用放置策略
3. 构建放置到张量维度的映射

---

## 放置策略类型

DTensor 支持以下放置策略：

| 策略                 | 说明                             | 示例                         |
| -------------------- | -------------------------------- | ---------------------------- |
| `Shard(dim)`         | 沿指定维度切分张量               | `Shard(0)` 表示沿第 0 维切分 |
| `Replicate()`        | 在该网格维度上复制完整张量       | 每个设备持有完整副本         |
| `Partial(reduce_op)` | 部分归约状态，需要后续 AllReduce | 通常在矩阵乘法后出现         |

---

## 完整示例

### 示例 1：数据并行训练

```python
from hyper_parallel.core.dtensor import DTensor
from hyper_parallel.core.device_mesh import DeviceMesh
from hyper_parallel.core.placement_types import Shard, Replicate

# 创建数据并行的设备网格（4 个设备）
mesh = DeviceMesh((4,), ("dp",))

# 模型权重：在所有设备上复制
weights = Tensor(np.random.randn(1024, 1024))
dist_weights = DTensor.distribute_tensor(weights, mesh, [Replicate()])

# 输入数据：沿 batch 维度切分
batch_data = Tensor(np.random.randn(128, 1024))  # batch_size=128
dist_data = DTensor.distribute_tensor(batch_data, mesh, [Shard(0)])
# 每个设备得到 32 个样本
```

### 示例 2：张量并行（Tensor Parallelism）

```python
# 创建张量并行的设备网格
mesh = DeviceMesh((4,), ("tp",))

# 权重矩阵沿列切分
weights = Tensor(np.random.randn(1024, 4096))
dist_weights = DTensor.distribute_tensor(weights, mesh, [Shard(1)])
# 每个设备持有 (1024, 1024) 的分片
```

### 示例 3：混合并行（2D 并行）

```python
# 创建 2x2 的设备网格：数据并行 + 张量并行
mesh = DeviceMesh((2, 2), ("dp", "tp"))

# 权重：在 dp 维度复制，在 tp 维度切分
weights = Tensor(np.random.randn(1024, 2048))
dist_weights = DTensor.distribute_tensor(weights, mesh, [Replicate(), Shard(1)])

# 输入：在 dp 维度切分，在 tp 维度复制
inputs = Tensor(np.random.randn(64, 1024))
dist_inputs = DTensor.distribute_tensor(inputs, mesh, [Shard(0), Replicate()])
```

### 示例 4：重分布操作

```python
# 矩阵乘法后需要改变分布策略
# 假设 result 是 [Shard(0), Partial()] 状态

# 先归约 Partial 状态
result = result.reduce_partial()

# 再重分布到所需的布局
result = result.redistribute(mesh, [Replicate(), Shard(0)])
```

---

## 依赖模块

| 模块                                        | 说明                                             |
| ------------------------------------------- | ------------------------------------------------ |
| `hyper_parallel.core.layout`                | 提供 Layout、DeviceMesh 和张量切片功能           |
| `hyper_parallel.core.placement_types`       | 定义 Placement、Shard、Replicate、Partial 等类型 |
| `hyper_parallel.core.tensor_redistribution` | 实现张量重分布的底层逻辑                         |
| `hyper_parallel.core.shard._op_dispatch`    | DTensor 的算子分发机制                           |
| `hyper_parallel.platform`                   | 平台适配层，提供 DTensorBase 和 Tensor 基类      |

---

## 注意事项

1. **内存管理**：`full_tensor()` 会在每个 rank 上创建完整张量的副本，注意内存使用
2. **通信开销**：`redistribute()` 和 `full_tensor()` 涉及集合通信，应尽量减少调用频率
3. **放置策略一致性**：确保 `placements` 的长度与 `device_mesh` 的维度数一致

MindSpore / hyper-parallel关闭 .gitee-modal { width: 500px !important; }

搜索帮助

MindSpore / hyper-parallel
关闭