Switch bf16 quantize ops to at::kBFloat16 (#5735)

Summary: Pull Request resolved: #5735

Reviewed By: spcyppt

Differential Revision: D103884190

Pulled By: q10

fbshipit-source-id: b062b27c778c21711af8f8b29de36adb5daefc2d

Author: cyyever

fbgemm_gpu/src/quantize_ops/quantize_bfloat16.cu

@@ -24,20 +24,18 @@ namespace fbgemm_gpu {
 DLL_PUBLIC at::Tensor _float_to_bfloat16_gpu(const at::Tensor& input) {
   CUDA_DEVICE_GUARD(input);
 
-  // TODO: replace Half by BFloat16, after BFloat16 is supported by Nvidia
-  // NCCL input.options().dtype(at::kBFloat16)); // at::kBFloat16
-  auto output = at::empty({}, input.options().dtype(at::kHalf));
+  auto output = at::empty({}, input.options().dtype(at::kBFloat16));
   output.resize_(0);
 
   auto iter = at::TensorIteratorConfig()
                   .check_all_same_dtype(false)
                   .add_output(output)
                   .add_input(input)
                   .build();
-  at::native::gpu_kernel(iter, [] GPU_LAMBDA(float in) -> at::Half {
+  at::native::gpu_kernel(iter, [] GPU_LAMBDA(float in) -> at::BFloat16 {
     fbgemm_gpu::fint32 temp;
     temp.F = in;
-    return at::Half((temp.I + (1 << 15)) >> 16, at::Half::from_bits());
+    return at::BFloat16((temp.I + (1 << 15)) >> 16, at::BFloat16::from_bits());
   });
 
   return output;
@@ -62,7 +60,7 @@ DLL_PUBLIC at::Tensor _bfloat16_to_float_gpu(const at::Tensor& input) {
                   .add_input(input)
                   .build();
 
-  at::native::gpu_kernel(iter, [] GPU_LAMBDA(at::Half in) -> float {
+  at::native::gpu_kernel(iter, [] GPU_LAMBDA(at::BFloat16 in) -> float {
     fbgemm_gpu::fint32 temp;
     temp.I = in.x << 16;
     return temp.F;

fbgemm_gpu/src/sparse_ops/sparse_ops_cpu.cpp

@@ -593,24 +593,20 @@ void BFloat16QuantizedToFloat_ref(
   }
 }
 
-// TODO: replace Half by BFloat16, after BFloat16 is supported by Nvidia NCCL
 at::Tensor _float_to_bfloat16_cpu(const at::Tensor& input) {
   TENSOR_ON_CPU(input);
 
   const auto input_sizes = input.sizes();
-  auto output = at::empty(
-      input_sizes,
-      input.options().dtype(at::kHalf)); // at::kHalf
+  auto output = at::empty(input_sizes, input.options().dtype(at::kBFloat16));
 
   FloatToBFloat16Quantized_ref(
       input.const_data_ptr<float>(),
       input.numel(),
-      reinterpret_cast<uint16_t*>(output.mutable_data_ptr<at::Half>()));
+      reinterpret_cast<uint16_t*>(output.mutable_data_ptr<at::BFloat16>()));
 
   return output;
 }
 
-// TODO: replace Half by BFloat16, after BFloat16 is supported by Nvidia NCCL
 at::Tensor _bfloat16_to_float_cpu(const at::Tensor& input) {
   TENSOR_ON_CPU(input);
 
@@ -619,7 +615,7 @@ at::Tensor _bfloat16_to_float_cpu(const at::Tensor& input) {
   auto output = at::empty(input_sizes, input.options().dtype(at::kFloat));
 
   BFloat16QuantizedToFloat_ref(
-      reinterpret_cast<const at::BFloat16*>(input.const_data_ptr<at::Half>()),
+      input.const_data_ptr<at::BFloat16>(),
       input.numel(),
       output.mutable_data_ptr<float>());
 

fbgemm_gpu/test/quantize/bfloat16_test.py

@@ -27,6 +27,7 @@ class SparseNNOperatorsGPUTest(unittest.TestCase):
         k=st.integers(min_value=2, max_value=2),
         n=st.integers(min_value=2, max_value=2),
     )
+    @settings(deadline=10000, suppress_health_check=[HealthCheck.filter_too_much])
     def test_dense_mlp_quantize_ops(
         self, precision: str, batch_size: int, k: int, n: int
     ) -> None:
@@ -69,17 +70,15 @@ def test_quantize_op(self, nrows: int, ncols: int) -> None:
             return
         f = np.vectorize(lambda x: bfloat_quantize(x))
         reference = f(input_data.numpy())
-        quantized_data_uint16 = quantized_data.numpy()
-        quantized_data_uint16.dtype = np.uint16
+        quantized_data_uint16 = quantized_data.view(torch.uint16).numpy()
         np.testing.assert_array_almost_equal(quantized_data_uint16, reference)
 
         if torch.cuda.is_available():
             input_data_gpu = input_data.cuda()
             quantized_data_gpu = torch.ops.fbgemm.FloatToBfloat16Quantized(
                 input_data_gpu
             )
-            quantized_data_numpy = quantized_data_gpu.cpu().numpy()
-            quantized_data_numpy.dtype = np.uint16
+            quantized_data_numpy = quantized_data_gpu.view(torch.uint16).cpu().numpy()
             np.testing.assert_allclose(quantized_data_numpy, reference)
 
     # pyre-fixme[56]: Pyre was not able to infer the type of argument