PyPTO分布式共享内存加载

pypto.distributed.shmem_load

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

产品支持情况

功能说明

从输入的 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 数据类型和形状相同的 Tensor。

约束说明

1. 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]，则 m1，n1 分别用于切分 m，n 轴。

  pypto.set_vec_tile_shapes(4, 8)

接口调用示例

• 示例 1：从 pe = 1 的 shared memory tensor 的全部视图中获取数据并输出该数据，对应的输出数据 shape 为 [128, 256]。

  shmem_tensor = pypto.distributed.create_shmem_tensor(group_name="tp", n_pes=8, dtype=pypto.DT_FP16, shape=[128, 256]) pypto.set_vec_tile_shapes(128, 256) load_out = pypto.experimental.shmem_load( src=shmem_tensor, src_pe=1, pred=predToken, )

• 示例 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_name="tp", n_pes=8, dtype=pypto.DT_FP16, shape=[128, 256]) pypto.set_vec_tile_shapes(128, 256) load_out = pypto.experimental.shmem_load( src=shmem_tensor, src_pe=1, shape=[128, 128], offsets=[0, 0], pred=predToken, valid_shape=[128, 64], )

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