Merge branch 'rs-dgrad-overlap-bugfix' of github.com:denera/Transform…

…erEngine into rs-dgrad-overlap-bugfix
NVIDIA · Dec 5, 2024 · 340e033 · 340e033
2 parents 2535d77 + 4764591
commit 340e033
Showing 1 changed file with 2 additions and 3 deletions.
diff --git a/transformer_engine/pytorch/module/base.py b/transformer_engine/pytorch/module/base.py
@@ -418,9 +418,8 @@ def add_ub(
         if name in ub_cfgs:
             final_cfg = get_default_config(name)
             final_cfg.update(ub_cfgs[name])
-            final_cfg["fp8_buf"] = (
-                (name in layers_all_gather_overlap)
-                or ub_cfgs[name].get("fp8_buf", False)
+            final_cfg["fp8_buf"] = (name in layers_all_gather_overlap) or ub_cfgs[name].get(
+                "fp8_buf", False
             )
             add_ub(name, **final_cfg)