CUDA: Out of memory error when using multi-gpu

sanchitgn: <p>Hi all,<br>
I am trying to fine-tune the <a href="https://huggingface.co/transformers/model_doc/bart.html#bartforconditionalgeneration" rel="nofollow noopener">BART model from transformers</a> for language generation on a custom dataset (30K examples of 256 length. &lt;5MB on disk).</p>
<p>I have followed the Data parallelism guide. Here are the relevant parts of my code</p>
<pre><code class="lang-auto">args.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

if args.n_gpu &gt; 1:
        model = nn.DataParallel(model)
model.to(args.device)

# Training
args.per_gpu_train_batch_size * max(1, args.n_gpu)

for step, batch in enumerate(epoch_iterator):
            model.train()
            batch = tuple(t.to(args.device) for t in batch)
</code></pre>
<p>I am facing a CUDA: Out of memory issue when using a batch size (per gpu) of 4 on 2 gpus. However training works fine on a single GPU. I am trying to train on 2 Titan-X gpus with 12GB memory.</p>
<p>This is the error message</p>
<pre><code class="lang-auto">RuntimeError: CUDA out of memory. Tried to allocate 394.00 MiB (GPU 0; 11.93 GiB total capacity; 10.84 GiB already allocated; 289.81 MiB free; 277.07 MiB cached) (malloc at /pytorch/c10/cuda/CUDACachingAllocator.cpp:267)
</code></pre>
<p>If it helps, I am using AdamW optimizer, with linear warmup.<br>
I have tried setting CUDA_VISIBLE_DEVICES, which gives the same error.</p>
<p>Am I missing something?</p>

ptrblck: <p><code>nn.DataParallel</code> will use more memory on the default device as described <a href="https://medium.com/huggingface/training-larger-batches-practical-tips-on-1-gpu-multi-gpu-distributed-setups-ec88c3e51255">here</a>.<br>
We generally recommend to use <code>nn.DistributedDataParallel</code> with a single process per GPU to get the best performance.</p>

Usman_Mahmood: <p><a class="mention" href="/u/ptrblck">@ptrblck</a> more memory compared to other gpus or more memory compared to if you were only using 1 gpu?<br>
When I run my code with 1 gpu and batch size 16, it works. But when I run same code with same batch size using 2 gpus (with equal memory) I get out of memory error, and on GPU 1 not on GPU 0, which is strange because my default device is GPU 0.<br>
My issue looks similar to one discussed here:</p><aside class="quote quote-modified" data-post="1" data-topic="19302">
  <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/r/e9a140/48.png" class="avatar">
    <a href="https://discuss.pytorch.org/t/dataparallel-memory-consumption-in-pytorch-0-4/19302">DataParallel memory consumption in PyTorch 0.4</a> 
  </div>
  <blockquote>
    I believe I’m seeing a certain loss of functionality after upgrading from PyTorch 0.3.1 to 0.4.0.  Specifically I’m trying to use nn.DataParallel to train, on two GPU’s, a model with a parameter that takes up over half the memory of either GPU.  When the DataParallel library code attempts to replicate the model over both GPU’s it broadcasts the parameters to both, and runs out of GPU memory during the broadcast operation.  In 0.3.1 the broadcast operation was implemented in Python, and contained…
  </blockquote>
</aside>


ptrblck: <p>The default device should use more memory than the other GPUs, since e.g. you would increase the batch size.<br>
Could you post your model definition so that we could take a look, please?</p>

Usman_Mahmood: <p>I made another post here</p>
<aside class="quote quote-modified" data-post="29" data-topic="22551">
  <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/usman_mahmood/48/22826_2.png" class="avatar">
    <a href="https://discuss.pytorch.org/t/dataparallel-imbalanced-memory-usage/22551/29">DataParallel imbalanced memory usage</a> 
  </div>
  <blockquote>
    I am having the same imbalance issue but the problem is that my gpu 1 not gpu 0 is going out of memory. Both gpus have 32GB of memory. With NVIDIA-SMI i see that gpu 0 is only using 6GB of memory whereas, gpu 1 goes to 32. 
I could have understood if it was other way around with gpu 0 going out of memory but this is weird. 
I only pass my model to the DataParallel so it’s using the default values. 
Also, if I use only 1 GPU, i don’t get any out of memory issues. This is also strange for me. 
Any…
  </blockquote>
</aside>

<p>Here is my encoder</p>
<pre><code>class Encoder(nn.Module):
  def __init__(self, input_channels, args):
    super().__init__()
    self.feature_size = args.fMRI_feature_size
    self.hidden_size = self.feature_size
    self.downsample = not args.no_downsample
    self.input_channels = input_channels
    self.two_d = args.fMRI_twoD
    self.end_with_relu = args.end_with_relu
    self.args = args
    init_ = lambda m: init(m,
                           nn.init.orthogonal_,
                           lambda x: nn.init.constant_(x, 0),
                           nn.init.calculate_gain('relu'))
    self.flatten = Flatten()

    if self.two_d:
        self.final_conv_size = 128 * 24 * 30
        self.final_conv_shape = (128, 24, 30)
        self.main = nn.Sequential(
            init_(nn.Conv2d(self.input_channels, 32, (9,10), stride=1)),
            nn.ReLU(),
            init_(nn.Conv2d(32, 64, (9,10), stride=1)),
            nn.ReLU(),
            init_(nn.Conv2d(64, 128, (8,9), stride=1)),
            nn.ReLU(),
            init_(nn.Conv2d(128, 128, (7,8), stride=1)),
            nn.ReLU(),
            Flatten(),
            init_(nn.Linear(self.final_conv_size, self.feature_size))
            #nn.ReLU()
        )
    else:
        self.final_conv_size = 10 * 24 * 30 * 12
        self.final_conv_shape = (10, 24, 30, 12)
        self.main = nn.Sequential(
            init_(nn.Conv3d(self.input_channels, 3, (9, 10, 4), stride=(1, 1, 1))),
            nn.ReLU(),
            init_(nn.Conv3d(3, 5, (9, 10, 3), stride=(1, 1, 1))),
            nn.ReLU(),
            init_(nn.Conv3d(5, 8, (8, 9, 3), stride=(1, 1, 1))),
            nn.ReLU(),
            init_(nn.Conv3d(8, 10, (7, 8, 2), stride=(1, 1, 1))),
            nn.ReLU(),
            Flatten(),
            init_(nn.Linear(self.final_conv_size, self.feature_size)),

            #nn.ReLU()
        )
    self.train()

