pypto.distributed 模块介绍

【免费下载链接】pyptoPyPTO（发音: pai p-t-o）：Parallel Tensor/Tile Operation编程范式。项目地址: https://gitcode.com/cann/pypto

1. 概述

pypto.distributed模块提供了分布式场景下的共享内存通信能力，支持多个 PE（Processing Element）之间的数据传输、同步和协同计算。该模块基于对称内存概念设计，实现了高效的跨卡数据交换机制。

2. 核心概念：ShmemTensor

2.1 ShmemTensor 的设计理念

ShmemTensor（Shared Memory Tensor）是分布式通信的核心数据结构，与普通 Tensor 有以下关键区别：

• 对称内存访问：通过指定访问的 PE，可以访问其他卡的 ShmemTensor，实现了跨 PE 的数据共享
• 视图操作支持：与普通 Tensor 一样，ShmemTensor 支持 View 操作，允许对共享内存张量的部分视图进行操作
• 通信组隔离：通过group_name参数实现不同通信域的隔离，支持多个独立的分布式任务

2.2 与普通 Tensor 的对比

3. 分布式通信设计模式

3.1 通信模型

pypto.distributed采用基于共享内存的通信模型，主要包含以下操作类型：

1. 数据传输：通过shmem_put和shmem_get实现 PE 间的数据读写
2. 信号同步：通过shmem_signal和shmem_wait_until实现 PE 间的同步通知
3. 视图操作：通过shmem_view创建共享内存的部分视图
4. 集合通信：通过shmem_barrier_all实现全局同步

3.2 同步机制

信号同步是分布式通信的关键机制，确保数据的一致性和正确性：

• 信号发送：使用shmem_signal向目标 PE 发送信号通知
• 信号等待：使用shmem_wait_until等待信号满足指定条件
• 原子操作：支持 SET（覆盖）和 ADD（累加）两种原子操作类型
• 广播支持：支持向所有 PE 广播信号

3.3 典型通信流程

一个典型的分布式通信流程包含以下步骤：

1. 创建 ShmemTensor（数据张量和信号张量） 2. 设置 TileShape（必须步骤） 3. 数据写入（shmem_put） 4. 信号发送（shmem_signal） 5. 等待信号（shmem_wait_until） 6. 数据读取（shmem_get） 7. 可选：全局同步（shmem_barrier_all）

4. 应用场景

4.1 AllReduce/AllGather 等集合通信

在分布式训练和推理中，AllReduce 和 AllGather 是常见的集合通信模式：

• AllReduce：所有 PE 的数据聚合后分发到所有 PE
• AllGather：收集所有 PE 的数据并分发到所有 PE
• ReduceScatter：聚合数据后分片分发到不同 PE

4.2 MoE（Mixture of Experts）分布式推理

MoE 模型需要动态路由到不同的专家，分布式通信模块支持：

• 专家路由和分发
• 专家计算结果的聚合
• 多专家协同计算

4.3 自定义分布式算法

用户可以基于pypto.distributed实现自定义的分布式算法：

• 自定义通信模式
• 特定应用的数据交换
• 多阶段协同计算

5. 使用最佳实践

5.1 TileShape 设置

重要性：TileShape 的正确设置是分布式通信正常工作的前提

• 必须设置：在调用任何 ShmemTensor 相关函数前，必须通过set_vec_tile_shapes设置 TileShape
• 维度匹配：TileShape 的维度应与数据张量的维度一致
• 一致性要求：shmem_signal和shmem_wait_until的 TileShape 设置必须保持一致

# 正确示例 pypto.set_vec_tile_shapes(16, 64) # 对应 [m, n] 形状的数据

5.2 依赖关系管理

重要性：正确管理操作依赖关系确保数据一致性和执行顺序

• 依赖传递：通过pred参数正确传递操作依赖关系
• 操作顺序：shmem_get通常在shmem_wait_until之后执行，确保数据已写入
• 流水优化：在切块数据大于 1 的场景下，保持相同的切块配置以优化流水排布

5.3 性能优化建议

• 批量传输：尽量使用较大的数据块进行传输，减少通信次数
• 视图复用：合理使用shmem_view避免重复创建共享内存张量
• 信号合并：在可能的情况下，合并多个信号操作
• 流水并行：合理设置 TileShape 以优化流水排布

6. 使用示例

本节提供的快速入门示例仅用于演示pypto.distributed模块的基本用法。更详细的使用示例和完整的实现案例，请参考以下文档：

• combine_shmem_implementation.md - 使用通信shmem API实现combine算子

6.1 基础数据传输

import pypto # 设置 TileShape（必须步骤） pypto.set_vec_tile_shapes(16, 64) # 创建本地数据 local_data = pypto.tensor([16, 64], pypto.DT_FP32, "local_data") # 创建共享内存张量，形状与 local_data 一致 shmem_tensor = pypto.distributed.create_shmem_tensor( group_name="example", n_pes=4, dtype=pypto.DT_FP32, shape=[16, 64] ) # 将数据写入目标 PE 的共享内存 put_out = pypto.distributed.shmem_put( src=local_data, offsets=[0, 0], dst=shmem_tensor, dst_pe=1, put_op=pypto.AtomicType.SET, ) # 从目标 PE 的共享内存读取数据 get_out = pypto.distributed.shmem_get( src=shmem_tensor, src_pe=1, pred=[put_out], )

6.2 信号同步

import pypto # 创建信号张量 signal_tensor = pypto.distributed.create_shmem_signal( group_name="sync", n_pes=4 ) # 设置 TileShape pypto.set_vec_tile_shapes(32, 64) # PE 0 发送信号通知 PE 1 signal_out = pypto.distributed.shmem_signal( src=signal_tensor, src_pe=0, signal=1, target_pe=1, sig_op=pypto.AtomicType.SET, ) # PE 1 等待信号 wait_out = pypto.distributed.shmem_wait_until( src=signal_tensor, src_pe=0, cmp=pypto.OpType.EQ, cmp_value=1, clear_signal=True, pred=[signal_out], )

7. 相关文档

• 分布式 API 详细文档 - 查看各个函数的详细说明和约束条件
• 分布式故障排查 - 常见问题和解决方案

【免费下载链接】pyptoPyPTO（发音: pai p-t-o）：Parallel Tensor/Tile Operation编程范式。项目地址: https://gitcode.com/cann/pypto