Upload folder using huggingface_hub

debug_sdpa2.py

@@ -0,0 +1,62 @@
+import torch
+import math
+
+query = torch.load("query_states_sdpa_3.pt")
+key = torch.load("key_states_sdpa_3.pt")
+value = torch.load("value_states_sdpa_3.pt")
+
+print("query", query.device, query.dtype, query.is_contiguous(), query.shape)
+print("key", key.device, key.dtype, key.is_contiguous(), key.shape)
+print("value", value.device, value.dtype, value.is_contiguous(), value.shape)
+
+torch.set_printoptions(threshold=1000000, precision=6)
+
+def scaled_dot_product_attention(query, key, value, is_causal: bool, custom_cast: bool):
+    scale_factor = math.sqrt(1 / math.sqrt(query.size(-1)))
+    assert not is_causal
+
+    softmax_inp = ((query * scale_factor) @ ((key * scale_factor).transpose(-2, -1)))
+
+    if custom_cast:
+        attn_weight = torch.softmax(
+            softmax_inp,
+            dim=-1,
+            dtype=torch.float32
+        ).to(query.dtype)
+    else:
+        attn_weight = torch.softmax(softmax_inp, dim=-1)
+
+    return attn_weight @ value, softmax_inp
+
+is_causal = False
+
+with torch.no_grad():
+    with torch.backends.cuda.sdp_kernel(enable_flash=False, enable_math=True, enable_mem_efficient=False):
+        res_sdpa = torch.nn.functional.scaled_dot_product_attention(query, key, value, is_causal=is_causal, attn_mask=None)
+
+        res_eager_cast, softmax_inp = scaled_dot_product_attention(query, key, value, is_causal=is_causal, custom_cast=True)
+
+        res_eager_no_cast, _ = scaled_dot_product_attention(query, key, value, is_causal=is_causal, custom_cast=False)
+
+res_softmax_0 = torch.softmax(softmax_inp, dim=-1)
+res_softmax_1 = torch.softmax(softmax_inp, dim=-1, dtype=torch.float32).to(torch.float16)
+
+print("-----")
+
+absdiff = (res_softmax_0 - res_softmax_1).abs()
+print("max absdiff softmax", absdiff.max())
+print("median absdiff softmax", absdiff.median())
+
+print("-----")
+
+# These cast do not seem to matter.
+res_sdpa = res_sdpa.to(torch.float32)
+res_eager_cast = res_eager_cast.to(torch.float32)
+res_eager_no_cast = res_eager_no_cast.to(torch.float32)
+
+absdiff_nocast = (res_sdpa - res_eager_no_cast).abs()
+absdiff_cast = (res_sdpa - res_eager_cast).abs()
+print("SDPA max absdiff (no cast):", absdiff_nocast.max())
+print("SDPA max absdiff (with cast):", absdiff_cast.max())
+print("argwhere absdiff no cast > 0.0001", torch.argwhere(absdiff_nocast > 1e-4))
+print("argwhere absdiff with cast > 0.0001", torch.argwhere(absdiff_cast > 1e-4))

key_states_sdpa_3.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b3bae5605ebe7843d99b5ca477993c6584544f1825663c26569e3a1b22c1341d
+size 10142926

query_states_sdpa_3.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5dad381f5ce6bbb3a1ec26b8d578d92395de8d57b11860a9c1726f9caa15887a
+size 10142936

value_states_sdpa_3.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:149317f8c4c45572bde47f6aeb704219cd9eccf7ca311d7ec66fb8e08de63219
+size 30426328