Device = "mps" is producing Nan weights in nn.Embedding

FALL-ML: <p>I have an NLP model that trains fine in the following contexts:</p>
<ul>
<li>Windows11 + CPU</li>
<li>Windows11 + CUDA</li>
<li>Ubuntu20.04 + CPU</li>
<li>Ubuntu20.04 + CUDA</li>
<li>macOS12.5 + CPU</li>
</ul>
<p>However, my attempts to run the same model using “mps” as the device are resulting in unexpected behavior: the nn.Embedding layers in my model are being initialized but then the weights quickly train to Nan values. There is no specific error that occurs - the loss just never improves. Since the model trains fine if I simply change my device to “cpu”, I believe this is likely an issue with my virtual environment setup. Does anyone have insight on what I am doing wrong?</p>
<hr>
<p><strong>Environment Setup Information</strong><br>
To setup a M1 compatible environment I used the following commands:</p>
<blockquote>
<p>CONDA_SUBDIR=osx-arm64 conda create -n test_environment python=3.9 -c conda-forge</p>
</blockquote>
<blockquote>
<p>conda env config vars set CONDA_SUBDIR=osx-arm64</p>
</blockquote>
<blockquote>
<p>pip3 install torch torchvision torchaudio</p>
</blockquote>
<hr>
<p><strong>Output from Simple Sanity Checks</strong><br>
These sanity checks seems to suggest the environment should work:</p>
<pre><code class="lang-auto">import torch
import platform

print(f'Platform: {platform.platform()}')
print(f'torch.has_mps: {torch.has_mps}')
print(f'MPS is available: {torch.backends.mps.is_available()}')
print(f'Pytorch was built with MPS: {torch.backends.mps.is_built()}')
</code></pre>
<blockquote>
<p>Platform: macOS-12.5-arm64-arm-64bit<br>
torch.has_mps: True<br>
MPS is available: True<br>
Pytorch was built with MPS: True</p>
</blockquote>

FALL-ML: <p>Here is some additional debugging context for this problem. I’m arbitrarily examining the weights during the 5th mini-batch in my first epoch.</p>
<p><strong>With device = “cpu”.</strong><br>
<code>self.pos_embedding.weight</code> produces:</p>
<blockquote>
<p>Parameter containing:<br>
tensor([[ 0.0549, -0.0284, -0.1278,  …, -0.0946, -0.0928, -0.0088],<br>
[ 0.0771, -0.0170, -0.0525,  …,  0.1236,  0.0622,  0.0237],<br>
[-0.0847, -0.0200, -0.1097,  …, -0.0541,  0.0832,  0.0280],<br>
…,<br>
[ 0.1198, -0.0781, -0.0658,  …, -0.0878,  0.0403,  0.1020],<br>
[ 0.1255,  0.0605,  0.0063,  …, -0.0283, -0.0824, -0.1120],<br>
[ 0.0942, -0.0514, -0.0760,  …, -0.0519,  0.0670, -0.0561]],<br>
requires_grad=True)</p>
</blockquote>
<p><strong>With device = “mps”.</strong><br>
<code>self.pos_embedding.weight</code> produces:</p>
<blockquote>
<p>Parameter containing:<br>
tensor([[    nan,     nan,     nan,  …,     nan,     nan,     nan],<br>
[    nan,     nan,     nan,  …,     nan,     nan,     nan],<br>
[    nan,     nan,     nan,  …,     nan,     nan,     nan],<br>
…,<br>
[-0.1200, -0.0589, -0.0015,  …, -0.0278,  0.1136,  0.0352],<br>
[-0.0968,  0.0428, -0.0388,  …, -0.1297, -0.0833, -0.1068],<br>
[ 0.1078, -0.0055,  0.0330,  …, -0.0212, -0.0376,  0.1046]],<br>
device=‘mps:0’, requires_grad=True)</p>
</blockquote>

philipturner: <p>Please look: <a href="https://github.com/pytorch/pytorch/issues/83015" class="inline-onebox" rel="noopener nofollow ugc">Incorrect tensor conversion to m1 MPS. · Issue #83015 · pytorch/pytorch · GitHub</a></p>

FALL-ML: <p>Thank you for the response. Unfortunately, all my .to(device) calls are already using the default non_blocking = False.</p>
<p>I’ll play around a bit and see if I can extract a minimal example of the error I’m getting.</p>

pedram: <p><a class="mention" href="/u/fall-ml">@FALL-ML</a> any luck in resolving this issue?</p>

NightMachinery: <p>I also have this issue. Using <code>model().to(device=device, non_blocking=False)</code> did not solve it.</p>