Different value between BertSDPA scratch and pretrained sequence classification

jeyan9: <p>Hello, I am using the bert pretrained parameter as an input to accelerator. But the problem is that the result of BertSdpaSelfAttention in BertForSequenceClassifcation gives different output value with the torch F.scaled_dot_product_attention or scratch_sdpa. Am I missing something? I disabled dropout and layernorm for simplicity.</p>
<pre><code class="lang-auto">num_block = 12
len_seq = 128
dim_model = 768
num_head = 12

def disable_dropout(model):
    for module in model.modules():
        if isinstance(module, torch.nn.Dropout):
            module.p = 0


def disable_layernorm(m):
    if isinstance(m, nn.LayerNorm):
        m.reset_parameters()
        m.eval()
        with torch.no_grad():
            m.weight.fill_(1.0)
            m.bias.zero_()

def scratch_sdpa(query, key, value, mask):
    scale = 64 ** -0.5
    query = query.reshape(num_head, len_seq, int(dim_model/num_head))
    key = key.reshape(num_head, len_seq, int(dim_model/num_head))
    value = value.reshape(num_head, len_seq, int(dim_model/num_head))

    extend_mask = mask.expand(1, len_seq, len_seq).unsqueeze(1).repeat(1, 12, 1, 1)
    boolean_mask = (extend_mask == 0)

    scale_dot_prod = torch.einsum('bid, bjd -&gt; bij', query, key) * scale        # shape 12, 128, 128
    scale_dot_prod = scale_dot_prod.masked_fill_(boolean_mask[0], -1e9)
    attention = torch.softmax(scale_dot_prod, dim=-1)
    # attention = torch.softmax(scale_dot_prod+mask.expand(1, 128, 128), dim=-1)
    attention_output = torch.einsum("bij, bjd-&gt;bid", attention, value)
    attention_output = attention_output.reshape(len_seq, dim_model)
    return attention_output

def bert_sequence_classification():
    '''
    Example 1. Bert for sequence classificatioin (positive or negative)
    '''
    tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
    # model = BertModel.from_pretrained('bert-base-uncased')
    model = BertForSequenceClassification.from_pretrained('bert-base-uncased')
    
    disable_dropout(model)
    model.apply(disable_layernorm)

    sentence = "A visually stunning rumination on love"
    encoded_dict = tokenizer.encode_plus(
                    sentence,
                    add_special_tokens=True,
                    max_length=len_seq,
                    padding="max_length",
                    truncation=True,
                    return_attention_mask=True,
                    return_tensors="pt"
                    )
    input_ids = encoded_dict["input_ids"]
    attention_mask = encoded_dict["attention_mask"]
    labels = torch.tensor([1]).unsqueeze(0)  # Batch size 1

    # outputs = model(input_ids=input_ids, labels=labels)
    outputs = model(input_ids=input_ids, attention_mask=attention_mask, labels=labels)
    loss, logits = outputs[:2]
    # encoded_dict["attention_mask"] = torch.ones_like(encoded_dict["attention_mask"])
    return model, encoded_dict, outputs

def dump_bert(model, inputs):
    for l in range(num_block):
        block = model.bert.encoder.layer[l]
        
    # generating qkv
        query_w = block.attention.self.query.weight
        query_b = block.attention.self.query.bias
        query_o = activations[f"bert.encoder.layer.{l}.attention.self.query"]["output"][0]
        key_w = block.attention.self.key.weight
        key_b = block.attention.self.key.bias
        key_o = activations[f"bert.encoder.layer.{l}.attention.self.key"]["output"][0]
        value_w = block.attention.self.value.weight
        value_b = block.attention.self.value.bias
        value_o = activations[f"bert.encoder.layer.{l}.attention.self.value"]["output"][0]
        np.save(f"./bert/l{l}_in", activations[f"bert.encoder.layer.{l}.attention"]["input"][0].detach().numpy())
        np.save(f"./bert/l{l}_query_weight", query_w.detach().numpy())
        np.save(f"./bert/l{l}_query_bias", query_b.detach().numpy())
        np.save(f"./bert/l{l}_query_out", query_o.detach().numpy())
        np.save(f"./bert/l{l}_key_weight", key_w.detach().numpy())
        np.save(f"./bert/l{l}_key_bias", key_b.detach().numpy())
        np.save(f"./bert/l{l}_key_out", key_o.detach().numpy())
        np.save(f"./bert/l{l}_value_weight", value_w.detach().numpy())
        np.save(f"./bert/l{l}_value_bias", value_b.detach().numpy())
        np.save(f"./bert/l{l}_value_out", value_o.detach().numpy())
        
    # calculating attention
    # should be same as np.load("/home/je/Desktop/2024_rampc/pytorch/bert/bert.encoder.layer.0.attention.output.input.npy")
        sw_attn_in = activations[f"bert.encoder.layer.{l}.attention"]["input"][0]
        sw_query_o = torch.einsum("ijk,lk-&gt;ijl", sw_attn_in, query_w) + query_b
        sw_key_o = torch.einsum("ijk,lk-&gt;ijl", sw_attn_in, key_w) + key_b
        sw_value_o = torch.einsum("ijk,lk-&gt;ijl", sw_attn_in, value_w) + value_b
        sw_attn_out = scratch_sdpa(sw_query_o, sw_key_o, sw_value_o, inputs["attention_mask"])
        F.scaled_dot_product_attention(sw_query_o.reshape(12, 128, 64), sw_key_o.reshape(12, 128, 64), sw_value_o.reshape(12, 128, 64))
        np.save(f"./bert/l{l}_attn_out", activations[f"bert.encoder.layer.{l}.attention.output"]["input"][0].detach().numpy())
</code></pre>