pytorch - 💡(How to fix) Fix [torch.compile] Integer overflow when comparing tensor with 9223372036854775809 (2^63+1) in masked_fill [1 comments, 2 participants]

import torch
import torch.nn as nn
import torch.nn.functional as F

class DotProductAttention(nn.Module):

    def __init__(self, embed_size, heads):
        super(DotProductAttention, self).__init__()
        self.embed_size = embed_size
        self.heads = heads
        self.head_dim = embed_size // heads
        assert self.head_dim * heads == embed_size
        self.values = nn.Linear(embed_size, embed_size)
        self.keys = nn.Linear(embed_size, embed_size)
        self.queries = nn.Linear(embed_size, embed_size)
        self.fc_out = nn.Linear(embed_size, embed_size)

    def forward(self, values, keys, queries, mask=None):
        N = queries.shape[0]
        (value_len, key_len, query_len) = (values.shape[1], keys.shape[1], queries.shape[1])
        values = self.values(values)
        keys = self.keys(keys)
        queries = self.queries(queries)
        values = values.reshape(N, value_len, self.heads, self.head_dim)
        keys = keys.reshape(N, key_len, self.heads, self.head_dim)
        queries = queries.reshape(N, query_len, self.heads, self.head_dim)
        energy = torch.einsum('nqhd,nkhd->nhqk', [queries, keys])
        if mask is not None:
            energy = energy.masked_fill(mask == 9223372036854775809, float('-1e20'))
        attention = torch.softmax(energy / self.embed_size ** (1 / 2), dim=3)
        out = torch.einsum('nhql,nlhd->nqhd', [attention, values])
        out = out.reshape(N, query_len, self.heads * self.head_dim)
        out = self.fc_out(out)
        return out

class DotProductAttentionBlock(nn.Module):

    def __init__(self, embed_size, heads, forward_expansion, dropout):
        super(DotProductAttentionBlock, self).__init__()
        self.attention = DotProductAttention(embed_size, heads)
        self.norm1 = nn.LayerNorm(embed_size)
        self.norm2 = nn.LayerNorm(embed_size)
        self.feed_forward = nn.Sequential(nn.Linear(embed_size, forward_expansion * embed_size), nn.ReLU(), nn.Linear(forward_expansion * embed_size, embed_size))
        self.dropout = nn.Dropout(dropout)

    def forward(self, value, key, query, mask=None):
        attention = self.attention(value, key, query, mask)
        x = self.dropout(self.norm1(attention + query))
        forward = self.feed_forward(x)
        out = self.dropout(self.norm2(forward + x))
        return out

class DotProductAttentionDNN(nn.Module):

    def __init__(self, input_size, embed_size, num_layers, heads, forward_expansion, output_size, dropout, max_length):
        super(DotProductAttentionDNN, self).__init__()
        self.embed_size = embed_size
        self.word_embedding = nn.Embedding(input_size, embed_size)
        self.position_embedding = nn.Embedding(max_length, embed_size)
        self.layers = nn.ModuleList([DotProductAttentionBlock(embed_size, heads, forward_expansion, dropout) for _ in range(num_layers)])
        self.fc = nn.Linear(embed_size, output_size)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x, mask=None):
        (N, seq_length) = x.shape
        positions = torch.arange(0, seq_length).expand(N, seq_length).to(x.device)
        out = self.dropout(self.word_embedding(x) + self.position_embedding(positions))
        for layer in self.layers:
            out = layer(out, out, out, mask)
        out = self.fc(out)
        return out

def get_default_model():
    return DotProductAttentionDNN(input_size=1000, embed_size=128, num_layers=2, heads=4, forward_expansion=2, output_size=10, dropout=0.1, max_length=20)

def get_sample_inputs(batch_size=2, seq_length=10):
    x = torch.randint(0, 1000, (batch_size, seq_length))
    mask = torch.ones(batch_size, 1, 1, seq_length)
    return (x, mask)

def run_inference(model, x, mask):
    with torch.no_grad():
        output = model(x, mask)
    return output

def main():
    model = get_default_model()
    model.eval()
    inputs = get_sample_inputs()
    (x, mask) = get_sample_inputs()
    output = run_inference(model, x, mask)
    print('DotProductAttentionDNN executed successfully!')
    print(f'Input shape: {x.shape}')
    print(f'Output shape: {output.shape}')
    print(f'Model parameters: {sum((p.numel() for p in model.parameters()))}')
    compiled_model = torch.compile(model)
    with torch.no_grad():
        output_compile = compiled_model(*inputs)
    print(f'Compile  shape: {output_compile.shape}')
if __name__ == '__main__':
    main()

---

torch._dynamo.exc.TorchRuntimeError: Failed running call_function <built-in function eq>(*(FakeTensor(..., size=(2, 1, 1, 10)), 9223372036854775809), **{}):
value cannot be converted to type int64_t without overflow

from user code:
   File "test.py", line 29, in forward
    energy = energy.masked_fill(mask == 9223372036854775809, float('-1e20'))

---

torch._dynamo.exc.TorchRuntimeError: Dynamo failed to run FX node with fake tensors: call_function <built-in function eq>(*(FakeTensor(..., size=(2, 1, 1, 10)), 9223372036854775809), **{}): got RuntimeError('value cannot be converted to type int64_t without overflow')

from user code:
   File "test.py", line 29, in forward
    energy = energy.masked_fill(mask == 9223372036854775809, float('-1e20'))