Attention mask is applied incorrectly

[jlwitthuhn]

I was working on my own port of nanogpt to mlx when I found this one. While cross-referencing my work with yours I think I found a problem with how the attention is applied. I am quite new to this so I might be entirely off here, let me know if this doesn't add up or is accounted for elsewhere.

This issue is specifically with how the attention mask is applied while training.

In the original nanogpt model.py we see the mask is created like this: torch.tril(torch.ones(config.block_size, config.block_size)).view(1, 1, config.block_size, config.block_size) and then later applied like this att = att.masked_fill(self.bias[:,:,:T,:T] == 0, float('-inf'))

For a small example value of T = 4, this produces an output that resembles:

tensor([[[[-0.9124,    -inf,    -inf,    -inf],
          [-0.3895, -1.7567,    -inf,    -inf],
          [-0.2186, -0.4644, -1.2544,    -inf],
          [-1.2129,  0.1759, -0.3573,  2.3461]]]])

In this repo, the mask is created with mx.tril(mx.ones([N, N])).reshape(1, 1, N, N).astype(dtype) and then applied with att = mx.where(mask[:,:,:T,:T] == 0, att, float('-1e9'))

which produces output that looks like (-1e9 replaced with -inf for clarity):

array([[[[-inf, -0.865919, 1.19201, 0.53468],
         [-inf, -inf, -1.09948, 0.285406],
         [-inf, -inf, -inf, -0.390758],
         [-inf, -inf, -inf, -inf]]]], dtype=float32)

It looks like the mask is being applied in the opposite way of what is desired. Looks easily fixable by swapping the last two arguments to mx.where. MLX selects the first of those arguments when the condition is true and the second when the condition is false.