Repository: pytorch
Issue: LayerNorm is very slow (almost frozen) in CPU of multiprocessing
Labels: module: cpu,triaged

The context of the use case is doing ES over small network in multiprocessing. 

It turns out that with `LayerNorm` it becomes extremely slow on CPU.

To reproduce the effect, the code is attached below:

```python
import os
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.nn.utils import vector_to_parameters
from concurrent.futures import ProcessPoolExecutor
import time
import numpy as np


global use_ln
use_ln = True
global device
device = 'cpu'

class Model(nn.Module):
    def __init__(self, use_ln):
        super().__init__()
        
        self.use_ln = use_ln
        
        self.fc1 = nn.Linear(17, 64)
        if use_ln:
            self.ln1 = nn.LayerNorm(64)
        self.fc2 = nn.Linear(64, 64)
        if use_ln:
            self.ln2 = nn.LayerNorm(64)
        self.out = nn.Linear(64, 6)
    
    def forward(self, x):
        if self.use_ln:
            x = self.ln1(F.relu(self.fc1(x)))
            x = self.ln2(F.relu(self.fc2(x)))
        else:
            x = F.relu(self.fc1(x))
            x = F.relu(self.fc2(x))
        x = torch.tanh(self.out(x))
        return x


def initializer():
    print(f'Initializing, PID: {os.getpid()}')
    t = time.perf_counter()
    global model
    model = Model(use_ln).to(device)
    print(f'Finish initializing, PID: {os.getpid()}, taken {time.perf_counter() - t} s')
    
def fitness(solution):
    vector_to_parameters(solution.to(device), model.parameters())
    for i in range(10):
        for t in range(1000):
            observation = np.random.randn(1, 17)
            action = model(torch.from_numpy(observation).float().unsqueeze(0).to(device))
    print(f'PID: {os.getpid()}, finished')
        
        
with ProcessPoolExecutor(max_workers=16, initializer=initializer) as executor:
    t = time.perf_counter()
    num_param = sum([param.numel() for param in Model(use_ln).parameters() if param.requires_grad])
    solutions = [torch.randn(num_param) for _ in range(16)]
    list(executor.map(fitness, solutions))
    print(f'Total time: {time.perf_counter() - t}')
```
The benchmark is:

- **`use_ln=False, device='cpu'`:**
```
Initializing, PID: 2009192
Finish initializing, PID: 2009192, taken 0.001776786521077156 s
Initializing, PID: 2009194
Initializing, PID: 2009196
Finish initializing, PID: 2009194, taken 0.001952311024069786 s
Finish initializing, PID: 2009196, taken 0.0018356721848249435 s
Initializing, PID: 2009203
Finish initializing, PID: 2009203, taken 0.0015176795423030853 s
Initializing, PID: 2009210
Finish initializing, PID: 2009210, taken 0.0021672528237104416 s
Initializing, PID: 2009216
Initializing, PID: 2009226
Finish initializing, PID: 2009216, taken 0.0018634404987096786 s
Finish initializing, PID: 2009226, taken 0.001307731494307518 s
Initializing, PID: 2009232
Initializing, PID: 2009237
Finish initializing, PID: 2009232, taken 0.0021293386816978455 s
Finish initializing, PID: 2009237, taken 0.0014764349907636642 s
Initializing, PID: 2009242
Finish initializing, PID: 2009242, taken 0.0015203002840280533 s
Initializing, PID: 2009253
Finish initializing, PID: 2009253, taken 0.0014950130134820938 s
Initializing, PID: 2009259
Initializing, PID: 2009262
Finish initializing, PID: 2009259, taken 0.002156071364879608 s
Initializing, PID: 2009269
Finish initializing, PID: 2009272, taken 0.0021638404577970505 s
Finish initializing, PID: 2009262, taken 0.0012727994471788406 s
Finish initializing, PID: 2009269, taken 0.0021071340888738632 s
Initializing, PID: 2009272
Initializing, PID: 2009280
Finish initializing, PID: 2009280, taken 0.0012570079416036606 s
PID: 2009259, finished
PID: 2009232, finished
PID: 2009192, finished
PID: 2009242, finished
PID: 2009216, finished
PID: 2009196, finished
PID: 2009280, finished
PID: 2009203, finished
PID: 2009210, finished
PID: 2009262, finished
PID: 2009226, finished
PID: 2009253, finished
PID: 2009272, finished
PID: 2009237, finished
PID: 2009194, finished
PID: 2009269, finished
Total time: 3.3697756361216307
```
- **`use_ln=True, device='cpu'`:**
```
Initializing, PID: 2013343
Initializing, PID: 2013346
Finish initializing, PID: 2013343, taken 0.002532508224248886 s
Initializing, PID: 2013354
Initializing, PID: 2013359
Finish initializing, PID: 2013346, taken 0.001596800982952118 s
Finish initializing, PID: 2013354, taken 0.002076219767332077 s
Finish initializing, PID: 2013359, taken 0.0015829745680093765 s
Initializing, PID: 2013364
Finish initializing, PID: 2013364, taken 0.0018675383180379868 s
Initializing, PID: 2013369
Initializing, PID: 2013372
Finish initializing, PID: 2013369, taken 0.00201280415058136 s
Finish initializing, PID: 2013372, taken 0.0018042325973510742 s
Initializing, PID: 2013379
Finish initializing, PID: 2013379, taken 0.0017306245863437653 s
Initializing, PID: 2013384
Finish initializing, PID: 2013384, taken 0.0018676463514566422 s
Initializing, PID: 2013389
Initializing, PID: 2013392
Finish initializing, PID: 2013389, taken 0.0027825143188238144 s
Finish initializing, PID: 2013392, taken 0.0016419757157564163 s
Initializing, PID: 2013397
Finish initializing, PID: 2013397, taken 0.0018332768231630325 s
Initializing, PID: 2013402
Finish initializing, PID: 2013402, taken 0.0018747877329587936 s
Initializing, PID: 2013409
Finish initializing, PID: 2013409, taken 0.0016566626727581024 s
Initializing, PID: 2013414
Finish initializing, PID: 2013414, taken 0.0016422048211097717 s
Initializing, PID: 2013419
Finish initializing, PID: 2013419, taken 0.0017681997269392014 s

...
```
**Completely frozen after waiting a few minutes, and `htop` shows all 80 cores CPU are 100% busy**


