Why pytorch needs so much memory to execute distributed training?

khalil: <p>Today I want to do distributed training in many nodes, however the memory decrease rapidly.</p>
<p><img src="https://discuss.pytorch.org/uploads/default/original/3X/7/6/76626668b2bf3663b47e7dca2f9db6a8556818a2.png" alt="image" data-base62-sha1="gTha55U4DvSwfMRL7IRdLtPS1uW" width="526" height="280"></p>
<p>I execute the command 'dstat -c -m '.We can see the available memory decrease from 64GB to 177MB.</p>
<p>My batch_size is 1 and datasets include 4000 samples,every sample is about 16MB.I use the gloo as the backend.<br>
data_loader code like this</p>
<pre><code class="lang-auto">class HDF5Dataset(data.Dataset):
    """Represents an abstract HDF5 dataset.
    
    Input params:
        file_path: Path to the folder containing the dataset (one or multiple HDF5 files).
        recursive: If True, searches for h5 files in subdirectories.
        transform: PyTorch transform to apply to every data instance (default=None).
        divide: use only 1//divide dataset
    """
    def __init__(self, file_path, recursive, transform=None, divide =1):
        super().__init__()
        self.data_info = []
        self.data_cache = {} #dict
        self.transform = transform
        self.length = 0
        self.divide = divide

        # Search for all h5 files
        p = Path(file_path)
        assert(p.is_dir())
        
        
        if recursive:
            files = sorted(p.glob('**/*.hdf5'))
        else:
            files = sorted(p.glob('*.hdf5'))
        if len(files) &lt; 1:
            raise RuntimeError('No hdf5 datasets found')
        
        self.length = len(files) // divide
        self.data_info = files[0: self.length]
               
        #print("File Prepared !\n")
            
    def __getitem__(self, index):
        # get data
        y, x = self.get_data(index)
        
        if self.transform:
            x = self.transform(x)
        else:
            x = torch.from_numpy(x)
            x = torch.reshape(x, (1,128,128,128))
            x = x.type(torch.FloatTensor)

        # get label
        y = torch.from_numpy(y)
        y = y.type(torch.FloatTensor)
        return (x, y)

    def get_data(self, i):
        fp = self.data_info[i] #list - dict
        #print(fp)
        
        try:
            with h5py.File(fp,'r') as h5_file:
        
                label = h5_file.get('label')[()]            
                dataset = h5_file.get('dataset')[()]
            
                return (label, dataset)
        except IOError:
            print('Error:File {filepath} is broken.'.format(filepath=fp))
            i=i+1
            return self.get_data(i)
        
        
    def __len__(self):
        return self.length
</code></pre>
<p>I want to know why pytorch DDP training need more memory with the nodes expand. And how to reduce memory usage？</p>

mrshenli: <p>What is the <code>world_size</code> in your application? Each process will have its own model replica, so the total memory consumption is expected to be larger than <code>world_size</code> X (non-DDP memory footprint). Besides, DDP also creates buffers for communication, which will contribute another ~1.5X if not considering optimizer.</p>

khalil: <p>Thanks for your answer.I set the parameter nproc_per_node to 1 and I guess the world_size is 1.There is only one process per node.However, it still consumes a lot of memory. DDP adopts the data parallelism, I think when I train this model in more nodes, there is little memory to be consumed.<br>
I found that the node cache always uses more than 20GB of memory and if I can not reduce the memory usage,I should choose another framework to do distributed training = = !</p>

mrshenli: <p>Hey <a class="mention" href="/u/khalil">@khalil</a></p>
<p>How large is your model?</p>
<blockquote>
<p>My batch_size is 1 and datasets include 4000 samples,every sample is about 16MB.</p>
</blockquote>
<p>Does this actually mean the memory is consumed by the data loader? Can you check the memory consumption when using the same model and same data loader <strong>without</strong> DDP?</p>

mrshenli: <p>cc <a class="mention" href="/u/vincentqb">@vincentqb</a> <a class="mention" href="/u/alband">@albanD</a> for dataloder questions</p>