def forward(self, inputs, fmaps=False):
    f5 = self.main[:6](inputs)
    f7 = self.main[6:8](f5)
    out = self.main[8:](f7)
    if self.end_with_relu:
        assert self.args.method != "vae", "can't end with relu and use vae!"
        out = F.relu(out)
    if fmaps:
        return {
            'f5': f5.permute(0, 2, 3, 1),
            'f7': f7.permute(0, 2, 3, 1),
            'out': out
        }
    return out
</code></pre>

Usman_Mahmood: <p>I was using slurm on server and I found out that<br>
gpu 0 in nvidia smi != gpu :0 in the codeI made following two changes</p>
<p>device = torch.device(“cuda:” + “1”)<br>
encoder = torch.nn.DataParallel(encoder, device_ids=[1, 0])</p>
<p>and now the error says gpu 0 is out of memory<br>
so 0 in code is actually 1 , and vice versa.</p>
<p>so, the only issue remains that why using 2 gpus with DataParallel goes out of memory while using 1 gpu with same data and batch size doesn’t<br>
Does doing DataParallel requires extra memory to manage parallelism ?</p>

ptrblck: <p><code>nn.DataParallel</code> might create an overhead, thus we recommend to use DistributedDataParallel with a single process for each device.<br>
This would be the fastest approach without creating the imbalanced memory usage.</p>

Olivier-CR: <p>could the PyTorch team fix it then? it’s important to have a single-process multi-device solution, especially as multi-GPU nodes get bigger and bigger, it’s less and less useful to do multi-node training. We need something simpler than DDP please</p>

ptrblck: <p><code>DDP</code> is not only used for multi-node training, but is also speeding up single-node multi-GPU workloads.<br>
The current proposal is to <a href="https://github.com/pytorch/pytorch/issues/65936">deprecate DataParllel</a> and in this sense to ramp up the documentation on DDP.</p>

Olivier-CR: <p>yes I know! I’m just suggesting, since single-node multi-device training is the entry point to dist training, and will likely be enough for most users, why having it done with DDP (which is so <a href="https://discuss.pytorch.org/u/olivier-cr/activity/topics">insanely complex</a> to get working), while torch.nn.DataParallel API is so compact and friendly. I know that DataParallel in its current form <a href="https://discuss.pytorch.org/u/olivier-cr/activity/topics">performs badly</a>, but looking at things from the customer perspective, this is the experience the Pytorch team should aim for for DDP future. No torchrun, no MPI, no processes to spawn manually etc. It’s better to have a single-code, single process experience, like <a href="https://mxnet.apache.org/versions/1.5.0/tutorials/gluon/multi_gpu.html" rel="noopener nofollow ugc">MXNet Gluon offers.</a>. Again, I’m not saying it’s possible today in Pytorch, I’m just saying this is what you should aim for, so that the data-parallelism experience is better. Think about how easy it is to write a single-file PySpark data processing code and run it over multiple machines. Why things aren’t this simpler in the DL training world yet? We’re clearly lagging behind the big data world in terms of abstracting complexity</p>

Jackie_Li: <p>DDP can solve the imbalanced memory use, but it will use more memory in each GPU for communication, right?<br>
In the case with large sample and large model (e.g. I can only use batch_size=1), is there any suggestion for parallel training?<br>
Thanks a lot!</p>

ptrblck: <aside class="quote no-group" data-username="Jackie_Li" data-post="11" data-topic="72333">
<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/jackie_li/48/45177_2.png" class="avatar"> Jackie_Li:</div>
<blockquote>
<p>but it will use more memory in each GPU for communication, right?</p>
</blockquote>
</aside>
<p>No, I don’t think additional data is needed as the gradients are calculated in any case. Why would this be the case?</p>

Jackie_Li: <p>Because when I ran the same program with DDP, it comsumes more GPU memory and threw an OOM error. Without DDP, the memory usage is about 18000M, but with DDP 24000M is not enough</p>

ptrblck: <p>That sounds like a lot of overhead, but might also depend on the actual model or your setup.<br>
I’ve checked the memory overhead used for the buckets and see that the memory increases approx. from 13.4GB to 14.1GB for a resnet so unsure where your ~33% overhead is coming from</p>

Jackie_Li: <p>I checked with some code from others work, the problem is caused by not setting “torch.cuda.set_device()”. After setting this, DDP works fine.</p>
<p>Thanks a lot for your help!</p>