- **`use_ln=False, device='cuda'`:**
```
Initializing, PID: 2010102
Initializing, PID: 2010103
Initializing, PID: 2010108
Initializing, PID: 2010111
Initializing, PID: 2010114
Initializing, PID: 2010117
Initializing, PID: 2010119
Initializing, PID: 2010122
Initializing, PID: 2010126
Initializing, PID: 2010129
Initializing, PID: 2010132
Initializing, PID: 2010135
Initializing, PID: 2010138
Initializing, PID: 2010141
Initializing, PID: 2010144
Initializing, PID: 2010147
Finish initializing, PID: 2010111, taken 17.613565219566226 s
Finish initializing, PID: 2010102, taken 17.802420677617192 s
Finish initializing, PID: 2010114, taken 17.798802657052875 s
Finish initializing, PID: 2010103, taken 17.81860247440636 s
Finish initializing, PID: 2010117, taken 17.89303245767951 s
Finish initializing, PID: 2010108, taken 17.97427663579583 s
Finish initializing, PID: 2010132, taken 18.15790667384863 s
Finish initializing, PID: 2010126, taken 18.16597461514175 s
Finish initializing, PID: 2010141, taken 18.14934244006872 s
Finish initializing, PID: 2010122, taken 18.205005967989564 s
Finish initializing, PID: 2010119, taken 18.209256762638688 s
Finish initializing, PID: 2010129, taken 18.243990190327168 s
Finish initializing, PID: 2010147, taken 18.25442030467093 s
Finish initializing, PID: 2010135, taken 18.283051615580916 s
Finish initializing, PID: 2010138, taken 18.324560744687915 s
Finish initializing, PID: 2010144, taken 18.545221898704767 s
PID: 2010111, finished
PID: 2010114, finished
PID: 2010102, finished
PID: 2010103, finished
PID: 2010117, finished
PID: 2010108, finished
PID: 2010126, finished
PID: 2010122, finished
PID: 2010141, finished
PID: 2010119, finished
PID: 2010132, finished
PID: 2010147, finished
PID: 2010129, finished
PID: 2010135, finished
PID: 2010138, finished
PID: 2010144, finished
Total time: 35.38426893763244
```

- **`use_ln=True, device='cuda'`:**
```
Initializing, PID: 2010657
Initializing, PID: 2010658
Initializing, PID: 2010663
Initializing, PID: 2010666
Initializing, PID: 2010669
Initializing, PID: 2010670
Initializing, PID: 2010675
Initializing, PID: 2010678
Initializing, PID: 2010679
Initializing, PID: 2010684
Initializing, PID: 2010685
Initializing, PID: 2010690
Initializing, PID: 2010693
Initializing, PID: 2010694
Initializing, PID: 2010698
Initializing, PID: 2010702
Finish initializing, PID: 2010657, taken 16.785767970606685 s
Finish initializing, PID: 2010658, taken 17.48567765392363 s
Finish initializing, PID: 2010663, taken 17.629254495725036 s
Finish initializing, PID: 2010679, taken 17.669835798442364 s
Finish initializing, PID: 2010684, taken 17.806519005447626 s
Finish initializing, PID: 2010666, taken 17.94813333079219 s
Finish initializing, PID: 2010678, taken 17.9285822045058 s
Finish initializing, PID: 2010675, taken 17.933384394273162 s
Finish initializing, PID: 2010685, taken 17.93839507550001 s
Finish initializing, PID: 2010669, taken 17.989747492596507 s
Finish initializing, PID: 2010670, taken 18.010178688913584 s
Finish initializing, PID: 2010690, taken 18.07465230114758 s
Finish initializing, PID: 2010698, taken 18.061042606830597 s
Finish initializing, PID: 2010693, taken 18.082385893911123 s
Finish initializing, PID: 2010702, taken 18.17208438925445 s
Finish initializing, PID: 2010694, taken 18.283266445621848 s
PID: 2010657, finished
PID: 2010658, finished
PID: 2010663, finished
PID: 2010679, finished
PID: 2010684, finished
PID: 2010678, finished
PID: 2010675, finished
PID: 2010685, finished
PID: 2010670, finished
PID: 2010669, finished
PID: 2010666, finished
PID: 2010693, finished
PID: 2010698, finished
PID: 2010690, finished
PID: 2010702, finished
PID: 2010694, finished
Total time: 39.30256709083915
```


```
 - PyTorch Version: 1.0.0.dev20190323
 - OS: Ubuntu 18.04
 - How you installed PyTorch: pip
 - Python version: 3.7
```