Torch >= 2.2.0 inference issues on MPS

[davmacario]

When running

python sample.py --init_from=gpt2 --num_samples=2 --max_new_tokens=100

having set device = 'mps' on my M1 Pro MacBook (MacOS 14.4), with Torch 2.2.1 and 2.2.0, I get this output:

Overriding: init_from = gpt2
Overriding: num_samples = 2
Overriding: max_new_tokens = 100
loading weights from pretrained gpt: gpt2
forcing vocab_size=50257, block_size=1024, bias=True
overriding dropout rate to 0.0
number of parameters: 123.65M
No meta.pkl found, assuming GPT-2 encodings...

!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
---------------

!!!!!!。!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
---------------

The character ! corresponds to token 0, meaning that the model only generates 0's at the output.
This does not happen when using Torch 2.1.x.

I know this is probably a Torch bug, but I could use the help trying to pinpoint the actual cause of this issue (and possibly submit a bug report to Torch).
Let me know if someone has had the same issue.

[adriankobras]

I had the same issue on a M1 Pro Macbook with Torch 2.2.0

[sun1638650145]

I encountered a similar issue as well. I reproduced a Transformer on mps and encountered a similar error, but it worked fine on cpu. I upgraded to torch 2.3.0 and it seems to have fixed it(Although I didn't see any mention of bug fixes in the logs).

[davmacario]

Alright, thanks @sun1638650145, I'll give it a try and possibly update the issue!