khalil: <p>Thank you,<a class="mention" href="/u/mrshenli">@mrshenli</a>. I have trained this model in one node and it consumed 30GB on the cache.Just like this，<br>
<img src="https://discuss.pytorch.org/uploads/default/original/3X/8/3/83463676db34a618fb87c6b003f5627d023bdd37.png" alt="image" data-base62-sha1="iJiYGkyd6ds9Ib1KxGWfayvYRPV" width="286" height="43"> .<br>
<img src="https://discuss.pytorch.org/uploads/default/original/3X/c/b/cbd593a08180fde0f1065621564743696a6b07d9.png" alt="image" data-base62-sha1="t5cvLFFW77anWI92DDzYci3ABTz" width="286" height="132"></p>
<p>I guess there are so many data to be loaded.Could you help me check my data_loader code?</p>

mrshenli: <p>Hey <a class="mention" href="/u/khalil">@khalil</a></p>
<p>I am trying to identify which component hogs memory. Could you please share some more details about the sentence below.</p>
<blockquote>
<p>I have trained this model in one node and it consumed 30GB on the cache</p>
</blockquote>
<p>By “trained this model in one node”, do you still use DDP and the same data loader? It will be helpful to know the memory footprint of the following cases:</p>
<ol>
<li>Train a local model <strong>without</strong> data loader and <strong>without</strong> DDP. Just feed the model with some randomly generated tensors.</li>
<li>Train a local model <strong>with</strong> data loader but without DDP.</li>
<li>Wrap the model with DDP, and train it with data loader.</li>
</ol>

khalil: <p>Thank you,<a class="mention" href="/u/mrshenli">@mrshenli</a>,I will try it.</p>

khalil: <p>I have created a small dataset which only contains 30 samples.However the memory still decreased rapidly.(batch_size=1)</p>
<p><img src="https://discuss.pytorch.org/uploads/default/original/3X/d/c/dc85b85614e1c8dc78df1f16e1a6c9119b4d37d6.png" alt="image" data-base62-sha1="vsPxjOH7Cc1Zyhkg3myC6QSC1HU" width="424" height="151"></p>
<p>So I think this issue is not caused by dataloader and I am sure reason is the node expansion.</p>
<p>When I do DDP training in one node,the memory problem is not serious.When I expand the nodes,the memory problem is serious.I want to know if this is pytorch’s problem?</p>
<p>My model like this<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/1/c/1ce6a5d7b981a495c1044f4c6bfcbcec8914f8d9.png" data-download-href="https://discuss.pytorch.org/uploads/default/1ce6a5d7b981a495c1044f4c6bfcbcec8914f8d9" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/1/c/1ce6a5d7b981a495c1044f4c6bfcbcec8914f8d9_2_690x250.png" alt="image" data-base62-sha1="47FwQL5zC3p4GlZ8sZ3YqU1BJjz" width="690" height="250" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/1/c/1ce6a5d7b981a495c1044f4c6bfcbcec8914f8d9_2_690x250.png, https://discuss.pytorch.org/uploads/default/optimized/3X/1/c/1ce6a5d7b981a495c1044f4c6bfcbcec8914f8d9_2_1035x375.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/1/c/1ce6a5d7b981a495c1044f4c6bfcbcec8914f8d9.png 2x" data-dominant-color="EEF1F7"><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">image</span><span class="informations">1139×413 72.2 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></p>

mrshenli: <p>Hey <a class="mention" href="/u/khalil">@khalil</a></p>
<p>Could you please provide details of the following question?</p>
<ol start="0">
<li>
<p>Which version of PyTorch are you using?</p>
</li>
<li>
<p>What’s the size of your model (how many GB)? You can get this by doing the following:</p>
<pre><code class="lang-python">model_size = 0
for p in model.parameters():
    model_size += p.numel() * p.element_size()
</code></pre>
<p>Given the picture shown in the previous post, it looks like the model is about 27GB? If that is the case, then, yes, DDP would use another 27GB as comm buffers. We are working on improving this: <a href="https://github.com/pytorch/pytorch/issues/39022" rel="nofollow noopener">https://github.com/pytorch/pytorch/issues/39022</a></p>
<p>I am curious how did you train this model locally without DDP? If the model is 27GB, after the backward pass, the grads will also consume 27GB. So local training without DDP will also use &gt;54GB memory?</p>
<p>And if your optimizer uses any momentum etc., it’s likely the optimizer will also consume a few more X of 27GB. But looks like there are only 64GB memory available? Is it because your optimizer does not contain any states?</p>
</li>
<li>
<p>Does this getting worse when you use more nodes (e.g., scale from 2 nodes to 3 nodes)?</p>
</li>
</ol>

khalil: <p><a class="mention" href="/u/mrshenli">@mrshenli</a>,Thanks for your general help.The details are as follows:</p>
<ol start="0">
<li>
<p>The version of PyTorch is 1.5.0</p>
</li>
<li>
<p>I use your code to test the size of my model and it printed this sentence: model_size is 9865484.So I think the size of model are 10MB…</p>
</li>
<li>
<p>When I scaled from 2 nodes to 4 nodes,the available memory started to decrease.The memory usage went up with the number of nodes increased.</p>
</li>
</ol>

mrshenli: <blockquote>
<p>I use your code to test the size of my model and it printed this sentence: model_size is 9865484.So I think the size of model are 10MB</p>
</blockquote>
<p>In this case DDP should only consume 10MB more memory for communication buffers.</p>
<p>BTW, as the model is just 10MB, do you know why even without DDP it consumes 30GB? And have you seen OOM errors?</p>
<blockquote>
<p>When I scaled from 2 nodes to 4 nodes, the available memory started to decrease.The memory usage went up with the number of nodes increased.</p>
</blockquote>
<p>This is also not expected with DDP. DDP’s memory footprint should be constant, regardless of how many nodes are used.</p>
<p>Do you have a minimum repro that we can investigate?</p>

khalil: <p>Dear <a class="mention" href="/u/mrshenli">@mrshenli</a>,forgive me for not being clear,I said it consumes 30GB which means the cache<br>
consumes 30GB.I guess the memory is used to load the disk data.I tried to use a small datasets<br>
which only contains 40 samples(Every sample is 16MB).And in this case,the cache consumption is about 2.5GB.However,the memory problem still exist.You can see this picture:<br>
<img src="https://discuss.pytorch.org/uploads/default/original/3X/d/c/dc85b85614e1c8dc78df1f16e1a6c9119b4d37d6.png" alt="image" data-base62-sha1="vsPxjOH7Cc1Zyhkg3myC6QSC1HU" width="424" height="151"></p>
<p>Now I think the problem is not caused by dataloader.</p>
<p>And you say DDP’s memory footprint should be constant,regardless of how many nodes are used.I test the maximum memory footprint as the nodes expand,it just like this<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/2/0/20bcaa1469b6cd976ebbba3850b544ed7a05958a.png" data-download-href="https://discuss.pytorch.org/uploads/default/20bcaa1469b6cd976ebbba3850b544ed7a05958a" title="node-memory"><img src="https://discuss.pytorch.org/uploads/default/original/3X/2/0/20bcaa1469b6cd976ebbba3850b544ed7a05958a.png" alt="node-memory" data-base62-sha1="4FBupUr3Hig4dyHyFcLvdebUCsW" width="690" height="202" data-dominant-color="D0CDE2"><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">node-memory</span><span class="informations">1366×401 5.92 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></p>
<p>The memory footprint does not include the cache consumption.</p>
<p>And when I test the memory footprint about 32 nodes,I found a memory allocated error which indicates the memory has run out.</p>

mrshenli: <p>Thanks for the detailed information. This is weird. This is the first time that we saw the memory footprint per machine increases significantly with the total number of machines.</p>
<blockquote>
<p>I guess the memory is used to load the disk data.I tried to use a small datasets<br>
which only contains 40 samples(Every sample is 16MB).And in this case,the cache consumption is about 2.5GB.</p>
</blockquote>
<p>Q1: This does not add up. The model is 10MB, and the dataset is 40X16MB = 640MB.  I assume you do not use CUDA, as you mainly concerned about CPU memory. In this case, I would assume the total memory consumption to be less than 1GB. Do you know where does the other 1.5GB come from?</p>
<p>Q2: Can you try to the following to see if it is indeed used by the process instead of some os cache.</p>
<pre data-code-wrap="python"><code class="lang-python">import psutil

process = psutil.Process(os.getpid())

for _ in get_batch():
    ....
    # print this in every iteration
    print("Memory used: ", process.memory_info().rss)
</code></pre>
<p>Q3: Do you have a minimum reprodueable example that we can investigate locally?</p>
<p>Q4: BTW, how do you launch your distributed training job? Are you using <code>torch.distributed.launch</code>. If so, could you please share the command you are using?</p>

khalil: <p><a class="mention" href="/u/mrshenli">@mrshenli</a>,Sorry,I can not reply you in time.</p>
<p>Now I found the other 1.5GB comes from system.There is 1.5GB consumption before I start this job so I think cache consumption is right.</p>
<p>I use the psutil module to test the process memory footprint.The result just like the following photo:<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/5/9/59e36f955fa72bba53ddfe9d516fabda0acf955d.png" data-download-href="https://discuss.pytorch.org/uploads/default/59e36f955fa72bba53ddfe9d516fabda0acf955d" title="image"><img src="https://discuss.pytorch.org/uploads/default/original/3X/5/9/59e36f955fa72bba53ddfe9d516fabda0acf955d.png" alt="image" data-base62-sha1="cPbLZbfaMu0MGdl5K7V0vfYGJhX" width="690" height="166" data-dominant-color="0C0C0B"><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">image</span><span class="informations">1096×265 13.4 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>
(1 node)</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/4/c/4ccbc3baa2d75d32bee1e658b84dcad40cfa8622.png" data-download-href="https://discuss.pytorch.org/uploads/default/4ccbc3baa2d75d32bee1e658b84dcad40cfa8622" title="image"><img src="https://discuss.pytorch.org/uploads/default/original/3X/4/c/4ccbc3baa2d75d32bee1e658b84dcad40cfa8622.png" alt="image" data-base62-sha1="aXmQMuuUO0JyOcOiN1CMXMXzBWG" width="690" height="193" data-dominant-color="0B0A0A"><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">image</span><span class="informations">1226×343 16.2 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>
(16 node)</p>
<p>So the process memory footprint per machine does not increase with the total number of machines.However,the total memory footprint per machine increases significantly.</p>
<p>Then I found the process numbers increase significantly with the number of machines.<br>
<img src="https://discuss.pytorch.org/uploads/default/original/3X/c/4/c4e8268d4dfc2cb0f080deba2c25ca373d0f4e01.png" alt="image" data-base62-sha1="s5UXNLTgFwBFiExMifmECHNKdJD" width="494" height="276"></p>
<p>I set nproc_per_node to 1 so there should be one process in every machine.However,it is not.<br>
There is my command:<br>
bash.sh</p>
<pre><code class="lang-auto">MIP=$ip
MPORT=$port
NPROC_PER_NODE=1
HOSTLIST=$hostlist
COMMAND=$HOME/sample.py --epochs=120 --workers=0 --batch-size=1 --print-freq=50 --data=$HOME/datasets/v3
RANK=0
for node in $HOSTLIST; do
        echo $node
        ssh -q $node \
                python -m torch.distributed.launch \
                --nproc_per_node=$NPROC_PER_NODE \
                --nnodes=$SLURM_JOB_NUM_NODES \
                --node_rank=$RANK \
                --master_addr="$MIP" --master_port=$MPORT \
                $COMMAND &gt; "log_v1_"${SLURM_JOB_ID}"_"${RANK}".out" &amp;
        RANK=$((RANK+1))
done
</code></pre>
<p>sample.py</p>
<pre><code class="lang-auto">....
if __name__ == '__main__':
    if torch.distributed.is_available():
        dist.init_process_group(backend='gloo',init_method='env://',timeout=datetime.timedelta(seconds=600))
        main()
</code></pre>
<p>In addition,I do not use CUDA and I train this model in CPUs.</p>

khalil: <p>I do not think pytorch is the cause of the mistake.I will check my shell script,thank you <a class="mention" href="/u/mrshenli">@mrshenli</a></p>

asmuqing: <p>Hi <a class="mention" href="/u/khalil">@khalil</a>, have you solve this problem?<br>
I am actually meeting exactly the same problem as you.</p>

khalil: <p>Sorry abou the delay, I have solved this problem. The reason is that I have execute the script in same node.(ssh failed)  I think you can check to see if the process has started on each node.</p>