# DataParallel imbalanced memory usage

**URL:** <https://discuss.pytorch.org/t/dataparallel-imbalanced-memory-usage/22551>\
**Category:** Uncategorized\
**Created:** [August 6, 2018, 10:29pm UTC](https://discuss.pytorch.org/t/dataparallel-imbalanced-memory-usage/22551 "2018-08-06T22:29:49Z")\
**Posts on this page:** 20\
**Page:** 1

<div class="post-metadata">

**Author:** ![JuanFMontesinos](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/juanfmontesinos/32/76115_2.png) [@JuanFMontesinos](https://discuss.pytorch.org/u/JuanFMontesinos)\
**Post date:** [August 6, 2018, 10:29pm UTC](https://discuss.pytorch.org/t/dataparallel-imbalanced-memory-usage/22551/1 "2018-08-06T22:29:49Z")

</div>

Hi there,  
I’m going to re-edit the whole thread to introduce a unlikely behavior with `DataParallel`  
Right now there are several recent posts about this topic and I would like to summarize the problem.

> [@CUDA error: out of memory - huge embedding layer](https://discuss.pytorch.org/t/cuda-error-out-of-memory-huge-embedding-layer/22556):
>
> I am working on a CNN autoencoder for text with an architecture and loss function like this: [screen shot 2018-08-06 at 4 32 25 pm] def computeCrossEntropy(log\_prob, target): loss = [F.nll\_loss(sentence\_emb\_matrix, word\_ids, size\_average=False) for sentence\_emb\_matrix, word\_ids in zip(log\_prob, target)] average\_loss = sum([torch.sum(l) for l in loss]) / log\_prob.size()[0] return average\_loss I understand the embedding is huge and I can only have a batch size of 8. Otherwise, I wou…

> [@DataParallel not effective](https://discuss.pytorch.org/t/dataparallel-not-effective/18221):
>
> Hi all, I have a quite large model and need to do data parallel among multiple GPUs. I used: model = nn.DataParallel(model) And there are three visible GPUs. The GPU usage is: +-----------------------------------------------------------------------------+ | Processes: GPU Memory | | GPU PID Type Process name Usage | |=============================================================================|…

> [@Imbalanced GPU memory usage training LSTM](https://discuss.pytorch.org/t/imbalanced-gpu-memory-usage-training-lstm/15049):
>
> I’m training a language model using the code here: [https://github.com/pytorch/examples/blob/master/word\_language\_model/main.py](https://github.com/pytorch/examples/blob/master/word_language_model/main.py) I have made some slight changes so that the model can be trained across multiple GPUs. However, the GPU memory usage is extremely imbalanced. I can understand that one GPU is set to gather and store all outputs. I wonder if there is any way I can balance the memory usage? Or can I set one GPU for gathering outputs and the rest for training on batc…

Right now it seems there is an imbalaced usage of GPUs when calling DataParallel.  
[](https://discuss.pytorch.org/uploads/default/optimized/2X/5/5aba409800ca4813fc11eead25b891e4befe3dd8_1_690x119.png)

From my experience and other users’ explanations I will explain why this happens:  
Using DataParallel there are 4 types of data to consider:

1. Input
2. Output
3. Ground-truth
4. Optimizer Parameters

DataParallel has a main GPU, which is the GPU where the model is stored.  
As @Yuzhou_Song indicates in another post, DataParallel splits the batch in as many GPUs as choosen, copy the model in each of them, compute the forward pass independently and then **collect back to one GPU the outputs of each GPU to calculate the loss instead of computing loss independently in each GPU**. This is the main cause of imbalance memory usage. Consider that ground-truth and output (target of loss) must be in the same GPU

I discovered by my self that some optimizers requires lot of memory to save their parameters. However all **these parameters are located in the already mentioned main GPU**. This makes the problem worse.

There is a last reason. Model inputs are usually allocated to GPU by using `.cuda()`, which usually points to the main GPU and generates more imbalance.

There are some ways to minimize this:  
DataParallel have 2 arguments, `device_ids` which allows to choose in which GPUs the model will be trained out of all available GPUs (`CUDA_VISIBLE_DEVICES`) and `output_device` which allows to choose in which GPU output will be stored.

Let’s imagine we have 3 devices, [0,1,2].  
calling  
`model = DataParallel(model).cuda()` would set device0 as main gpu. Then, optimizers parameters will be stored here.  
calling `model = DataParallel(model,output_device=1).cuda()` and `grountruth.cuda(1)` will collect all the outputs and compute loss in cuda:1  
lastly, you can allocate inputs to cuda2.  
This way the memory usage is distributed as much as possible.

Is there a way of solving this problem? @smth @ptrblck @albanD@soulless  
I guess this behavior is very inconvenient.

---

<div class="post-metadata">

**Author:** ![soulless](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/soulless/32/40008_2.png) [@soulless](https://discuss.pytorch.org/u/soulless)\
**Post date:** [August 8, 2018, 4:47am UTC](https://discuss.pytorch.org/t/dataparallel-imbalanced-memory-usage/22551/2 "2018-08-08T04:47:24Z")

</div>

I am having a similar issue. Could this be because the loss calculation is not done in the forward function?

> [@Imbalanced GPU memory usage training LSTM](https://discuss.pytorch.org/t/imbalanced-gpu-memory-usage-training-lstm/15049/2):
>
> I found the reason, it is because we collect the output back to one gpu and calculate loss there. If move loss calculation into model.forward(), the problem is resolved.

---

<div class="post-metadata">

**Author:** ![JuanFMontesinos](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/juanfmontesinos/32/76115_2.png) [@JuanFMontesinos](https://discuss.pytorch.org/u/JuanFMontesinos)\
**Post date:** [August 8, 2018, 3:31pm UTC](https://discuss.pytorch.org/t/dataparallel-imbalanced-memory-usage/22551/3 "2018-08-08T15:31:32Z")

</div>

Well the op you cited is right, in the end pytorch collect the whole batch output in one gpu to calculate the loss. That’s what generates the imbalanced usage. At the same time optimizers’ parameters are stored in by-default gpu, what makes the problem worse. There is no apparent solution since dataparallel allows you to choose output gpu but the problem remains. There are several open threads asking about the same issue. I tried to play with dataparallel arguments: device\_ids and output\_device but no way.

---

<div class="post-metadata">

**Author:** ![albanD](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/alband/32/215_2.png) [@albanD](https://discuss.pytorch.org/u/albanD)\
**Post date:** [August 8, 2018, 4:30pm UTC](https://discuss.pytorch.org/t/dataparallel-imbalanced-memory-usage/22551/4 "2018-08-08T16:30:49Z")

</div>

If you want the loss to be splitted among gpus, just make your loss layer part of the DataParallel and add a `sum` or `mean` operation on what you get out of it. That way if you use DataParallel on 4 devices, only 4 extra numbers will be allocated on the output\_device. Is that a good solution for you?

---

<div class="post-metadata">

**Author:** ![JuanFMontesinos](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/juanfmontesinos/32/76115_2.png) [@JuanFMontesinos](https://discuss.pytorch.org/u/JuanFMontesinos)\
**Post date:** [August 8, 2018, 4:48pm UTC](https://discuss.pytorch.org/t/dataparallel-imbalanced-memory-usage/22551/5 "2018-08-08T16:48:39Z")

</div>

Thanks it seems a good solution.  
If i call loss=DataParallel(loss).cuda() I would get an array of losses?

Apart of that, is it possible to split optimizers among all gpus?

---

<div class="post-metadata">

**Author:** ![albanD](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/alband/32/215_2.png) [@albanD](https://discuss.pytorch.org/u/albanD)\
**Post date:** [August 8, 2018, 4:59pm UTC](https://discuss.pytorch.org/t/dataparallel-imbalanced-memory-usage/22551/6 "2018-08-08T16:59:08Z")

</div>

You would need to have the following:

```auto
class FullModel(nn.Module):
  def __init__ (self, model, loss):
    self.model = model
    self.loss = loss

  def forward(self, inputs, targets):
    outputs = self.model(inputs)
    loss = self.loss(outputs, targets)

full_model = DataParallel(FullModel(model, loss), device_ids=[0, 1, 2], ***)
full_model.cuda()

loss = full_model(fake_input, fake_target)
print(loss.size) # returns [3] (If you loss returns a 0-dimensional tensor containing the loss value)
final_loss = loss.sum()

```

For the optimizer that would depend on the optimizer I guess.  
For Adam for example, it need to store one (if I remember correctly?) extra copy of all the weights. And it need this copy to update the weights. That means that if all the gradients are accumulated on a single gpu, then this state should be there as well.  
That being said. The size of the weights of your network should not be a very large part of the memory consumption for classic nets. Intermediary states are the most demanding. So hopefully this extra memory as big as the set of weights should not be too big a problem.

---

<div class="post-metadata">

**Author:** ![Yuzhou\_Song](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/yuzhou_song/32/3937_2.png) [@Yuzhou\_Song](https://discuss.pytorch.org/u/Yuzhou_Song)\
**Post date:** [August 8, 2018, 5:08pm UTC](https://discuss.pytorch.org/t/dataparallel-imbalanced-memory-usage/22551/7 "2018-08-08T17:08:33Z")

</div>

I had tried put the loss inside of forward() before. It did mitigate this imbalanced issue but not fully resolved it. BTW, I wonder if PyTorch has plans to implement the functionality to store model parameters in CPU instead of one main GPU? I’m training production scale data which is super big. If the GPU memory is used by 99% after mounting the model, there is little room for me to mount data anymore, which totally lose the scalability of multiple GPU training.

---

<div class="post-metadata">

**Author:** ![albanD](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/alband/32/215_2.png) [@albanD](https://discuss.pytorch.org/u/albanD)\
**Post date:** [August 8, 2018, 5:11pm UTC](https://discuss.pytorch.org/t/dataparallel-imbalanced-memory-usage/22551/8 "2018-08-08T17:11:24Z")

</div>

Well you need to store model parameters on the gpu to be able to do the forward.  
The backward computes gradients on the gpu so they will be there.  
Not sure what you mean by “store model parameters in CPU instead of one main GPU” ?

---

<div class="post-metadata">

**Author:** ![Yuzhou\_Song](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/yuzhou_song/32/3937_2.png) [@Yuzhou\_Song](https://discuss.pytorch.org/u/Yuzhou_Song)\
**Post date:** [August 8, 2018, 5:31pm UTC](https://discuss.pytorch.org/t/dataparallel-imbalanced-memory-usage/22551/9 "2018-08-08T17:31:45Z")

</div>

According to my understanding, one/multiple CPUs could serve as Parameter Server storing model parameters. In forward step, each GPU could mount a batch of data and part of model parameters (let’s say first 10 layers’), do the calculation. Then, abandon the first 10 layers’ and mount 11-20 layers’ and do calculation again. We repeat this until get the final output of forward(). The benefit is model size is no longer limited by GPU memory. I guess typical distributed training framework is implemented this way, right?

If my understanding is incorrect, what’s the best way to train a model with parameter size exceeds GPU memory?

---

<div class="post-metadata">

**Author:** ![albanD](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/alband/32/215_2.png) [@albanD](https://discuss.pytorch.org/u/albanD)\
**Post date:** [August 8, 2018, 5:40pm UTC](https://discuss.pytorch.org/t/dataparallel-imbalanced-memory-usage/22551/10 "2018-08-08T17:40:48Z")

</div>

Hi,

You could do that but transferring few GB of data many times between the cpu and gpu will be quite slow.  
But in practice, as I said above, what is going to take a lot of space are the intermediary results. And that is going to be even worst.

That being said, if your problem is that your net is too big, even for a batch size of 1, then you don’t want data parallelism but model parallelism.  
This can be done very easily by adding some `.cuda(1)` then `.cuda(2)`… during your forward pass, and sending part of your net to different gpus. according to this split.  
There is no builtin method to do this as it need to be done in your forward method directly and you need to ensure that the right part of each model is moved to the right gpu.

---

<div class="post-metadata">

**Author:** ![Yuzhou\_Song](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/yuzhou_song/32/3937_2.png) [@Yuzhou\_Song](https://discuss.pytorch.org/u/Yuzhou_Song)\
**Post date:** [August 8, 2018, 5:49pm UTC](https://discuss.pytorch.org/t/dataparallel-imbalanced-memory-usage/22551/11 "2018-08-08T17:49:55Z")

</div>

That makes sense to me! It could be great if you guys can make a post/blog/sample code about what we discussing here. It must be very helpful!

Thanks a lot,  
Best,  
Yuzhou

---

<div class="post-metadata">

**Author:** ![JuanFMontesinos](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/juanfmontesinos/32/76115_2.png) [@JuanFMontesinos](https://discuss.pytorch.org/u/JuanFMontesinos)\
**Post date:** [August 9, 2018, 2:22am UTC](https://discuss.pytorch.org/t/dataparallel-imbalanced-memory-usage/22551/12 "2018-08-09T02:22:37Z")

</div>

Hi there @albanD, @Yuzhou_Song  
I noticed there is an small mistake with the code you provided:  
It’s necessary to unsqueeze loss inside forward pass to DataParallel were able to build loss back. Loss provided by PyTorch loss functions seems not to have dimensions, and DataParallel mount batch back in dim 0 by default. It also requires to class super class not to raise up with an error. Anyway thank you very much. Awesome help.

```auto
class FullModel(nn.Module):
  def __init__ (self, model, loss):
    super(FullModel, self). __init__ ()
    self.model = model
    self.loss = loss

  def forward(self, targets, *inputs):
    outputs = self.model(*inputs)
    loss = self.loss(outputs, targets)
    return torch.unsqueeze(loss,0),outputs
    

def DataParallel_withLoss(model,loss,**kwargs):
    model=FullModel(model, loss)
    if 'device_ids' in kwargs.keys():
        device_ids=kwargs['device_ids']
    else:
        device_ids=None
    if 'output_device' in kwargs.keys():
        output_device=kwargs['output_device']
    else:
        output_device=None
    if 'cuda' in kwargs.keys():
        cudaID=kwargs['cuda'] 
        model=torch.nn.DataParallel(model, device_ids=device_ids, output_device=output_device).cuda(cudaID)
    else:
        model=torch.nn.DataParallel(model, device_ids=device_ids, output_device=output_device).cuda()
    return model

```

I was trying this implementation and it works.  
However imbalance keeps happening

 ![Selection_037](https://discuss.pytorch.org/uploads/default/original/2X/c/cade79ccad98b96837d067587ca953362ee3dd95.png)

So, as far as i could check printing parameter.device model weights are stored in gpu by-default gpu. In the case I run, we can roughly say each mini-batch requires ~5Gb and model weights ~4Gb (8500-4500)  
There is a gain in computing loss in a distributed way.

Is not possible to share model weights among all gpus? Cos in this example, i can increase batch size until gpu uses 12Gb, but it would mean gpu1 and gpu2 would be using only 8.

Even a worse case, if i were using Adam (you mentioned it copies all the model weigths), memory usage in GPU0 would be 12 Gb meanwhile gpu1 and gpu2 would be using 4.

If this were not possible due to pytorch requirements (I guess pytorch requires all tensors to be in the same gpu to be able to operate them)  
Is there are way of penalizing workload, this is, an imbalanced amount of samples per gpu to reduce by-default gpu memory requirements?

Toy example  
BS=30  
gpu0–\>4  
gpu1–\>13  
gpu2–\>13

---

<div class="post-metadata">

**Author:** ![albanD](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/alband/32/215_2.png) [@albanD](https://discuss.pytorch.org/u/albanD)\
**Post date:** [August 9, 2018, 9:03am UTC](https://discuss.pytorch.org/t/dataparallel-imbalanced-memory-usage/22551/13 "2018-08-09T09:03:33Z")

</div>

I am by no mean a specialist on DataParallel and I have a cpu-only install so I couldn’t test the code.  
The doc states explicitly that if the module returns a 0dim tensor, then you will get a 1D tensor with size the number of gpus. That does not work? Maybe that’s only supported in the latest version.

Sharing the weights among gpus is not possible as they need to be on the device to perform computations.

I am not sure at the moment what creates this inbalance…  
Maybe @smth knows? 🙂

---

<div class="post-metadata">

**Author:** ![soulless](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/soulless/32/40008_2.png) [@soulless](https://discuss.pytorch.org/u/soulless)\
**Post date:** [August 10, 2018, 8:28am UTC](https://discuss.pytorch.org/t/dataparallel-imbalanced-memory-usage/22551/14 "2018-08-10T08:28:07Z")

</div>

I am able to reach a relatively balanced memory usage by incorporating the loss calculation in the forward pass. Thanks for the help  
 ![19%20PM](https://discuss.pytorch.org/uploads/default/original/2X/a/ad4295b39872f17fb4b560bc2fedddcedea8a9fc.png)

---

<div class="post-metadata">

**Author:** ![albanD](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/alband/32/215_2.png) [@albanD](https://discuss.pytorch.org/u/albanD)\
**Post date:** [August 10, 2018, 9:04am UTC](https://discuss.pytorch.org/t/dataparallel-imbalanced-memory-usage/22551/15 "2018-08-10T09:04:04Z")

</div>

Perfect !

That means that the inbalance that @JuanFMontesinos sees is due to something in his code that creates that. Not a problem with DataParallel.

---

<div class="post-metadata">

**Author:** ![JuanFMontesinos](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/juanfmontesinos/32/76115_2.png) [@JuanFMontesinos](https://discuss.pytorch.org/u/JuanFMontesinos)\
**Post date:** [August 10, 2018, 10:54am UTC](https://discuss.pytorch.org/t/dataparallel-imbalanced-memory-usage/22551/16 "2018-08-10T10:54:55Z")

</div>

Well glad to hear that, then i will find a way to solve it ^^ Ty @albanD for your help.  
Just for discarding, how many parameters does your net have @soulless?  
Cos I’m using a very deep multimodal net

---

<div class="post-metadata">

**Author:** ![soulless](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/soulless/32/40008_2.png) [@soulless](https://discuss.pytorch.org/u/soulless)\
**Post date:** [August 12, 2018, 4:09am UTC](https://discuss.pytorch.org/t/dataparallel-imbalanced-memory-usage/22551/17 "2018-08-12T04:09:39Z")

</div>

I am using a CNN autoencoder with around 8 million trainable parameters. How many parameters do you have?

---

<div class="post-metadata">

**Author:** ![JuanFMontesinos](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/juanfmontesinos/32/76115_2.png) [@JuanFMontesinos](https://discuss.pytorch.org/u/JuanFMontesinos)\
**Post date:** [August 21, 2018, 5:57pm UTC](https://discuss.pytorch.org/t/dataparallel-imbalanced-memory-usage/22551/18 "2018-08-21T17:57:22Z")

</div>

I guess, not confirmed yet, that if we save/load a model and this model is directly stored from gpu, the memory is not freed.

I mean, I’ve been working is several clusters and saving model directly from dataparallel module. I’ve had all these troubles loading pretrained weights from this dataparallel state dic. However I found that saving the model prepared for cpu, loading the model again and then moving it to cuda uses less memory.

---

<div class="post-metadata">

**Author:** ![soulless](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/soulless/32/40008_2.png) [@soulless](https://discuss.pytorch.org/u/soulless)\
**Post date:** [August 22, 2018, 2:48am UTC](https://discuss.pytorch.org/t/dataparallel-imbalanced-memory-usage/22551/19 "2018-08-22T02:48:31Z")

</div>

Do you mean it would save more memory by moving the dataparallel model back to the cpu model before saving it ?

---

<div class="post-metadata">

**Author:** ![Hjules](https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/h/dc4da7/32.png) [@Hjules](https://discuss.pytorch.org/u/Hjules)\
**Post date:** [November 14, 2018, 11:42am UTC](https://discuss.pytorch.org/t/dataparallel-imbalanced-memory-usage/22551/20 "2018-11-14T11:42:59Z")

</div>

Dear soulless,

Do you by any chance still have the script around for which you made this work?

I am trying to do what albanD suggested above, but it just won’t work. It’d be really nice to have an example to look at.

Regards,  
Hjules

[Next page](https://discuss.pytorch.org/t/dataparallel-imbalanced-memory-usage/22551.md?page=2)