Training fails (sometimes) when switching from GPU to CPU

Ada: <h1><a name="p-456514-summary-1" class="anchor" href="#p-456514-summary-1"></a>Summary</h1>
<p>I am training a Pytorch implementation of Basenji (a genomics model).  Switching from training on GPU to CPU on one machine causes the model to fail to train: loss does not reduce and NaNs eventually appear.  On another machine, however, training on CPU works as expected <strong>using the same code</strong>.</p>
<p>I have generally been working on issues of reproducibility, so it is concerning that training behavior on CPU vs. GPU can be so different, and that behavior can differ so starkly between machines.  I would have expected behavior on CPU vs. GPU to be nearly identical (up to rounding errors), and certainly that CPU behavior would be the same across machines.  Does anybody know why this is not the case?</p>
<h1><a name="p-456514-detailed-description-2" class="anchor" href="#p-456514-detailed-description-2"></a>Detailed description</h1>
<p>On one machine (machine <span class="hashtag-raw">#1</span> in the figure below), training with GPU works as expected (green line, hidden behind red, in plot below).  Internally, in the code, data and models are sent to GPU using <code>.to(device)</code>, and <code>device</code> is set in an outer of code with <code>device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")</code>.  When I switch to training to CPU using <code>device = torch.device("cpu")</code>, loss does not reduce, and NaNs are introduced to the model spontaneously, and at different times for different runs (orange vs. blue line in plot below).</p>
<p><strong>This issue seems to be machine-specific, as opposed to a bug in the code.</strong>  I have another machine (machine <span class="hashtag-raw">#2</span>, red line, below) where, when loading the exact same code, training on CPU behaves the same as training on GPU with machine <span class="hashtag-raw">#1</span>.  Both machines are virtual machines through Google Cloud, so I actually detached the external drive with my code and data from one machine and attached it to the other.</p>
<p>The figure below compares performance on CPU v. GPU, and on the two machines.<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/c/3/c3501a81e83250b1824e7a5ef66f2a3b2f3ffc13.png" data-download-href="https://discuss.pytorch.org/uploads/default/c3501a81e83250b1824e7a5ef66f2a3b2f3ffc13" title="CPU_v_GPU_performance"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/c/3/c3501a81e83250b1824e7a5ef66f2a3b2f3ffc13_2_503x375.png" alt="CPU_v_GPU_performance" data-base62-sha1="rROJ7CSRk05Hf3xuCtkH3taeCDV" width="503" height="375" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/c/3/c3501a81e83250b1824e7a5ef66f2a3b2f3ffc13_2_503x375.png, https://discuss.pytorch.org/uploads/default/original/3X/c/3/c3501a81e83250b1824e7a5ef66f2a3b2f3ffc13.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/c/3/c3501a81e83250b1824e7a5ef66f2a3b2f3ffc13.png 2x" data-dominant-color="E5D8DA"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">CPU_v_GPU_performance</span><span class="informations">715×532 102 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div><br>
The orange line shows one run on CPU on machine <span class="hashtag-raw">#1</span>.  It stops when loss became NaN.  The blue line is a rerun of the exact same code, it went much longer before loss became NaN.  The green line (hidden behind the red line) is a GPU run on the same machine, and this loss/behavior is as expected.  The red line is a CPU run on machine <span class="hashtag-raw">#2</span> using the exact same code as what produced the blue and orange lines.</p>
<h1><a name="p-456514-limiting-nondeterministic-behavior-3" class="anchor" href="#p-456514-limiting-nondeterministic-behavior-3"></a>Limiting nondeterministic behavior</h1>
<p>Nondeterministic behavior is controlled by the following lines of code, which were pulled from a Pytorch implementation of another model:</p>
<pre><code class="lang-auto">torch.backends.cudnn.deterministic = True
random.seed(hash("setting random seeds") % 2**32 - 1)
np.random.seed(hash("improves reproducibility") % 2**32 - 1)
torch.manual_seed(hash("by removing stochasticity") % 2**32 - 1)
torch.cuda.manual_seed_all(hash("so runs are repeatable") % 2**32 - 1)
</code></pre>
<p>Data/batches are pulled deterministically.  I think there is still some randomness in the system because, although the blue/orange lines track, and the green/red lines track, the actual loss values differ slightly (order of ~2, which is about 0.05% of the loss value).</p>
<h1><a name="p-456514-machine-details-4" class="anchor" href="#p-456514-machine-details-4"></a>Machine details</h1>
<p>The machines are virtual machines in Google Cloud and created from public images for deep learning on Linux, available through Google Cloud.  I’m not 100% sure which Google image Machine <span class="hashtag-raw">#1</span> was built from, but my best guess is listed in the table below.  Other select package versions from the environments are also listed below - I’m not sure which ones would be most relevant to the problem here.</p>
<div class="md-table">
<table>
<thead>
<tr>
<th></th>
<th>Machine <span class="hashtag-raw">#1</span></th>
<th>Machine <span class="hashtag-raw">#2</span></th>
</tr>
</thead>
<tbody>
<tr>
<td>Machine image</td>
<td>c2-deeplearning-pytorch-1-12-cu113-v20220928-debian-10 (I think)</td>
<td>c0-deeplearning-common-cu123-v20240708-debian-11-py310</td>
</tr>
<tr>
<td>Machine image description</td>
<td>Google, Debian 10 based Deep Learning VM for PyTorch CPU/GPU with CUDA 11.3, M97, Deep Learning VM Image with PyTorch 1.12 and fast.ai preinstalled.</td>
<td>Google, Deep Learning VM with CUDA 12.3, M123, Debian 11, Python 3.10. With CUDA 12.3 preinstalled.</td>
</tr>
<tr>
<td>Python</td>
<td>3.7.12</td>
<td>3.10.14</td>
</tr>
<tr>
<td>Pytorch</td>
<td>1.12.1 (build py3.7_cuda11.3_cudnn8.3.2_0)</td>
<td>1.12.1 (build cuda112py310he33e0d6_201)</td>
</tr>
<tr>
<td>cuda-version</td>
<td>not installed</td>
<td>11.8</td>
</tr>
<tr>
<td>cudatoolkit</td>
<td>11.3.1</td>
<td>11.8.0</td>
</tr>
<tr>
<td>cudnn</td>
<td>not installed</td>
<td>8.9.7.29</td>
</tr>
</tbody>
</table>
</div><h1><a name="p-456514-code-imports-5" class="anchor" href="#p-456514-code-imports-5"></a>Code imports</h1>
<p>Here are the packages I load into the code.</p>
<pre><code class="lang-auto"># data loading
import pickle
import json

# deep learning / computation
import wandb
import torch
import torch.nn as nn
import torch.utils.data as Data
import torch.nn.functional as F
import numpy as np

# Pytorch implementation of Basenji
import sys
sys.path.insert(0,basenji2_dir)
from model import Basenji2

# useful types / datasetructures
from torch import Tensor
from typing import Any, List, Tuple
from collections import defaultdict, OrderedDict, namedtuple
import bunch
import pandas as pd

# standard library functions
import os       # get file paths
import time     # time performance
import glob
from natsort import natsorted
from bisect import bisect_right
import random
import math
import shutil

# visualizing
from tqdm import notebook as tqdm
tqdm.tqdm().pandas()
from matplotlib import pyplot as plt
import seaborn as sns
</code></pre>