From 646a2c2bed4eeca8f5a88a2ee8a551c78c82dac1 Mon Sep 17 00:00:00 2001 From: ssit Date: Sun, 24 Sep 2023 15:40:52 -0400 Subject: [PATCH] More detailed weighted attention error logs --- fabric/fabric.py | 1 + fabric/weighted_attn.py | 3 +++ 2 files changed, 4 insertions(+) diff --git a/fabric/fabric.py b/fabric/fabric.py index 78f3b11..b1f912a 100644 --- a/fabric/fabric.py +++ b/fabric/fabric.py @@ -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 diff --git a/fabric/weighted_attn.py b/fabric/weighted_attn.py index 267a9ce..1405388 100644 --- a/fabric/weighted_attn.py +++ b/fabric/weighted_attn.py @@ -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