Interactions between Mixed Precision Training and Memory when using CUDA

TheeNinja: <p>My understanding of mixed precision training is that there is a tensor of master weights that is FP32. Each iteration of training, a local copy of FP16 weights is made, which is used to do the forward-propagation and back-propagation.</p>
<p>This is the first interaction I’ve got a question about: If all the propagation and matrix multiplications (done on the GPU) only use this local copy of FP16 weights, does that mean that the FP32 master weights can be stored in CPU memory, with the FP16 weights being a copy from CPU memory -&gt; GPU memory (that also halves precision) performed each iteration? I believe this is an implementation detail in PyTorch (all I do on my end is set <code>precision=16</code> on the <code>Trainer</code> along with an <code>amp_level</code>), but an answer would help me understand the pipeline as a whole.</p>
<p>Moving from weights-side of things to the data side: I’ve defined this dataset:</p>
<pre><code class="lang-auto">class InMemoryLabeledAudioDataset(LabeledAudioDataset):
    def __init__(self, labeled_journal_entry_ids, feature_len, max_steps, model_category):
        super().__init__(labeled_journal_entry_ids, feature_len, max_steps)

        self.model_category = model_category
        expected_y_len = SIZED_SCORES[model_category]

        self.x_batch = torch.empty((len(self), self.feature_len, self.max_steps))
        self.y_batch = torch.empty((len(self), expected_y_len))

        for index in range(len(self)):
            self.x_batch[index], self.y_batch[index] = super().__getitem__(index)

    def __getitem__(self, index):
        return self.x_batch[index], self.y_batch[index]
</code></pre>
<p><code>LabeledAudioDataset</code> (specifically its <code>__getitem__</code>) reads audio from disk, that’s just abstracted away. All the data I have, as F32, can all fit in CPU memory, but not GPU memory. If all data is as F16, it can fit in both CPU memory and GPU memory. I can think of 3 different ways to go about feeding this data into my network:</p>
<ol>
<li>
<p>Store all data in CPU memory as F32, perform half-precision one-time copy to GPU memory. All mini-batches of data are now in GPU memory as F16, and can be read from each iteration.</p>
</li>
<li>
<p>Store all data in GPU memory as F16 (and no instance of the data will exist as F32, CPU memory will be freed).</p>
</li>
<li>
<p>Store all data in CPU memory as F32, perform full-precision local copy each iteration to GPU memory. Now each mini-batch of data will have be copied from CPU as F32, but because not all F32 mini-batches can exist in GPU memory at the same time, this local copy would have to be done each iteration.</p>
</li>
</ol>
<p>This is the second interaction I’ve got a question about: which of these options is the reality <code>DataLoader</code> executes (could also be an option outside of the 3), and how explicit am I supposed to be in <code>LabeledAudioDataset</code>? By explicit, I mean whether I should be calling <code>half()</code> on <code>self.x_batch</code> and <code>self.y_batch</code> (even though my <code>Trainer</code> is already with <code>precision=16</code>), and whether I should call <code>to(my_cuda_device)</code> or not even though my <code>Trainer</code> was provided with <code>gpus=1</code>. In other words, what will be taken care of of vs. what am I supposed to do on my end explicitly?</p>
<p>I know that the goal of mixed precision training is to only have some specific tensors as F16, but is data one of those tensors? Will options 1 and 2 make the network perform sub-optimally loss-wise compared to option 3 because of the data’s lower precision? Finally, how do the  <code>pin_memory</code> and <code>num_workers</code> <code>DataLoader</code> arguments change up things here?</p>
<p>The questions I have are all over the place, but I’ve read discussions about <a href="https://discuss.pytorch.org/t/pin-memory-vs-sending-direct-to-gpu-from-dataset/33891">how to store data in GPU memory</a> and what is done in <code>Dataset</code> vs. <code>DataLoader</code>, and haven’t been able to piece together the big picture.</p>

