DistributedDataParallel: GRU module gets additional processes on GPU 0 (1st GPU) and takes more memory

username: Chenguang_Wan
cooked: <h1><a name="problem-description-1" class="anchor" href="#problem-description-1"></a>Problem description</h1>
<p>Hi, When I am testing a simple example with DistributedDataParallel, using a single node with 4 gpus, I found that when I used the GRU or LSTM module was taking additional processes and more memory on GPU 0. while using the Linear was not gotten these problems. The test code snippets are as follows:</p>
<pre><code class="lang-auto">def run_demo(demo_fn, world_size):
    mp.spawn(demo_fn,
             args=(world_size,),
             nprocs=world_size,
             join=True)
def demo_basic(rank, world_size):
    print(f"Running basic DDP example on rank {rank}.")
    dist.init_process_group("nccl", rank=rank, world_size=world_size)
    # The GRU or LSTM gets additional processes on GPU 0.
    ToyModel = nn.GRU(10, 10, 1)
    # The Linear does not get these problems.
    # ToyModel = nn.Linear(10,1)
    model = ToyModel.to(rank)
    ddp_model = DDP(model, device_ids=[rank])
    pbar_len = int(1e10 / 2)
    for _ in range(pbar_len):
        input_seq = torch.randn(4, 20,10)
        input_seq = input_seq.float().to(rank)
        ddp_model(input_seq)
    dist.destroy_process_group()
if __name__ == "__main__":
    world_size = torch.cuda.device_count()
    os.environ['MASTER_ADDR'] = 'localhost'
    os.environ['MASTER_PORT'] = '12355'
    run_demo(demo_basic, world_size)
</code></pre>
<p>I called the script like python XX.py.  And I got the two results as follows:<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/9/1/9143b9e0cedc8e41b5690b9fd2b72fef72b50dc8.png" data-download-href="https://discuss.pytorch.org/uploads/default/9143b9e0cedc8e41b5690b9fd2b72fef72b50dc8" title="GPU_results"><img src="https://discuss.pytorch.org/uploads/default/original/3X/9/1/9143b9e0cedc8e41b5690b9fd2b72fef72b50dc8.png" alt="GPU_results" data-base62-sha1="kJ4kOns7gAjakqOnrnt99YVibXi" width="690" height="485" data-dominant-color="1B1B1B"><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">GPU_results</span><span class="informations">1002×705 20.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>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/9/9/9944e8a5994ab67fb2b2889d38de4fe4da93c81c.png" data-download-href="https://discuss.pytorch.org/uploads/default/9944e8a5994ab67fb2b2889d38de4fe4da93c81c" title="Linear_results"><img src="https://discuss.pytorch.org/uploads/default/original/3X/9/9/9944e8a5994ab67fb2b2889d38de4fe4da93c81c.png" alt="Linear_results" data-base62-sha1="lRSGJRVxfYgUWmI4RDyi7czMqEk" width="690" height="484" data-dominant-color="1B1B1B"><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">Linear_results</span><span class="informations">969×680 19.6 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>
<h1><a name="versions-2" class="anchor" href="#versions-2"></a>Versions</h1>
<p>Collecting environment information…<br>
PyTorch version: 1.10.0<br>
Is debug build: False<br>
CUDA used to build PyTorch: 10.2<br>
ROCM used to build PyTorch: N/A</p>
<p>OS: CentOS Linux release 7.7.1908 (Core) (x86_64)<br>
GCC version: (GCC) 4.8.5 20150623 (Red Hat 4.8.5-44)<br>
Clang version: 3.4.2 (tags/RELEASE_34/dot2-final)<br>
CMake version: Could not collect<br>
Libc version: glibc-2.17</p>
<p>Python version: 3.9.7 (default, Sep 16 2021, 13:09:58)  [GCC 7.5.0] (64-bit runtime)<br>
Python platform: Linux-3.10.0-514.el7.x86_64-x86_64-with-glibc2.17<br>
Is CUDA available: True<br>
CUDA runtime version: Could not collect<br>
GPU models and configuration:<br>
GPU 0: Tesla P100-PCIE-16GB<br>
GPU 1: Tesla P100-PCIE-16GB<br>
GPU 2: Tesla P100-PCIE-16GB<br>
GPU 3: Tesla P100-PCIE-16GB</p>
<p>Nvidia driver version: 460.27.04<br>
cuDNN version: Could not collect<br>
HIP runtime version: N/A<br>
MIOpen runtime version: N/A</p>
<p>Versions of relevant libraries:<br>
[pip3] mypy-extensions==0.4.3<br>
[pip3] numpy==1.21.2<br>
[pip3] torch==1.10.0<br>
[pip3] torch-tb-profiler==0.1.0<br>
[pip3] torchaudio==0.10.0<br>
[pip3] torchinfo==1.5.4<br>
[pip3] torchvision==0.11.1<br>
[conda] blas                      1.0                         mkl    defaults<br>
[conda] cudatoolkit               10.2.89              hfd86e86_1    defaults<br>
[conda] ffmpeg                    4.3                  hf484d3e_0    pytorch<br>
[conda] mkl                       2021.4.0           h06a4308_640    defaults<br>
[conda] mkl-service               2.4.0            py39h7f8727e_0    defaults<br>
[conda] mkl_fft                   1.3.1            py39hd3c417c_0    defaults<br>
[conda] mkl_random                1.2.2            py39h51133e4_0    defaults<br>
[conda] mypy_extensions           0.4.3            py39h06a4308_0    defaults<br>
[conda] numpy                     1.21.2           py39h20f2e39_0    defaults<br>
[conda] numpy-base                1.21.2           py39h79a1101_0    defaults<br>
[conda] pytorch                   1.10.0          py3.9_cuda10.2_cudnn7.6.5_0    pytorch<br>
[conda] pytorch-mutex             1.0                        cuda    pytorch<br>
[conda] torch-tb-profiler         0.1.0                    pypi_0    pypi<br>
[conda] torchaudio                0.10.0               py39_cu102    pytorch<br>
[conda] torchinfo                 1.5.4              pyhd8ed1ab_0    conda-forge<br>
[conda] torchvision               0.11.1               py39_cu102    pytorc</p>

username: changtaoli
cooked: <p>I met the same problem when training my ResNet101 on CIFAR100. Do you figure it out?</p>

username: Chenguang_Wan
cooked: <p>I find the solution by the help with <a href="https://github.com/pytorch/pytorch/issues/70404#issuecomment-1001113109" rel="noopener nofollow ugc">@ngimel</a>. This is the solution learned from her.</p>
<pre><code class="lang-python">
def demo_basic(rank, world_size):

    print(f"Running basic DDP example on rank {rank}.")

    dist.init_process_group("nccl", rank=rank, world_size=world_size)

    torch.cuda.set_device(rank)

    ToyModel = nn.GRU(10, 10, 1)

    model = ToyModel.cuda()

    ddp_model = DDP(model, device_ids=[rank])

    pbar_len = int(1e10 / 2)

    for _ in range(pbar_len):

        input_seq = torch.randn(4, 20,10)

        input_seq = input_seq.float().cuda()

        ddp_model(input_seq)

    dist.destroy_process_group()

</code></pre>