Support FP16 for user buffer (#690)

* add_dtype_for_userbuf
Signed-off-by: Hongbin Liu <hongbinl@nvidia.com>

* Update transformer_engine/pytorch/module/base.py
Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>

* Fix syntax
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Fix lint
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

---------
Signed-off-by: Hongbin Liu <hongbinl@nvidia.com>
Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Co-authored-by: Hongbin Liu <hongbinl@nvidia.com>
Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

transformer_engine/pytorch/module/base.py

@@ -111,6 +111,7 @@ def initialize_ub(
     shape: list,
     tp_size: int,
     use_fp8: bool = False,
+    dtype: torch.dtype = torch.bfloat16,
     ub_cfgs: Optional[dict] = None
 ) -> None:
     """Initialize communicators for TP comm overlap using userbuffers."""
@@ -151,8 +152,10 @@ def initialize_ub(
         num_splits: int = 4,
         aggregate: int = 0,
     ) -> None:
-        dtype = torch.uint8 if (use_fp8 and name in fp8_buf) else torch.bfloat16
-        sample_buffer = torch.empty(shape, dtype=dtype, device='cuda')
+        sample_buffer = torch.empty(
+            shape,
+            dtype=torch.uint8 if (use_fp8 and name in fp8_buf) else dtype,
+            device='cuda')
         if method == 'ring_exchange':
             ub_obj = tex.UbufP2PCommOverlap(
                     sample_buffer,          # Sample userbuffer