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

for more information, see https://pre-commit.ci

Author: pre-commit-ci[bot]

tests/pytorch/attention/run_attention_with_cp.py

@@ -90,17 +90,22 @@ def generate_input_shapes(
         cu_seqlens_kv_padded = None
     elif qkv_format == "thd":
         # seqlens_q = torch.randint(0, config.max_seqlen_q + 1, [config.batch_size]).to(torch.int32)
-        seqlens_q = torch.ones([config.batch_size], dtype=torch.int32).to(torch.int32) * config.max_seqlen_q
+        seqlens_q = (
+            torch.ones([config.batch_size], dtype=torch.int32).to(torch.int32) * config.max_seqlen_q
+        )
         seqlens_q_padded = (seqlens_q + 2 * world_size - 1) // (world_size * 2) * (world_size * 2)
         cu_seqlens_q_padded = torch.cat(
             [
                 torch.zeros([1], dtype=torch.int32),
                 seqlens_q_padded.cumsum(0, dtype=torch.int32),
-                #torch.tensor([q_input_shape[0]], dtype=torch.int32),
+                # torch.tensor([q_input_shape[0]], dtype=torch.int32),
             ]
         ).cuda()
         cu_seqlens_q = torch.clone(cu_seqlens_q_padded)
-        print(f"dev {torch.cuda.current_device()} cu_seqlens_q: {cu_seqlens_q}, cu_seqlens_q_padded: {cu_seqlens_q_padded}")
+        print(
+            f"dev {torch.cuda.current_device()} cu_seqlens_q: {cu_seqlens_q}, cu_seqlens_q_padded:"
+            f" {cu_seqlens_q_padded}"
+        )
         q_input_shape = (
             cu_seqlens_q_padded[-1],
             config.num_heads,
@@ -266,9 +271,9 @@ def run_dpa_with_cp(
         cu_seqlens_q_padded,
         cu_seqlens_kv_padded,
     ) = generate_input_shapes(qkv_format, config, world_size, kernel_backend)
-    q_orig = torch.clamp(2*torch.ones(q_input_shape, dtype=dtypes[dtype]), min=-1, max=2).cuda()
-    k_orig = torch.clamp(2*torch.ones(k_input_shape, dtype=dtypes[dtype]), min=-1, max=2).cuda()
-    v_orig = torch.clamp(2*torch.ones(v_input_shape, dtype=dtypes[dtype]), min=-1, max=2).cuda()
+    q_orig = torch.clamp(2 * torch.ones(q_input_shape, dtype=dtypes[dtype]), min=-1, max=2).cuda()
+    k_orig = torch.clamp(2 * torch.ones(k_input_shape, dtype=dtypes[dtype]), min=-1, max=2).cuda()
+    v_orig = torch.clamp(2 * torch.ones(v_input_shape, dtype=dtypes[dtype]), min=-1, max=2).cuda()
     dout_orig = torch.clamp(
         torch.randn(attn_output_shape, dtype=dtypes[dtype]), min=-1, max=1
     ).cuda()
@@ -448,7 +453,9 @@ def run_dpa_with_cp(
         if not fp8_bwd:
             tensors[0], tensors[4] = tensors_to_deq
     for i, tensor in enumerate(tensors):
-        print(f"dev {torch.cuda.current_device()} tensor {i} is nan: {torch.isnan(tensor).nonzero()}")
+        print(
+            f"dev {torch.cuda.current_device()} tensor {i} is nan: {torch.isnan(tensor).nonzero()}"
+        )
         # print(f"dev {torch.cuda.current_device()} tensor {i} is inf: {torch.isinf(tensor).non_zero()}")
         assert torch.all(~torch.isnan(tensor))
         assert torch.all(~torch.isinf(tensor))

tests/pytorch/attention/test_attention.py

@@ -240,7 +240,7 @@ def test_dot_product_attention(
         if config.return_max_score:
             torch.testing.assert_close(fused_max_score, unfused_max_score, **tols)
         for i, _ in enumerate(unfused_attn_bwd):
-            print(f'iiiiii {i}')
+            print(f"iiiiii {i}")
             torch.testing.assert_close(fused_attn_bwd[i], unfused_attn_bwd[i], **tols)
     if fused_attn_supported and flash_attn_supported:
         logging.info("[test_dot_product_attention]: fused attn vs flash attn")
@@ -282,7 +282,9 @@ def test_dpa_max_score(dtype, model_configs, model):
     """Test DotProductAttention module with checkpointing"""
     config = model_configs[model]
     config.return_max_score = True
-    test_dot_product_attention(dtype, model_configs, model, False, True, "thd_thd_thd", False, False)
+    test_dot_product_attention(
+        dtype, model_configs, model, False, True, "thd_thd_thd", False, False
+    )
 
 
 model_configs_softmax = {

tests/pytorch/attention/test_attention_with_cp.py

@@ -183,7 +183,7 @@ def test_cp_with_flash_attention(dtype, model, qkv_format, cp_comm_type):
 qkv_formats = ["bshd", "sbhd", "thd"]
 cp_comm_types = ["p2p", "all_gather", "a2a", "a2a+p2p"]
 if test_essential:
-    configs = ["cp_1_0", "cp_1_1"] #, "cp_2_0", "cp_2_2", "cp_3_2", "cp_4_2"]
+    configs = ["cp_1_0", "cp_1_1"]  # , "cp_2_0", "cp_2_2", "cp_3_2", "cp_4_2"]
     model_configs_fused_attn = {k: model_configs_fused_attn[k] for k in configs}
     dtypes = ["bf16", "fp8"]
     qkv_formats = ["sbhd", "thd"]

transformer_engine/common/fused_attn/utils.h

@@ -122,8 +122,7 @@ struct FADescriptor_v1 {
                     page_size_v, max_pages_per_seq_k, max_pages_per_seq_v, bias_b, bias_h,
                     attnScale, isTraining, dropoutProbability, layout, mask_type, softmax_type,
                     window_size_left, window_size_right, deterministic, bias_type, qkv_tensor_type,
-                    o_tensor_type, do_tensor_type, dqkv_tensor_type,
-                    generate_max_sum_exp) <
+                    o_tensor_type, do_tensor_type, dqkv_tensor_type, generate_max_sum_exp) <
            std::tie(rhs.b, rhs.h, rhs.hg, rhs.s_q, rhs.s_kv, rhs.d_qk, rhs.d_v, rhs.num_pages_k,
                     rhs.num_pages_v, rhs.page_size_k, rhs.page_size_v, rhs.max_pages_per_seq_k,
                     rhs.max_pages_per_seq_v, rhs.bias_b, rhs.bias_h, rhs.attnScale, rhs.isTraining,

transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py

@@ -1617,7 +1617,10 @@ def forward(
                                 softmax_lse_per_step[i - 1],
                             )
                     if return_max_score:
-                        print(f"dev={torch.cuda.current_device()} i={i}, max_score_per_step={max_score_per_step[i - 1]}")
+                        print(
+                            f"dev={torch.cuda.current_device()} i={i},"
+                            f" max_score_per_step={max_score_per_step[i - 1]}"
+                        )
                         if i == 1:
                             max_score = torch.clone(max_score_per_step[0])
                         else:

transformer_engine/pytorch/cpp_extensions/fused_attn.py

@@ -324,8 +324,12 @@ def fused_attn_fwd(
     )
 
     if return_max_score:
-        qkv_format = qkv_layout.replace("3","").replace("2","").split("_")[0]
-        print(f"dev {torch.cuda.current_device()} qkv_format: {qkv_format}, cu_seqlens_q: {cu_seqlens_q}, cu_seqlens_kv: {cu_seqlens_kv}, cu_seqlens_q_padded: {cu_seqlens_q_padded}, cu_seqlens_kv_padded: {cu_seqlens_kv_padded}")
+        qkv_format = qkv_layout.replace("3", "").replace("2", "").split("_")[0]
+        print(
+            f"dev {torch.cuda.current_device()} qkv_format: {qkv_format}, cu_seqlens_q:"
+            f" {cu_seqlens_q}, cu_seqlens_kv: {cu_seqlens_kv}, cu_seqlens_q_padded:"
+            f" {cu_seqlens_q_padded}, cu_seqlens_kv_padded: {cu_seqlens_kv_padded}"
+        )
         # print(f"dev {torch.cuda.current_device()} q: {q.shape}, k: {k.shape}, v: {v.shape}")
         # print(f"dev {torch.cuda.current_device()} output_tensors[0] is nan: {torch.isnan(output_tensors[0]).sum()}, output_tensors[0]: {output_tensors[0].shape}, output_tensors[0].min(): {output_tensors[0].min()}, output_tensors[0].max(): {output_tensors[0].max()}")
         # print(f"dev {torch.cuda.current_device()} output_tensors[1] is nan: {torch.isnan(output_tensors[1]).sum()}, output_tensors[1]: {output_tensors[1].shape}, output_tensors[1].min(): {output_tensors[1].min()}, output_tensors[1].max(): {output_tensors[1].max()}")
@@ -335,7 +339,10 @@ def fused_attn_fwd(
             stats = output_tensors[1] + torch.log(output_tensors[2])
             zero_indices_1 = (output_tensors[1] == 0).nonzero()
             zero_indices_2 = (output_tensors[2] == 0).nonzero()
-            print(f"dev {torch.cuda.current_device()} zero_indices_1: {zero_indices_1}, zero_indices_2: {zero_indices_2}")
+            print(
+                f"dev {torch.cuda.current_device()} zero_indices_1: {zero_indices_1},"
+                f" zero_indices_2: {zero_indices_2}"
+            )
             if torch.cuda.current_device() == 0 and not os.path.exists("output_tensors1.pt"):
                 torch.save(output_tensors[1], "output_tensors1.pt")
                 torch.save(output_tensors[2], "output_tensors2.pt")
@@ -344,16 +351,39 @@ def fused_attn_fwd(
             # Max [tq, h, 1] -> max_score [h]
             max_score = torch.amax(output_tensors[1], dim=(0, 2)).to(dtype=output_tensors[0].dtype)
             print(f"dev {torch.cuda.current_device()} max_score: {max_score}")
-            print(f"dev {torch.cuda.current_device()} output_tensors[0] is nan: {torch.isnan(output_tensors[0]).sum()}, output_tensors[0]: {output_tensors[0].shape}, output_tensors[0].min(): {output_tensors[0].min()}, output_tensors[0].max(): {output_tensors[0].max()}")
-            print(f"dev {torch.cuda.current_device()} output_tensors[1] is nan: {torch.isnan(output_tensors[1]).sum()}, output_tensors[1]: {output_tensors[1].shape}, output_tensors[1].min(): {output_tensors[1].min()}, output_tensors[1].max(): {output_tensors[1].max()}")
-            print(f"dev {torch.cuda.current_device()} output_tensors[2] is nan: {torch.isnan(output_tensors[2]).sum()}, output_tensors[2]: {output_tensors[2].shape}, output_tensors[2].min(): {output_tensors[2].min()}, output_tensors[2].max(): {output_tensors[2].max()}")
-            print(f"dev {torch.cuda.current_device()} stats is nan: {torch.isnan(stats).sum()}, stats: {stats.shape}, stats.min(): {stats.min()}, stats.max(): {stats.max()}")
-            print(f"dev {torch.cuda.current_device()} max_score is nan: {torch.isnan(max_score).sum()}, max_score: {max_score.shape}    ")
+            print(
+                f"dev {torch.cuda.current_device()} output_tensors[0] is nan:"
+                f" {torch.isnan(output_tensors[0]).sum()}, output_tensors[0]:"
+                f" {output_tensors[0].shape}, output_tensors[0].min(): {output_tensors[0].min()},"
+                f" output_tensors[0].max(): {output_tensors[0].max()}"
+            )
+            print(
+                f"dev {torch.cuda.current_device()} output_tensors[1] is nan:"
+                f" {torch.isnan(output_tensors[1]).sum()}, output_tensors[1]:"
+                f" {output_tensors[1].shape}, output_tensors[1].min(): {output_tensors[1].min()},"
+                f" output_tensors[1].max(): {output_tensors[1].max()}"
+            )
+            print(
+                f"dev {torch.cuda.current_device()} output_tensors[2] is nan:"
+                f" {torch.isnan(output_tensors[2]).sum()}, output_tensors[2]:"
+                f" {output_tensors[2].shape}, output_tensors[2].min(): {output_tensors[2].min()},"
+                f" output_tensors[2].max(): {output_tensors[2].max()}"
+            )
+            print(
+                f"dev {torch.cuda.current_device()} stats is nan: {torch.isnan(stats).sum()},"
+                f" stats: {stats.shape}, stats.min(): {stats.min()}, stats.max(): {stats.max()}"
+            )
+            print(
+                f"dev {torch.cuda.current_device()} max_score is nan:"
+                f" {torch.isnan(max_score).sum()}, max_score: {max_score.shape}    "
+            )
         else:
             # output_tensors: out [b, sq, h, d] or [sq, b, h, d], Max [b, h, sq, 1], Sum_Exp [b, h, sq, 1]
             stats = output_tensors[1] + torch.log(output_tensors[2])
             # Max [b, h, sq, 1] -> max_score [h]
-            max_score = torch.amax(output_tensors[1], dim=(0, 2, 3)).to(dtype=output_tensors[0].dtype)
+            max_score = torch.amax(output_tensors[1], dim=(0, 2, 3)).to(
+                dtype=output_tensors[0].dtype
+            )
         aux_ctx_tensors = [stats]
         aux_ctx_tensors.extend(output_tensors[3:])
         return output_tensors[0], aux_ctx_tensors, max_score