Skip to content

GitLab

Unverified Commit d0b7bf92 authored by zhangyunze's avatar zhangyunze Committed by GitHub
Browse files

feat:hccl support bf16

parent ade3b5da
Hide whitespace changes
Inline Side-by-side
Showing with 3 additions and 3 deletions
src/infiniccl/ascend/infiniccl_ascend.cc
View file @ d0b7bf92
...@@ -27,6 +27,8 @@ inline HcclDataType getAscendDtype(infiniDtype_t datatype) { ...@@ -27,6 +27,8 @@ inline HcclDataType getAscendDtype(infiniDtype_t datatype) {
return HCCL_DATA_TYPE_FP32; return HCCL_DATA_TYPE_FP32;
case INFINI_DTYPE_F16: case INFINI_DTYPE_F16:
return HCCL_DATA_TYPE_FP16; return HCCL_DATA_TYPE_FP16;
case INFINI_DTYPE_BF16:
return HCCL_DATA_TYPE_BFP16;
default: default:
std::cerr << "Unsupported data type: " << datatype << std::endl; std::cerr << "Unsupported data type: " << datatype << std::endl;
std::abort(); std::abort();
...@@ -86,9 +88,7 @@ infiniStatus_t allReduce( ...@@ -86,9 +88,7 @@ infiniStatus_t allReduce(
infinicclComm_t comm, infinicclComm_t comm,
infinirtStream_t stream) { infinirtStream_t stream) {
if (datatype != INFINI_DTYPE_F32 && datatype != INFINI_DTYPE_F16) { CHECK_DTYPE(datatype, INFINI_DTYPE_F32, INFINI_DTYPE_F16, INFINI_DTYPE_BF16);
return INFINI_STATUS_BAD_PARAM;
}
CHECK_HCCL(HcclAllReduce(sendbuf, recvbuf, (uint64_t)count, CHECK_HCCL(HcclAllReduce(sendbuf, recvbuf, (uint64_t)count,
getAscendDtype(datatype), getHcclRedOp(op), getAscendDtype(datatype), getHcclRedOp(op),
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment