Repository: huggingface/transformers

Issue:
Multiprocessing CUDA issues when importing transformers
## Environment info
- `transformers` version: 4.2.0dev0
- Platform: Linux-4.15.0-128-generic-x86_64-with-Ubuntu-18.04-bionic
- Python version: 3.6.9
- PyTorch version (GPU?): 1.6.0 (True)
- Tensorflow version (GPU?): not installed (NA)
- Using GPU in script?: Yes
- Using distributed or parallel set-up in script?: Partly


### Who can help
@stas00

## Information

When using multiprocessing, importing the transformers package causes 
`RuntimeError: Cannot re-initialize CUDA in forked subprocess. To use CUDA with multiprocessing, you must use the 'spawn' start method`

The problem arises when using:
* `import transformers` is used in the main process.
This is due to the [following line](https://github.com/huggingface/transformers/blob/master/src/transformers/models/fsmt/modeling_fsmt.py#L268) in `modeling_fsmt.py`,  removing `torch.cuda.is_available()` call resolves the issue.



## To reproduce

```
import multiprocessing
import transformers  # NOQA
# You can also call torch.cuda instead of import transformers to get the same error
import torch.nn as nn


class Net(nn.Module):

    def __init__(self):
        super().__init__()
        self.linear = nn.Linear(10, 10)

    def forward(self, x):
        return self.linear(x)


def to_cuda(i):
    net = Net().cuda()
    print(f'Called {i} process')


try:
    cpus = multiprocessing.cpu_count()
except NotImplementedError:
    cpus = 2  # arbitrary default

pool = multiprocessing.Pool(processes=cpus)
pool.map(to_cuda, range(10))
```

## Expected behavior
The code snippet above runs without issues.




Discussion:
Hey @maxjeblick,

We do need to be able to call `torch.cuda` in our model code. We cannot "forbid" calls to `torch.cuda` in order to allow quite specific use cases where `.cuda()` is mixed with multiprocessing. IMO, using multiprocessing in combination with `.cuda()` is quite an edge case. Are you trying to run your model on multiple GPUs? 

I'd suggest to fork the repo and delete this line if you really need this feature.
Hey @patrickvonplaten thanks for the fast response! 

I agree that using multiprocessing with `.cuda` calls is not that common, one usecase would be multi-gpu inference without DDP.  The `torch.cuda.is_available()` call in `modeling_fsmt.py` is currently the only place which causes the `RuntimeError`; all other `.cuda` calls are either fine (e.g.  `from torch.cuda.amp import autocast` in `trainer.py`) or not executed during the `import transformers` statement.
It affects bart as well. 


I see ProphetNetLayerNorm solved it by using a runtime wrapper.

https://github.com/huggingface/transformers/blob/912f6881d2b69f180522172a5283702bd8c41d9c/src/transformers/models/prophetnet/modeling_prophetnet.py#L513-L521

So it's modeling_bart and modeling_fsmt are the only 2 that have this check at import time

I wonder if we can just remove this check altogether. Won't `from apex.normalization import FusedLayerNorm` fail w/o cuda? And then we actually don't need that check. We need to verify that.
heh, it looks like `FusedProphetNetLayerNorm` importing is silently failing, since it doesn't exist in `apex` ;) looks like untested code. Fix proposed at https://github.com/huggingface/transformers/pull/9349

Hmm, it works just fine without cuda:

```
CUDA_VISIBLE_DEVICES="" python -c "from apex.normalization import FusedLayerNorm; print(FusedLayerNorm(10))"
FusedLayerNorm(torch.Size([10]), eps=1e-05, elementwise_affine=True)
```

Do you know why did we need that check in first place?

The doc page https://nvidia.github.io/apex/layernorm.html doesn't say anything about needing cuda.
@t-vi figured it out - `apex.normalization.FusedLayerNorm` falls back on to non-cuda gracefully: https://github.com/NVIDIA/apex/blob/master/apex/normalization/fused_layer_norm.py#L154 so the `if torch.cuda.is_available()` check is not needed in first place.
> needed

Nice, so we could actually remove this import statement then from all `FusedLayerNorm` classes no? 

Resolving patch:
diff --git a/src/transformers/models/bart/modeling_bart.py b/src/transformers/models/bart/modeling_bart.py
--- a/src/transformers/models/bart/modeling_bart.py
+++ b/src/transformers/models/bart/modeling_bart.py
@@ -110,13 +110,12 @@ def _expand_mask(mask: torch.Tensor, dtype: torch.dtype, tgt_len: Optional[int]
 
 
 def BartLayerNorm(normalized_shape: torch.Size, eps: float = 1e-5, elementwise_affine: bool = True):
-    if torch.cuda.is_available():
-        try:
-            from apex.normalization import FusedLayerNorm
+    try:
+        from apex.normalization import FusedLayerNorm
 
-            return FusedLayerNorm(normalized_shape, eps, elementwise_affine)
-        except ImportError:
-            pass
+        return FusedLayerNorm(normalized_shape, eps, elementwise_affine)
+    except ImportError:
+        pass
     return torch.nn.LayerNorm(normalized_shape, eps, elementwise_affine)
 
 
diff --git a/src/transformers/models/fsmt/modeling_fsmt.py b/src/transformers/models/fsmt/modeling_fsmt.py
--- a/src/transformers/models/fsmt/modeling_fsmt.py
+++ b/src/transformers/models/fsmt/modeling_fsmt.py
@@ -265,14 +265,12 @@
 
 
 have_fused_layer_norm = False
-if torch.cuda.is_available():
-    try:
-        from apex.normalization import FusedLayerNorm
-
-        have_fused_layer_norm = True
-    except ImportError:
-        pass
+try:
+    from apex.normalization import FusedLayerNorm
 
+    have_fused_layer_norm = True
+except ImportError:
+    pass
 LayerNorm = FusedLayerNorm if have_fused_layer_norm else torch.nn.LayerNorm
 
 
diff --git a/src/transformers/models/prophetnet/modeling_prophetnet.py b/src/transformers/models/prophetnet/modeling_prophetnet.py
--- a/src/transformers/models/prophetnet/modeling_prophetnet.py
+++ b/src/transformers/models/prophetnet/modeling_prophetnet.py
@@ -511,13 +511,12 @@ class ProphetNetDecoderLMOutput(ModelOutput):
 
 
 def ProphetNetLayerNorm(normalized_shape, eps=1e-5, elementwise_affine=True):
-    if torch.cuda.is_available():
-        try:
-            from apex.normalization import FusedLayerNorm
+    try:
+        from apex.normalization import FusedLayerNorm
 
-            return FusedLayerNorm(normalized_shape, eps, elementwise_affine)
-        except ImportError:
-            pass
+        return FusedLayerNorm(normalized_shape, eps, elementwise_affine)
+    except ImportError:
+        pass
     return torch.nn.LayerNorm(normalized_shape, eps, elementwise_affine)
 
 


Test patch: