Unittest for the MultiHeadSelfAttention

[vipinanandcpp]

Hi guys I am having issues with the unittest for the MultiHeadSelfAttention. I have attached the implementation class below along with the adaptor function (run_multihead_self_attention). My unittest function (test_multihead_self_attention) fails with 2817/3072 unmatched components between the output Tensors. Any idea where I could be going wrong with this?

adapters.py
test_model.py
attention.py

class MultiHeadSelfAttention(torch.nn.Module):
    def __init__(self, embed_size: int, d_k: int, d_v: int, num_heads: int, dropout: float=0.0, device: torch.device=None, dtype: torch.dtype=None):
        super().__init__()
        self.num_heads: int = num_heads
        self.W_Q = Linear(in_features=embed_size, out_features=num_heads*d_k, device=device, dtype=dtype)
        self.W_K = Linear(in_features=embed_size, out_features=num_heads*d_k, device=device, dtype=dtype)
        self.W_V = Linear(in_features=embed_size, out_features=num_heads*d_v, device=device, dtype=dtype)
        self.W_O = Linear(in_features=num_heads*d_v, out_features=embed_size, device=device, dtype=dtype)
        self.dropout = torch.nn.Dropout(dropout)
        self.d_k: int = d_k
        self.d_v: int = d_v
    
    # x: [batch_size, seq_len, embed_size]
    def _get_normalized_scores(self, x: torch.Tensor) -> torch.Tensor:
        # [batch_size, seq_len, num_heads*d_k]
        Q: torch.Tensor = self.W_Q(x)
        # [batch_size, seq_len, num_heads, d_k] -> [batch_size, num_heads, seq_len, d_k]
        Q = Q.view(Q.shape[0], Q.shape[1], self.num_heads, self.d_k).transpose(1, 2)
        # [batch_size, seq_len, num_heads*d_k]
        K: torch.Tensor = self.W_K(x)
        # [batch_size, seq_len, num_heads, d_k] -> [batch_size, num_heads, seq_len, d_k]
        K = K.view(K.shape[0], K.shape[1], self.num_heads, self.d_k).transpose(1, 2)
        # [batch_size, num_heads, seq_len, seq_len]
        return (Q @ K.transpose(-2, -1)) / (K.shape[-1] ** 0.5)

    # x: [batch_size, seq_len, embed_size]
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # [batch_size, seq_len, num_heads*d_v]
        V: torch.Tensor = self.W_V(x)
        # [batch_size, seq_len, num_heads, d_v] -> [batch_size, num_heads, seq_len, d_v]
        V = V.view(V.shape[0], V.shape[1], self.num_heads, self.d_v).transpose(1, 2)
        # [batch_size, num_heads, seq_len, seq_len]
        attention_scores: torch.Tensor = self._get_normalized_scores(x)
        # [batch_size, num_heads, seq_len, seq_len]
        attention_weights: torch.Tensor = self.dropout(softmax(attention_scores, dim=-1))
        # [batch_size, num_heads, seq_len, d_v] -> [batch_size, seq_len, num_heads, d_v]
        attention_weights = (attention_weights @ V).transpose(1, 2)
        # [batch_size, seq_len, num_heads*d_v]
        attention_weights = attention_weights.flatten(-2, -1)
        # [batch_size, seq_len, embed_size]
        attention_weights = self.W_O(attention_weights)
        return attention_weights


class Linear(torch.nn.Module):
    def __init__(self, in_features: int, out_features: int, device: torch.device=None, dtype: torch.dtype=None):
        super().__init__()
        self.weight = torch.nn.Parameter(random_normal(torch.empty(out_features, in_features,
                                                                   device=device, dtype=dtype),
                                                                   mean=0.0, std=(2/(in_features+out_features))**0.5, 
                                                                   a=-3*(2/(in_features+out_features))**0.5,
                                                                   b=3*(2/(in_features+out_features))**0.5))
        self.bias = torch.zeros(out_features, device=device, dtype=dtype)

    # x: [batch_size, in_features]
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        print("Shapes", x.shape, self.weight.T.shape)
        return (x @ self.weight.T) + self.bias


def run_multihead_self_attention(
    d_model: int,
    num_heads: int,
    q_proj_weight: Tensor,
    k_proj_weight: Tensor,
    v_proj_weight: Tensor,
    o_proj_weight: Tensor,
    in_features: Tensor
) -> Tensor:
    """
    Given the query, key, and value tensors, compute the scaled dot product attention.

    Args:
        d_model (int): Embedding dimension size for the query or key tensor.
        num_heads (int): Number of heads.
        q_proj_weight (Tensor): Query projection weights [d_k, d_in].
        k_proj_weight (Tensor): Key projection weights [d_k, d_in].
        v_proj_weight (Tensor): Value projection weights [d_v, d_in].
        o_proj_weight (Tensor): Output projection weights [d_out, d_in].
        in_features (Tensor): Input features to run MultiHeadSelfAttention on.

    Returns:
        Tensor: Scaled dot product attention output.
    """

    attention = MultiHeadSelfAttention(embed_size=d_model, d_k=d_model//num_heads, d_v=d_model//num_heads, num_heads=num_heads)
    attention.W_Q.weight = nn.Parameter(q_proj_weight)
    attention.W_K.weight = nn.Parameter(k_proj_weight)
    attention.W_V.weight = nn.Parameter(v_proj_weight)
    attention.W_O.weight = nn.Parameter(o_proj_weight)

    x: Tensor = in_features
    return attention(x)
    

def test_multihead_self_attention(numpy_snapshot: NumpySnapshot, in_embeddings: Tensor, d_model: int, n_heads: int, ts_state_dict: Dict[str, Tensor]):
    d, _ = ts_state_dict
    q_proj_weight, k_proj_weight, v_proj_weight, o_proj_weight = [
        d[f"layers.0.attn.{k}_proj.weight"] for k in ["q", "k", "v", "output"]
    ]
    actual_output = run_multihead_self_attention(
        d_model=d_model,
        num_heads=n_heads,
        q_proj_weight=q_proj_weight,
        k_proj_weight=k_proj_weight,
        v_proj_weight=v_proj_weight,
        o_proj_weight=o_proj_weight,
        in_features=in_embeddings
    )
    numpy_snapshot.assert_match(actual_output, atol=1e-6)