Adam optimizer.step CUDA OOM

username: djseivad
cooked: <p>What I know about the problem</p>
<ol>
<li>Adam is stateful and requires a memory space proportional to the parameters in your model.</li>
<li>Model parameters must be loaded onto device 0</li>
<li>OOM occurs at state[‘exp_avg_sq’] = torch.zeros_like(p.data) which seems to be the last allocation of memory in the optimizer source code.</li>
<li>Neither manually allocating or use of nn.DataParallel prevents OOM error</li>
<li>Moved loss to forward function to reduce memory in training</li>
</ol>
<p>Below are my training and forward methods</p>
<pre><code class="lang-auto">def train(dataloader, vocabulary_dict, epoch):
   
    model = ViralClassification(len(vocabulary_dict), 0.5, 6588 )#, device_ids=[0,1], output_device=1)
    model.to('cuda:0')
    model.print_gpu_memory_info()
    
    optimizer = optim.Adam(model.parameters(), lr=0.01)
    
    cudnn.benchmark = True
    cudnn.enabled = True


    for i in range(epoch):

        running_loss = 0.0

        for (i, (label, sequence)) in enumerate(dataloader):
            
            loss = model(sequence.to('cuda:0'), label.to('cuda:1'))#.to('cuda:0'))
            running_loss += loss.item()            
            loss.backward()
            del loss            
           
            print('gpu_memory_one in training_loop') 
            model.print_gpu_memory_info()
            torch.cuda.empty_cache()

            print('gpu_memory_two in training_loop')
            model.print_gpu_memory_info()
            
            print('bout to step')
            optimizer.step()
            
            
            optimizer.zero_grad() 
            print('bottom of training loop')
</code></pre>
<pre><code class="lang-auto">def forward(self, inputs, labels):

        inputs = self.embedding(inputs)#.to('cuda:0')
        #print('embedded completed')

        (inputs, hidden_state) = self.bilstm_layer(inputs)
        #print('bilstm completed')
        inputs.to('cuda:1')
        
        self.bilstm_layer.flatten_parameters()
        
        torch.cuda.synchronize('cuda:0')
        torch.cuda.synchronize('cuda:1') 
        
        inputs = self.attention_layer(inputs)
        #print('attention completed')

        inputs = inputs.view(-1, self.row*2*self.lstm_dim)#.to('cuda:1')
        #print('view transformation completed')
        
        inputs = self.mlp_one(inputs, self.relu_one)
        #print('mlp one completed')

        inputs = self.mlp_two(inputs, self.relu_two)
        #print('mlp two completed')

        torch.cuda.synchronize('cuda:0')
        torch.cuda.synchronize('cuda:1')
        
        logits = self._classify(inputs)
        #print('logits completed')

        torch.cuda.empty_cache()
        torch.cuda.synchronize('cuda:0')
        torch.cuda.synchronize('cuda:1')
        #self.print_gpu_memory_info() 

        loss = self.criterion(logits, labels)
        
        return loss
</code></pre>
<p>My OOM occurs when I perform optimizer.step.</p>
<p>My problem is that before optimizer.step my memory on device 1 has plenty of open room but since the optimizer performs it calculations on device 0, the OOM occurs.</p>
<p>Is this a problem that checkpointing may be able to solve?<br>
Is it possible to change the location of the optimizer?</p>

username: ptrblck
cooked: <p>I’m not sure, how your model works.<br>
Your <code>forward</code> method seems to use model sharding, i.e. different parts of the model are located on different devices.<br>
However, you are never transferring the <code>inputs</code> to GPU1, since you are not reassigning it in this line of code:</p>
<pre><code class="lang-python">inputs.to('cuda:1')
</code></pre>
<p>Also, based on your code you are only calling <code>model.to('cuda:0')</code>, which would mean everything runs on GPU0.<br>
However, it seems that at least your <code>labels</code> are on GPU1, which should create an error in your loss calculation.</p>
<p>Are code parts missing or am I missing something? <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=9" title=":slight_smile:" class="emoji" alt=":slight_smile:"></p>

username: djseivad
cooked: <p>Don’t think you missed anything. Has given me a way forward to solving my problem. Thanks alot for the help.</p>
<p>So the reassignment of inputs acts like a pointer to a device memory space containing the tensors?</p>
<p>model.to(‘cuda:0’) will override explicit statements within the model itself?</p>
<p>Also the reason I am passing my labels to device 1 is because in theory (if I had reassigned) the logits should be output on device 1. I did this because I thought it would reduce memory use on device 0 allowing more room for the optimizer.</p>
<p>The model is for genomic sequence data and can be found at <a href="https://www.biorxiv.org/content/10.1101/694851v1" rel="nofollow noopener">https://www.biorxiv.org/content/10.1101/694851v1</a>. if your interested.</p>
<p>E: I think I pinpointed the error to when the optimizer attempts to step over the embedding layer it exceeds memory because of the size of my vocabulary. When I move it to device 1 I get a OOM on device 1 while the optimizer params stay on device 0. If I move the embedding to cpu. The model trains.</p>

username: ptrblck
cooked: <aside class="quote no-group" data-username="djseivad" data-post="3" data-topic="53198">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/d/71c47a/48.png" class="avatar"> djseivad:</div>
<blockquote>
<p>So the reassignment of inputs acts like a pointer to a device memory space containing the tensors?</p>
</blockquote>
</aside>
<p>Tensors return a copy on the specified device, while <code>model.to()</code> works recursively on all submodules.</p>
<p>How large is your embedding layer?</p>

username: djseivad
cooked: <p>Vocabulary dimension is 8,390,658.<br>
Embedding dimension is 100.</p>
<p>Embedding is supposed to represent all 12 char strings and their reverse complement with a nucleotide alphabet.</p>