Support FP16 for user buffer (#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]>
NVIDIA · Mar 7, 2024 · 8255f87 · 8255f87
1 parent b459ccc
commit 8255f87
Showing 1 changed file with 5 additions and 2 deletions.
diff --git a/transformer_engine/pytorch/module/base.py b/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