def all_to_all_v_c(send_data, send_count_matrix, rank, group="hccl_world_group")
集合通信域alltoallvc操作接口。向通信域内所有rank发送数据(数据量可以定制),并从所有rank接收数据。
alltoallvc通过输入参数send_count_matrix传入所有rank的收发参数,与alltoallv相比,性能更优。
参数名 |
输入/输出 |
描述 |
---|---|---|
send_data |
输入 |
待发送的数据。 tensorflow的tensor类型。 tensor支持的数据类型为 int8, uint8, uint16, int32, uint32, int64, uint64, float16, float32, float64 。 |
send_count_matrix |
输入 |
所有rank的收发参数,send_count_matrix[i][j]表示rank i发给rank j的数据量,基本单位是send_data_type的字节数。 例:send_data_type为int32,send_count_matrix[0][1]=1,表示rank0给rank1发送1个int32 tensorflow的tensor类型。tensor支持的数据类型为int64。 |
rank |
输入 |
int类型。 本节点的rank id,该id是group内的rank id。 |
group |
输入 |
group名称,可以为用户自定义group或者"hccl_world_group"。 String类型,最大长度为128字节,含结束符。 |
recv_data:对输入tensor执行完all_to_all_v_c操作之后的结果tensor。
from npu_bridge.npu_init import * result = hccl_ops.all_to_all_v_c(send_data_tensor, send_count_matrix_tensor, rank_tensor)