ptrblck: <aside class="quote no-group" data-username="TheeNinja" data-post="1" data-topic="79629">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/theeninja/48/23751_2.png" class="avatar"> TheeNinja:</div>
<blockquote>
<p>My understanding of mixed precision training is that there is a tensor of master weights that is FP32. Each iteration of training, a local copy of FP16 weights is made, which is used to do the forward-propagation and back-propagation.</p>
</blockquote>
</aside>
<p>Master parameters and gradients are used in <code>apex/amp</code> with <code>opt_level='O2'</code>.<br>
The current PyTorch master (and the nightly binaries) contain the native mixed-precision implementation, which we recommend to use now, and you can have a look at the documentation <a href="https://pytorch.org/docs/master/amp.html">here</a>.</p>
<aside class="quote no-group quote-modified" data-username="TheeNinja" data-post="1" data-topic="79629">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/theeninja/48/23751_2.png" class="avatar"> TheeNinja:</div>
<blockquote>
<p>If all the propagation and matrix multiplications (done on the GPU) only use this local copy of FP16 weights, does that mean that the FP32 master weights can be stored in CPU memory, with the FP16 weights being a copy from CPU memory → GPU memory (that also halves precision) performed each iteration?</p>
</blockquote>
</aside>
<p>While that might be possible, you would most likely lose the performance benefits from using TensorCores due to the data transfer, so the master params are stored on the GPU.</p>
<aside class="quote no-group" data-username="TheeNinja" data-post="1" data-topic="79629">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/theeninja/48/23751_2.png" class="avatar"> TheeNinja:</div>
<blockquote>
<p>By explicit, I mean whether I should be calling <code>half()</code> on <code>self.x_batch</code> and <code>self.y_batch</code> (even though my <code>Trainer</code> is already with <code>precision=16</code> ), and whether I should call <code>to(my_cuda_device)</code> or not even though my <code>Trainer</code> was provided with <code>gpus=1</code> . In other words, what will be taken care of of vs. what am I supposed to do on my end explicitly?</p>
</blockquote>
</aside>
<p>You should not call <code>half()</code> on the data or the model and <code>apex/amp</code> or <code>torch.cuda.amp.autocast</code> take care of the casting for you.<br>
I’m not sure where the <code>Trainer</code> class comes from, but I assume it’s from a high-level API?<br>
If so, I would expect the same behavior, i.e. no <code>half()</code> calls.</p>
<aside class="quote no-group" data-username="TheeNinja" data-post="1" data-topic="79629">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/theeninja/48/23751_2.png" class="avatar"> TheeNinja:</div>
<blockquote>
<p>Finally, how do the <code>pin_memory</code> and <code>num_workers</code> <code>DataLoader</code> arguments change up things here?</p>
</blockquote>
</aside>
<p>These arguments to the <code>DataLoader</code> are completely independent from mixed-precision training.<br>
<code>num_workers</code> will use multiprocessing, where each worker loads a batch in the background while your training is running.<br>
<code>pin_memory=True</code> will use page-locked memory to speed up the data transfer between the host and device and allows you to use <code>tensor.to(non_blocking=True)</code> to an asynchronous transfer.</p>

TheeNinja: <p>Thank you for the great answer, whole pipeline is a lot clearer now.</p>
<aside class="quote no-group" data-username="ptrblck" data-post="2" data-topic="79629">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/ptrblck/48/1823_2.png" class="avatar"> ptrblck:</div>
<blockquote>
<p>You should not call <code>half()</code> on the data or the model and <code>apex/amp</code> or <code>torch.cuda.amp.autocast</code> take care of the casting for you.<br>
I’m not sure where the <code>Trainer</code> class comes from, but I assume it’s from a high-level API?<br>
If so, I would expect the same behavior, i.e. no <code>half()</code> calls.</p>
</blockquote>
</aside>
<p>My bad, <code>Trainer</code> comes from PyTorch Lightning, but still sounds like casting will be taken care of because it uses <code>apex/amp</code> under the hood.</p>
<p>I have one followup question:</p>
<p>All my data when in FP16 can fit entirely on the GPU, which to my understanding would be more performant compared to transferring data from host to device even with <code>pin_memory=True</code> and multiple workers.</p>
<p>The problem is that I can never store all the data on the GPU in FP32, even intermediately, because it doesn’t fit. I need to have the <code>half()</code> call applied before the <code>to(cuda_device)</code> / <code>cuda()</code> call. Is there a way for me to move the data to GPU after the <code>half()</code> call that is done for me by <code>amp/apex</code> or <code>torch.cuda.amp</code>? In other words, is there a way for the <code>half()</code> to be done on the CPU tensor for the entire data batch, and have the tensor copied to GPU one-time after?</p>

ptrblck: <aside class="quote no-group" data-username="TheeNinja" data-post="3" data-topic="79629">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/theeninja/48/23751_2.png" class="avatar"> TheeNinja:</div>
<blockquote>
<p>In other words, is there a way for the <code>half()</code> to be done on the CPU tensor for the entire data batch, and have the tensor copied to GPU one-time after?</p>
</blockquote>
</aside>
<p>You could transform the data on the CPU to <code>half</code>, push them to the device, and transform the batch back to <code>float</code> before calling the forward function.<br>
Note that your training also needs to store intermediate tensors to be able to calculate the gradients, but I assume you’ve already checked it.</p>