<p dir="auto"><strong>Documentation issue</strong></p>
<p dir="auto">In the <a href="https://pytorch.org/docs/stable/nn.html#lstm" rel="nofollow">documentation</a> of PyTorch 1.0.1. It is mentioned that:</p>
<blockquote>
<p dir="auto">batch_first – If True, then the input and output tensors are provided as (batch, seq, feature). Default: False</p>
</blockquote>
<p dir="auto">So I would expect the ouput tensor h_n of shape: (batch, num_layers * num_directions, hidden_size). But it is not the case. The batch is the second dimension, which does not match with the above statement from the documentation.</p>
<h2 dir="auto">To Reproduce</h2>
<p dir="auto">Steps to reproduce the behavior:</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="input = torch.rand([60, 8, 256]) #batch size is 60, sequence size is 8, features are 256
lstm = torch.nn.LSTM(256, 512, batch_first=True)
_, (finalHiddenState, _) = lstm(input)
print(finalHiddenState.shape)
Ouput: torch.Size([1, 60, 512]) "><pre class="notranslate"><code class="notranslate">input = torch.rand([60, 8, 256]) #batch size is 60, sequence size is 8, features are 256
lstm = torch.nn.LSTM(256, 512, batch_first=True)
_, (finalHiddenState, _) = lstm(input)
print(finalHiddenState.shape)
Ouput: torch.Size([1, 60, 512]) 
</code></pre></div>
<p dir="auto">Batch (60) is now the second dimension here</p>
<h2 dir="auto"><g-emoji class="g-emoji" alias="bug" fallback-src="https://github.githubassets.com/images/icons/emoji/unicode/1f41b.png">🐛</g-emoji> Bug</h2>
<p dir="auto">When using MPI to launch processes, and mpi4py to assign ranks and get world size, DDP fails to initialize.</p>
<h2 dir="auto">To Reproduce</h2>
<p dir="auto">Steps to reproduce the behavior:</p>
<p dir="auto">Here is a complete reproducer script:</p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="import torch
import random

# This calls mpi.Init()
from mpi4py import MPI

# What rank?  World size?
rank       = MPI.COMM_WORLD.Get_rank()
world_size = MPI.COMM_WORLD.Get_size()

# Pytorch data parallel:
import torch
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
import torch.multiprocessing as mp


# Create a 1-layer network, put it on GPU:
n = torch.nn.Linear(1, 128).cuda()

# May need to change init method to not be 'cadams' home dir ...
init_method = f'file:///home/cadams/.torch_ddp_init'
dist.init_process_group(&quot;nccl&quot;, init_method=init_method,rank=rank, world_size=world_size)

n = DDP(n)

print(&quot;SUCCESS&quot;)"><pre class="notranslate"><span class="pl-k">import</span> <span class="pl-s1">torch</span>
<span class="pl-k">import</span> <span class="pl-s1">random</span>

<span class="pl-c"># This calls mpi.Init()</span>
<span class="pl-k">from</span> <span class="pl-s1">mpi4py</span> <span class="pl-k">import</span> <span class="pl-v">MPI</span>

<span class="pl-c"># What rank?  World size?</span>
<span class="pl-s1">rank</span>       <span class="pl-c1">=</span> <span class="pl-v">MPI</span>.<span class="pl-v">COMM_WORLD</span>.<span class="pl-v">Get_rank</span>()
<span class="pl-s1">world_size</span> <span class="pl-c1">=</span> <span class="pl-v">MPI</span>.<span class="pl-v">COMM_WORLD</span>.<span class="pl-v">Get_size</span>()

<span class="pl-c"># Pytorch data parallel:</span>
<span class="pl-k">import</span> <span class="pl-s1">torch</span>
<span class="pl-k">import</span> <span class="pl-s1">torch</span>.<span class="pl-s1">distributed</span> <span class="pl-k">as</span> <span class="pl-s1">dist</span>
<span class="pl-k">from</span> <span class="pl-s1">torch</span>.<span class="pl-s1">nn</span>.<span class="pl-s1">parallel</span> <span class="pl-k">import</span> <span class="pl-v">DistributedDataParallel</span> <span class="pl-k">as</span> <span class="pl-v">DDP</span>
<span class="pl-k">import</span> <span class="pl-s1">torch</span>.<span class="pl-s1">multiprocessing</span> <span class="pl-k">as</span> <span class="pl-s1">mp</span>


<span class="pl-c"># Create a 1-layer network, put it on GPU:</span>
<span class="pl-s1">n</span> <span class="pl-c1">=</span> <span class="pl-s1">torch</span>.<span class="pl-s1">nn</span>.<span class="pl-v">Linear</span>(<span class="pl-c1">1</span>, <span class="pl-c1">128</span>).<span class="pl-en">cuda</span>()

<span class="pl-c"># May need to change init method to not be 'cadams' home dir ...</span>
<span class="pl-s1">init_method</span> <span class="pl-c1">=</span> <span class="pl-s">f'file:///home/cadams/.torch_ddp_init'</span>
<span class="pl-s1">dist</span>.<span class="pl-en">init_process_group</span>(<span class="pl-s">"nccl"</span>, <span class="pl-s1">init_method</span><span class="pl-c1">=</span><span class="pl-s1">init_method</span>,<span class="pl-s1">rank</span><span class="pl-c1">=</span><span class="pl-s1">rank</span>, <span class="pl-s1">world_size</span><span class="pl-c1">=</span><span class="pl-s1">world_size</span>)

<span class="pl-s1">n</span> <span class="pl-c1">=</span> <span class="pl-v">DDP</span>(<span class="pl-s1">n</span>)

<span class="pl-en">print</span>(<span class="pl-s">"SUCCESS"</span>)</pre></div>
<p dir="auto">I launch this with <code class="notranslate">mpirun -n 2 python reproducer.py</code></p>
<p dir="auto">For each rank, I see this:</p>
<div class="highlight highlight-source-shell notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="Traceback (most recent call last):
  File &quot;pytorch_mpi4py_ddp_crash.py&quot;, line 25, in &lt;module&gt;
    n = DDP(n)
  File &quot;/opt/conda/lib/python3.6/site-packages/torch/nn/parallel/distributed.py&quot;, line 389, in __init__
    self.broadcast_bucket_size)
  File &quot;/opt/conda/lib/python3.6/site-packages/torch/nn/parallel/distributed.py&quot;, line 606, in _distributed_broadcast_coalesced
    dist._broadcast_coalesced(self.process_group, tensors, buffer_size)
RuntimeError: NCCL error in: ../torch/lib/c10d/ProcessGroupNCCL.cpp:558, invalid usage, NCCL version 2.7.8"><pre class="notranslate">Traceback (most recent call last):
  File <span class="pl-s"><span class="pl-pds">"</span>pytorch_mpi4py_ddp_crash.py<span class="pl-pds">"</span></span>, line 25, <span class="pl-k">in</span> <span class="pl-k">&lt;</span>module<span class="pl-k">&gt;</span>
    n = DDP(n)
  File <span class="pl-s"><span class="pl-pds">"</span>/opt/conda/lib/python3.6/site-packages/torch/nn/parallel/distributed.py<span class="pl-pds">"</span></span>, line 389, <span class="pl-k">in</span> __init__
    self.broadcast_bucket_size)
  File <span class="pl-s"><span class="pl-pds">"</span>/opt/conda/lib/python3.6/site-packages/torch/nn/parallel/distributed.py<span class="pl-pds">"</span></span>, line 606, <span class="pl-k">in</span> _distributed_broadcast_coalesced
    dist._broadcast_coalesced(self.process_group, tensors, buffer_size)
RuntimeError: NCCL error in: ../torch/lib/c10d/ProcessGroupNCCL.cpp:558, invalid usage, NCCL version 2.7.8</pre></div>
<p dir="auto">Is it possible that the mpi4py call to MPI_Init is being duplicated in pytorch? I know this was a bug in Summer 2020 but not sure if it persists.</p>
<h2 dir="auto">Expected behavior</h2>
<p dir="auto">I expect DDP to successfully init and wrap the networks, one per rank.</p>
<h2 dir="auto">Environment</h2>
<p dir="auto">Collecting environment information...<br>
PyTorch version: 1.7.0a0+8deb4fe<br>
Is debug build: False<br>
CUDA used to build PyTorch: 11.0<br>
ROCM used to build PyTorch: N/A</p>
<p dir="auto">OS: Ubuntu 18.04.5 LTS (x86_64)<br>
GCC version: (Ubuntu 7.5.0-3ubuntu1~18.04) 7.5.0<br>
Clang version: Could not collect<br>
CMake version: version 3.14.0</p>
<p dir="auto">Python version: 3.6 (64-bit runtime)<br>
Is CUDA available: True<br>
CUDA runtime version: Could not collect<br>
GPU models and configuration:<br>
GPU 0: A100-SXM4-40GB<br>
GPU 1: A100-SXM4-40GB<br>
GPU 2: A100-SXM4-40GB<br>
GPU 3: A100-SXM4-40GB<br>
GPU 4: A100-SXM4-40GB<br>
GPU 5: A100-SXM4-40GB<br>
GPU 6: A100-SXM4-40GB<br>
GPU 7: A100-SXM4-40GB</p>
<p dir="auto">Nvidia driver version: 450.51.05<br>
cuDNN version: Probably one of the following:<br>
/usr/lib/x86_64-linux-gnu/libcudnn.so.8.0.4<br>
/usr/lib/x86_64-linux-gnu/libcudnn_adv_infer.so.8.0.4<br>
/usr/lib/x86_64-linux-gnu/libcudnn_adv_train.so.8.0.4<br>
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_infer.so.8.0.4<br>
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_train.so.8.0.4<br>
/usr/lib/x86_64-linux-gnu/libcudnn_ops_infer.so.8.0.4<br>
/usr/lib/x86_64-linux-gnu/libcudnn_ops_train.so.8.0.4<br>
HIP runtime version: N/A<br>
MIOpen runtime version: N/A</p>
<p dir="auto">Versions of relevant libraries:<br>
[pip3] numpy==1.19.1<br>
[pip3] pytorch-transformers==1.1.0<br>
[pip3] torch==1.7.0a0+8deb4fe<br>
[pip3] torch-cluster==1.5.7<br>
[pip3] torch-geometric==1.6.1<br>
[pip3] torch-scatter==2.0.5<br>
[pip3] torch-sparse==0.6.7<br>
[pip3] torchtext==0.6.0<br>
[pip3] torchvision==0.8.0a0<br>
[conda] magma-cuda110             2.5.2                         5    local<br>
[conda] mkl                       2019.1                      144<br>
[conda] mkl-include               2019.1                      144<br>
[conda] nomkl                     3.0                           0<br>
[conda] numpy                     1.19.1           py36h30dfecb_0<br>
[conda] numpy-base                1.19.1           py36h75fe3a5_0<br>
[conda] pytorch-transformers      1.1.0                    pypi_0    pypi<br>
[conda] torch                     1.7.0a0+8deb4fe          pypi_0    pypi<br>
[conda] torchtext                 0.6.0                    pypi_0    pypi<br>
[conda] torchvision               0.8.0a0                  pypi_0    pypi</p>
<h2 dir="auto">Additional context</h2>
<p dir="auto">I am running this on 8xA100 DGX systems, using nvidia's  container: docker pull nvcr.io/nvidia/pytorch:20.09-py3</p>