Commit b8609df3 authored by Ceng's avatar Ceng Committed by Ceng2333
Browse files

issue/434 hccl support bf16


Signed-off-by: default avatarCeng <441651826@qq.com>
parent f9d16628
...@@ -23,6 +23,8 @@ inline hcclDataType_t getHcclDtype(infiniDtype_t datatype) { ...@@ -23,6 +23,8 @@ inline hcclDataType_t getHcclDtype(infiniDtype_t datatype) {
return hcclFloat; return hcclFloat;
case INFINI_DTYPE_F16: case INFINI_DTYPE_F16:
return hcclHalf; return hcclHalf;
case INFINI_DTYPE_BF16:
return hcclBfloat16;
default: default:
std::abort(); std::abort();
return hcclHalf; return hcclHalf;
...@@ -83,9 +85,7 @@ infiniStatus_t allReduce( ...@@ -83,9 +85,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, count, getHcclDtype(datatype), CHECK_HCCL(hcclAllReduce(sendbuf, recvbuf, count, getHcclDtype(datatype),
getHcclRedOp(op), getHcclComm(comm), getMacaStream(stream))); getHcclRedOp(op), getHcclComm(comm), getMacaStream(stream)));
......
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