[pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
NVIDIA · Dec 5, 2024 · 1036168 · 1036168
1 parent 340e033
commit 1036168
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 11 deletions.
diff --git a/tests/pytorch/distributed/run_layer_with_overlap.py b/tests/pytorch/distributed/run_layer_with_overlap.py
@@ -92,7 +92,7 @@ def _get_ub_cfg(config):
                 "qkv_fprop": dict(),
                 "qkv_dgrad": {
                     "method": "pipeline" if config.rs_dgrad_overlap else "bulk",
-                    "fp8_buf": True if config.fp8_buf and config.rs_dgrad_overlap else False
+                    "fp8_buf": True if config.fp8_buf and config.rs_dgrad_overlap else False,
                 },
             }
         )
@@ -101,9 +101,7 @@ def _get_ub_cfg(config):
     if config.layer_type in [te.Linear, te.MultiheadAttention, te.TransformerLayer]:
         ub_cfg.update(
             {
-                "proj_fprop": {
-                    "fp8_buf": True if config.fp8_buf else False
-                },
+                "proj_fprop": {"fp8_buf": True if config.fp8_buf else False},
                 "proj_dgrad": dict(),
             }
         )
@@ -113,7 +111,7 @@ def _get_ub_cfg(config):
                 "fc1_fprop": dict(),
                 "fc1_dgrad": {
                     "method": "pipeline" if config.rs_dgrad_overlap else "bulk",
-                    "fp8_buf": True if config.fp8_buf and config.rs_dgrad_overlap else False
+                    "fp8_buf": True if config.fp8_buf and config.rs_dgrad_overlap else False,
                 },
                 "fc2_fprop": {
                     "fp8_buf": True if config.fp8_buf else False,
@@ -150,7 +148,7 @@ def _parse_args(argv=None, namespace=None):
         "--fp8-buf",
         action="store_true",
         default=False,
-        help="Allocate FP8 communication buffers for layers that support it."
+        help="Allocate FP8 communication buffers for layers that support it.",
     )
     parser.add_argument(
         "--rs-dgrad-overlap",

diff --git a/tests/pytorch/distributed/test_comm_gemm_overlap.py b/tests/pytorch/distributed/test_comm_gemm_overlap.py
@@ -254,11 +254,7 @@ def test_bulk_overlaps(comm_type, fp8):
 )
 @pytest.mark.parametrize(
     "fp8,fp8_init",
-    [
-        (False, False),
-        (True, False),
-        (True, True)
-    ],
+    [(False, False), (True, False), (True, True)],
     ids=[
         " BF16 GEMM - BF16 PARAMS ",
         " FP8  GEMM - BF16 PARAMS ",