More detailed weighted attention error logs

This commit is contained in:
ssit
2023-09-24 15:40:52 -04:00
parent 96c21c5924
commit 646a2c2bed
2 changed files with 4 additions and 0 deletions
+1
View File
@@ -280,6 +280,7 @@ class FABRICPatcher:
def get_weights(pos_weight, neg_weight, q, num_pos, num_neg, cond_uncond_idxs):
"""
Prepare weights for the weighted attention
:return: Weights of shape [batch_size, nk] where nk is the size of the Key sequence length. batch_size = len(cond_uncond_idxs)
"""
input_dim = q.shape[1]
hs_dim = max(num_pos, num_neg) * input_dim
+3
View File
@@ -26,6 +26,7 @@ class Weighted_Attn_Patcher:
except Exception as e:
print(e)
print("[FABRIC] Encountered an exception. If this is not a memory issue, please report this issue.")
print(f"[FABRIC] weights b: {weights.shape[0]}, weights nk: {weights.shape[1]}, nq: {nq}, nk: {nk}, h: {h}, B: {B}")
self.unpatch()
raise e
@@ -45,6 +46,7 @@ class Weighted_Attn_Patcher:
except Exception as e:
print(e)
print("[FABRIC] Encountered an exception. If this is not a memory issue, please report this issue.")
print(f"[FABRIC] weights b: {weights.shape[0]}, weights nk: {weights.shape[1]}, nq: {nq}, nk: {nk}, h: {h}, bs: {bs}")
self.unpatch()
raise e
torch.nn.functional.scaled_dot_product_attention = pt_sdp
@@ -81,6 +83,7 @@ class Weighted_Attn_Patcher:
except Exception as e:
print(e)
print("[FABRIC] Encountered an exception. If this is not a memory issue, please report this issue.")
print(f"[FABRIC] weights b: {weights.shape[0]}, weights nk: {weights.shape[1]}, nq: {nq}, nk: {nk}, h: {h}, B: {B}")
self.unpatch()
raise e
torch.Tensor.softmax = softmax_method