The problem is located at torch.nn.MultiheadAttention module. Here is my model part code:
class CNNEncoder(nn.Module):
    def __init__(self, config):
        super(CNNEncoder, self).__init__()

        self.embedding = config['embedding_matrix']
        self.n_filters = config['model']['encoder']['num_layers']
        self.embedding_dim = config['model']['embed_size']
        self.filter_sizes = config['model']['encoder']['filter_size']
        self.head = config['model']['encoder']['attention_heads']
        self.dropout = config['model']['encoder']['dropout']
        # self.dropout = nn.Dropout(0.15)
        self.batch_size = config['model']['batch_size']

        self.convs = nn.ModuleList([
            nn.Conv2d(in_channels=1,
                      out_channels=self.n_filters,
                      kernel_size=(fs, self.embedding_dim))
            for fs in self.filter_sizes
        ])

        self.attn = nn.MultiheadAttention(embed_dim=self.n_filters * len(self.filter_sizes), num_heads=self.head, dropout=self.dropout) #embed_dim=300, num_heads=100
        

    def forward(self, text):
        # text = [batch size, sent len]

        embedded = self.embedding(text)

        # embedded = [batch size, sent len, emb dim]

        embedded = embedded.unsqueeze(1)

        # embedded = [batch size, 1, sent len, emb dim]

        conved = [F.relu(conv(embedded)).squeeze(3) for conv in self.convs]

        # conved_n = [batch size, n_filters, sent len - filter_sizes[n] + 1]

        pooled = [F.max_pool1d(conv, conv.shape[2]).squeeze(2) for conv in conved]

        # pooled_n = [batch size, n_filters]

        cat = torch.cat(pooled, dim=1)
        # cat = [batch size, n_filters * len(filter_sizes)]

        cat = cat.view(self.batch_size, -1, self.n_filters * len(self.filter_sizes))
        # cat = [batch size, 1, n_filters * len(filter_sizes)]

        cat = cat.permute(1,0,2) #[1, batch size, 300]

        attn_output, attn_output_weight = self.attn(query=cat, key=cat, value=cat) #k,q,v are same
        return attn_output.permute(1,0,2)

class Siamese_cnn_pair(nn.Module):
    def __init__(self, config):
        super(Siamese_cnn_pair, self).__init__()

        self.encoder = CNNEncoder(config)
        self.ling_dim = config['model']['ling_feature_size']
        # self.ling_embedding = config['ling_embedding_matrix']
        self.dropout = nn.Dropout(config['model']['encoder']['dropout'])
        self.input_dim = self.encoder.n_filters * len(self.encoder.filter_sizes)
        self.fc1 = nn.Linear(self.input_dim, 128)
        self.fc2 = nn.Linear(128,2)

    def forward(self, a1, a2):

        a = torch.cat([a1,a2], dim=1)
        if torch.cuda.is_available():
            a = a.cuda()

        output1 = self.encoder(a).squeeze(1) # batchsize * hidden_size(*2)

        output1 = self.fc1(output1)

        output1 = self.fc2(output1)  # batchsize * [c1, c2], binary classification.

        return output1

after loss.backward(), all other layer parts works normally, except encoder.attn.in_proj_weight.grad and encoder.attn.in_proj_bias.grad,
attn.in_proj_bias.grad have torch.Size([270,000]) (300 * 300 * 3(k,q,v)), but its first 180,000 are all zero, the last 90,000 have non-zero value. the situation is similar with attn.in_proj_weight.grad (torch.Size([900, 300])) it seems like only one of key, query, value matrix weight got gradient. And I don’t why it is, and how to fix it. Here is the output of attn.in_proj_weight.grad.