Support FP16 for user buffer (NVIDIA#690)

* add_dtype_for_userbuf

Signed-off-by: Hongbin Liu <[email protected]>

* Update transformer_engine/pytorch/module/base.py

Signed-off-by: Tim Moon <[email protected]>

* Fix syntax

Signed-off-by: Kirthi Shankar Sivamani <[email protected]>

* Fix lint

Signed-off-by: Kirthi Shankar Sivamani <[email protected]>

---------

Signed-off-by: Hongbin Liu <[email protected]>
Signed-off-by: Tim Moon <[email protected]>
Signed-off-by: Kirthi Shankar Sivamani <[email protected]>
Co-authored-by: Hongbin Liu <[email protected]>
Co-authored-by: Tim Moon <[email protected]>
Co-authored-by: Kirthi Shankar Sivamani <[email protected]>

Author: 4 people

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 add_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