I'm arising this issue for dlpack bfloat16 support.
Bfloat16 is a popular 16-bit floating point format for machine learning, supported by multiple hardware, e.g. TPU. Compared to fp16, bfloat16 has a greater dynamic range, so it's useful for things like gradients that can be outside the dynamic range of fp16. Compared to fp32, Using bfloat16 reduces the size of data in memory and allows larger models to fit in the same amount of memory. So there are many advantages of bfloat16, it's a trend for different frameworks to support bfloat16. Tensorflow has already supported bfloat16 data type. And we are now supporting bfloat16 in MXNet.
Dlpack is an open in-memory tensor structure for sharing tensors among deep learning frameworks. Supporting bfloat16 can make dlpack more flexible and integrated in data sharing between different frameworks.
Current status of dlpack and bfloat16 support in Frameworks:
1. Pytorch:
Pytorch has two interfaces for converting data from/to dlpack format. tsor = torch.utils.dlpack.from_dlpack(dl)
converts data from dlpack-defined tensor to pytorch-defined tensor. dl = torch.utils.dlpack.to_dlpack(tsor)
converts data from pytorch-defined tensor to dlpack-defined tensor. And when using to_dlpack
function, getDLDataType
is used to check the data types that have been enabled for data sharing in dlpack:
DLDataType getDLDataType(const Tensor& t) {
DLDataType dtype;
dtype.lanes = 1;
dtype.bits = t.element_size() * 8;
switch (t.scalar_type()) {
case ScalarType::Byte:
dtype.code = DLDataTypeCode::kDLUInt;
break;
case ScalarType::Char:
dtype.code = DLDataTypeCode::kDLInt;
break;
case ScalarType::Double:
dtype.code = DLDataTypeCode::kDLFloat;
break;
case ScalarType::Float:
dtype.code = DLDataTypeCode::kDLFloat;
break;
case …
case ScalarType::BFloat16:
throw std::logic_error("BFloat16 is not supported by dlpack");
break;
For now as dlpack has not supported bfloat16 yet, getDLDataType
throws an error when encountering bfloat16 data type. Once dlpack supports bfloat16 data type, this code can be easily changed.
2. MXNet:
Similar to pytorch, mxnet also has arr = mx.nd.from_dlpack(dl)
, dl = mx.nd.to_dlpack_for_read(arr)
and dl = mx.nd.to_dlpack_for_write(arr)
for dlpack/mxnet data sharing. Also DTypeTransform
is used to check the data types.
static DLDataType DTypeTransform(int type_flag) {
switch (type_flag) {
case mshadow::kFloat32: return DLDataType{kDLFloat, 32, 1};
case mshadow::kFloat64: return DLDataType{kDLFloat, 64, 1};
case mshadow::kFloat16: return DLDataType{kDLFloat, 16, 1};
case mshadow::kBfloat16: return DLDataType{kDLBfloat, 16, 1}; // add this line to support bfloat16
case ......
}
}
add bfloat16 data type support in this function and then we can use this data type as inputs, params or outputs for operator computation.
3. Tensorflow:
Tensorflow haven't support dlpack yet, but there's a discussion on it (issue). Tensorflow has already support bfloat16.
As discussed above, we can see that bfloat16 has a good support in various frameworks. On the other hand, dlpack is also becoming more and more popular. So it will be really great if dlpack can have bfloat16 data type support.
Proposal for supporting bfloat16 in dlpack:
Here is a draft proposal for supporting bfloat16 in dlpack. the modification in dlpack will be very simple, just add one single line in DLDataTypeCode:
typedef enum {
kDLInt = 0U,
kDLUInt = 1U,
kDLFloat = 2U,
kDLBfloat = 3U, // add this line to support bfloat16
} DLDataTypeCode;
And it's done.
Do you have any ideas? Thank you @soumith @piiswrong @Yangqing @naibaf7 @bhack @edgarriba @tqchen @prigoyal @zdevito @pengzhao-intel @ZhennanQin