Repository: pytorch
Issue: A bug in parallel.data_parallel when module_kwargs is not None
Labels: module: cuda,module: error checking,triaged,module: batching,module: data parallel

## 🐛 Bug

A bug in [parallel.data_parallel](https://github.com/pytorch/pytorch/blob/fdaa77ae8b084eaa7535075538e9d80f4c4a8d1a/torch/nn/parallel/data_parallel.py#L159-L189) when
- The batch size is smaller than the number of GPUs
- With some keyword arguments (module_kwargs) in the custom forward function.

## To Reproduce

Steps to reproduce the behavior:

1. Please execute the test script below (with more than 4 GPUs available).

<!-- If you have a code sample, error messages, stack traces, please provide it here as well -->
```python
import torch
from torch import nn
from torch.nn import parallel

def check(module, inputs, device_ids=None, output_device=None, dim=0, module_kwargs=None):
    '''
    This function is a just simple copy of nn.parallel.data_parallel:
    https://github.com/pytorch/pytorch/blob/fdaa77ae8b084eaa7535075538e9d80f4c4a8d1a/torch/nn/parallel/data_parallel.py#L159-L189
    '''
    if not isinstance(inputs, tuple):
        inputs = (inputs,)

    if device_ids is None:
        device_ids = list(range(torch.cuda.device_count()))

    if output_device is None:
        output_device = device_ids[0]

    inputs, module_kwargs = parallel.scatter_gather.scatter_kwargs(
        inputs, module_kwargs, device_ids, dim
    )
    print(module_kwargs)
    if len(device_ids) == 1:
        return module(*inputs[0], **module_kwargs[0])

    used_device_ids = device_ids[:len(inputs)]
    replicas = parallel.replicate(module, used_device_ids)
    outputs = parallel.parallel_apply(replicas, inputs, module_kwargs, used_device_ids)
    return parallel.gather(outputs, output_device, dim)

class Test(nn.Conv2d):

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

    def forward(self, x, msg='default'):
        print(msg)
        return super().forward(x)

if __name__ == '__main__':
    conv = Test(1, 1, 3, padding=1).cuda()
    x = torch.randn(4, 1, 4, 4).cuda()
    y = torch.randn(1, 1, 4, 4).cuda()

    print('x on 1 GPU with kwargs')
    check(conv, x, device_ids=range(1), module_kwargs={'msg': 'Hello'})
    print('x on 4 GPUs with kwargs')
    check(conv, x, device_ids=range(4), module_kwargs={'msg': 'Hello'})

    print('y on 1 GPU')
    check(conv, y, device_ids=range(1))
    print('y on 1 GPU with kwargs')
    check(conv, y, device_ids=range(1), module_kwargs={'msg': 'Hello'})
    print('y on 4 GPUs')
    check(conv, y, device_ids=range(4))
    print('y on 4 GPUs with kwargs')
    check(conv, y, device_ids=range(4), module_kwargs={'msg': 'Hello'}) # Error!
```

```bash
>>> python test.py
x on 1 GPU with kwargs
({'msg': 'Hello'},)
Hello
x on 4 GPUs with kwargs
({'msg': 'Hello'}, {'msg': 'Hello'}, {'msg': 'Hello'}, {'msg': 'Hello'})
Hello
Hello
Hello
Hello
y on 1 GPU
({},)
default
y on 1 GPU with kwargs
({'msg': 'Hello'},)
Hello
y on 4 GPUs
({},)
default
y on 4 GPUs with kwargs
({'msg': 'Hello'}, {'msg': 'Hello'}, {'msg': 'Hello'}, {'msg': 'Hello'})
Hello
Traceback (most recent call last):
  File "test.py", line 53, in <module>
    check(conv, y, device_ids=range(4), module_kwargs={'msg': 'Hello'})
  File "test.py", line 24, in check
    outputs = parallel.parallel_apply(replicas, inputs, module_kwargs, used_device_ids)
  File "/home/sanghyun/anaconda3/lib/python3.7/site-packages/torch/nn/parallel/parallel_apply.py", line 83, in parallel_apply
    raise output
  File "/home/sanghyun/anaconda3/lib/python3.7/site-packages/torch/nn/parallel/parallel_apply.py", line 59, in _worker
    output = module(*input, **kwargs)
  File "/home/sanghyun/anaconda3/lib/python3.7/site-packages/torch/nn/modules/module.py", line 489, in __call__
    result = self.forward(*input, **kwargs)
TypeError: forward() missing 1 required positional argument: 'x'
```
## Expected behavior

<!-- A clear and concise description of what you expected to happen. -->
``Test.forward()`` should be executed **only once** with the ``y on 4 GPUs with kwargs`` configuration (like ``y on 4 GPUs``). However, it seems that ``Test.forward()`` is called more than once, with an empty ``inputs``. We can see one ``Hello`` message below ``y on 4 GPUs with kwargs``, and I think ``parallel.parallel_apply`` should be terminated right after then.

## Environment

 - PyTorch Version (e.g., 1.0): 1.0
 - OS (e.g., Linux): Ubuntu 16.04
 - How you installed PyTorch (`conda`, `pip`, source): ``conda install pytorch torchvision -c pytorch``
 - Python version: Anaconda 3.7 (https://repo.continuum.io/archive/Anaconda3-2018.12-Linux-x86_64.sh)
 - CUDA/cuDNN version:
```bash
>>> nvcc --version
NVIDIA (R) Cuda compiler driver
Copyright (c) 2005-2017 NVIDIA Corporation
Built on Fri_Sep__1_21:08:03_CDT_2017
Cuda compilation tools, release 9.0, V9.0.176
```
```python
Python 3.7.1 (default, Dec 14 2018, 19:28:38)
[GCC 7.3.0] :: Anaconda, Inc. on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import torch
>>> torch.backends.cudnn.version()
7401
```
 - GPU models and configuration:
```bash
>>> nvidia-smi
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 384.98                 Driver Version: 384.98                    |
|-------------------------------+----------------------+----------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|===============================+======================+======================|
|   0  GeForce GTX 108...  On   | 00000000:04:00.0 Off |                  N/A |
| 33%   57C    P2   106W / 250W |   5691MiB / 11172MiB |     84%      Default |
+-------------------------------+----------------------+----------------------+
|   1  GeForce GTX 108...  On   | 00000000:05:00.0 Off |                  N/A |
| 41%   69C    P2   163W / 250W |   9115MiB / 11172MiB |     95%      Default |
+-------------------------------+----------------------+----------------------+
|   2  GeForce GTX 108...  On   | 00000000:08:00.0 Off |                  N/A |
|  0%   26C    P8    18W / 250W |      0MiB / 11172MiB |      0%      Default |
+-------------------------------+----------------------+----------------------+
|   3  GeForce GTX 108...  On   | 00000000:09:00.0 Off |                  N/A |
| 23%   28C    P8    16W / 250W |      0MiB / 11172MiB |      0%      Default |
+-------------------------------+----------------------+----------------------+
|   4  GeForce GTX 108...  On   | 00000000:83:00.0 Off |                  N/A |
| 23%   20C    P8     7W / 250W |      0MiB / 11172MiB |      0%      Default |
+-------------------------------+----------------------+----------------------+
|   5  GeForce GTX 108...  On   | 00000000:84:00.0 Off |                  N/A |
|  0%   24C    P8     8W / 250W |      0MiB / 11172MiB |      0%      Default |
+-------------------------------+----------------------+----------------------+
|   6  GeForce GTX 108...  On   | 00000000:87:00.0 Off |                  N/A |
| 23%   27C    P8     8W / 250W |      0MiB / 11172MiB |      0%      Default |
+-------------------------------+----------------------+----------------------+
|   7  GeForce GTX 108...  On   | 00000000:88:00.0 Off |                  N/A |
| 23%   21C    P8     8W / 250W |      0MiB / 11172MiB |      0%      Default |
+-------------------------------+----------------------+----------------------+

+-----------------------------------------------------------------------------+
| Processes:                                                       GPU Memory |
|  GPU       PID   Type   Process name                             Usage      |
|=============================================================================|
|    0      5671      C   python                                      5681MiB |
|    1      6349      C   python                                      9105MiB |
+-----------------------------------------------------------------------------+
```
(Used 4, 5, 6, 7 for the test script.)



cc @ngimel