PyPTO分布式共享内存加载

pypto.distributed.shmem_load【免费下载链接】pyptoPyPTO发音: pai p-t-oParallel Tensor/Tile Operation编程范式。项目地址: https://gitcode.com/cann/pypto产品支持情况产品是否支持Atlas A3 推理系列产品√Atlas A2 推理系列产品√功能说明从输入的 shared memory tensor 中取出部分视图到本地。函数原型shmem_load( src: ShmemTensor, src_pe: Union[int, SymbolicScalar], shape: list[int] None, offsets: list[Union[int, SymbolicScalar]] None, *, pred: list[Tensor] None, valid_shape: Optional[list[Union[int, SymbolicScalar]]] None, ) - Tensor参数说明参数名输入/输出说明src输入源操作数一个 shared memory tensor。src_pe输入shared memory tensor 所属的 pe。支持的数据类型为 int 或 SymbolicScalar 类型。0 src_pe n_pes。shape输入需要获取的视图大小。参数类型为 list[int] 类型。offsets输入需要获取的视图偏移量。支持 int 或 SymbolicScalar 类型的列表。offsets 的维度应与 src 的维度一致且每个维度的偏移量值应小于 src 对应维度的大小。pred输入用于控制操作执行的依赖关系张量列表。对数据类型无要求。不支持空 Tensor。valid_shape输入用于指定需要获取的有效数据大小。需要保证 valid_shape 小于 shape。返回值说明返回一个与 src 数据类型和形状相同的 Tensor。约束说明shmem_load 通常在 shmem_wait_until 之后执行以保证要获取的数据已经写入到了目标地址上。在 shmem_wait_until 切块数据大于1的场景下shmem_load 需要与其保持相同的切块配置以便两者能够形成更优的流水排布并保证精度正常。调用示例TileShape 设置示例说明调用该接口前应通过 set_vec_tile_shapes 设置 TileShape。TileShape 维度应和输出一致。示例 1输入的 shape 为 [m, n]输出的 shape 为 [m, n]TileShape设置为 [m1, n1]则 m1n1 分别用于切分 mn 轴。pypto.set_vec_tile_shapes(4, 8)接口调用示例示例 1从 pe 1 的 shared memory tensor 的全部视图中获取数据并输出该数据对应的输出数据 shape 为 [128, 256]。shmem_tensor pypto.distributed.create_shmem_tensor(group_nametp, n_pes8, dtypepypto.DT_FP16, shape[128, 256]) pypto.set_vec_tile_shapes(128, 256) load_out pypto.experimental.shmem_load( srcshmem_tensor, src_pe1, predpredToken, )示例 2从 pe 1 的 shared memory tensor 的部分视图中获取数据并输出该数据。该部分视图的 shape 为 [128, 128]offset 为 [0, 0]对应的输出数据 shape 为 [128, 128]实际获取的数据有效大小为 [128, 64]。shmem_tensor pypto.distributed.create_shmem_tensor(group_nametp, n_pes8, dtypepypto.DT_FP16, shape[128, 256]) pypto.set_vec_tile_shapes(128, 256) load_out pypto.experimental.shmem_load( srcshmem_tensor, src_pe1, shape[128, 128], offsets[0, 0], predpredToken, valid_shape[128, 64], )【免费下载链接】pyptoPyPTO发音: pai p-t-oParallel Tensor/Tile Operation编程范式。项目地址: https://gitcode.com/cann/pypto创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考