Repository: pytorch
Issue: Training slows down and memory usage increases when upgrading from PyTorch 1.6 to 1.7
Labels: module: performance,triaged,module: regression,module: amp (automated mixed precision)

## 🐛 Bug

[Original Post]
I noticed a significant slowdown of my training script this morning after upgrading from 1.6 to 1.7. GPU memory usage also has a noticeable increase. 

## To Reproduce

~~I will come back and add a MWE to reproduce this later.~~

[Edited on 11/20/2020]
First, we need to prepare two conda environments
```
conda create -n pt1.6 pytorch=1.6 torchvision torchaudio cudatoolkit=10.2 -c pytorch
conda create -n pt1.7 pytorch torchvision torchaudio cudatoolkit=11 -c pytorch
```

 Following the distributed training by [pycls](https://github.com/facebookresearch/pycls), wrote a short MNIST training similar to [PyTorch Examples](https://github.com/pytorch/examples/tree/master/mnist) as below:

```
"""train.py"""
import logging
from pathlib import Path
import time
import multiprocessing
import os
import random
import signal
import threading
import traceback
from typing import Callable
from typing import List
from typing import Optional

import torch
import torch.optim as optim
from torch import nn
from torch.nn import functional as F
from torch.optim.lr_scheduler import StepLR
from torchvision import datasets
from torchvision import transforms


################################################################################
# Parameters
################################################################################

name = Path(__file__).stem
num_gpus = 2
auto_resume = True
checkpoint_last = Path("checkpoint_" + name + ".pth")
num_epochs = 5
logger = logging.getLogger(__name__)
logging.basicConfig(level=logging.INFO)


################################################################################
# Distributed
################################################################################


def is_master_proc() -> bool:
    """Determines if the current process is the master process.

    Returns:
        Master process flag. True if the current process is the master process.
        False otherwise.

    Master process is responsible for logging, writing and loading checkpoints.
    In the multi GPU setting, we assign the master role to the rank 0 process.

    """
    return torch.distributed.get_rank() == 0


def init_process_group(proc_rank: int, world_size: int, host: str, port: int, backend: str) -> None:
    """Initializes the default process group.

    Args:
        proc_rank: Rank of the process, e.g. 0, 1, 2, ...
        world_size: Total number of processes.
        host: Hostname, e.g. "localhost"
        prot: Port number, e.g. 10000.
        backend: Distribution backend.

    """
    # Set the GPU to use
    err_msg = "Cannot use more GPU devices than available."
    if proc_rank >= torch.cuda.device_count():
        raise ValueError(err_msg)
    torch.cuda.set_device(proc_rank)
    # Initialize the process group
    torch.distributed.init_process_group(
        backend=backend, init_method="tcp://{}:{}".format(host, port), world_size=world_size, rank=proc_rank
    )


def destroy_process_group() -> None:
    """Destroys the default process group."""
    torch.distributed.destroy_process_group()


def all_reduce(tensors: List[torch.Tensor]) -> List[torch.Tensor]:
    """Performs the scaled all_reduce operation on the provided tensors.

    Args:
        tensors: List of tensors needs to be synced between GPUs.

    Returns:
        List of tensors that have been synced between GPUs.

    The input tensors are modified in-place. Currently supports only the sum
    reduction operator. The reduced values needs to be scaled by the number of
    GPUs outside this function.

    """

    # Queue the reductions
    reductions = []
    for tensor in tensors:
        reduction = torch.distributed.all_reduce(tensor, async_op=True)
        reductions.append(reduction)
    # Wait for reductions to finish
    for reduction in reductions:
        reduction.wait()
    return tensors


def scaled_all_reduce(tensors: List[torch.Tensor]) -> List[torch.Tensor]:
    """Performs the scaled all_reduce operation on the provided tensors.

    Args:
        tensors: List of tensors needs to be synced between GPUs.

    Returns:
        List of tensors that have been synced between GPUs.

    The input tensors are modified in-place. Currently supports only the sum
    reduction operator. The reduced values needs to be scaled by the number of
    GPUs outside this function.

    """

    num_gpus = float(torch.distributed.get_world_size())
    # Queue the reductions
    reductions = []
    for tensor in tensors:
        reduction = torch.distributed.all_reduce(tensor, async_op=True)
        reductions.append(reduction)
    # Wait for reductions to finish
    for reduction in reductions:
        reduction.wait()
    # Scale the results
    for tensor in tensors:
        tensor.mul_(1.0 / num_gpus)
    return tensors


class ChildException(Exception):
    """Wraps an exception from a child process."""


class ErrorHandler:
    """Multiprocessing error handler (based on fairseq's).

    Listens for errors in child processes and propagates the tracebacks to the
    parent.

    """

    def __init__(self, error_queue):
        # Shared error queue
        self.error_queue = error_queue
        # Children processes sharing the error queue
        self.children_pids = []
        # Start a thread listening to errors
        self.error_listener = threading.Thread(target=self.listen, daemon=True)
        self.error_listener.start()
        # Register the signal handler
        signal.signal(signal.SIGUSR1, self.signal_handler)

    def add_child(self, pid):
        """Registers a child process."""
        self.children_pids.append(pid)

    def listen(self):
        """Listens for errors in the error queue."""
        # Wait until there is an error in the queue
        child_trace = self.error_queue.get()
        # Put the error back for the signal handler
        self.error_queue.put(child_trace)
        # Invoke the signal handler
        os.kill(os.getpid(), signal.SIGUSR1)

    def signal_handler(self, _sig_num, _stack_frame):
        """Signal handler."""
        # Kill children processes
        for pid in self.children_pids:
            os.kill(pid, signal.SIGINT)
        # Propagate the error from the child process
        raise ChildException(self.error_queue.get())


def run(
    proc_rank: int,
    world_size: int,
    host: str,
    port: int,
    backend: str,
    error_queue,
    fun: Callable,
    fun_args: tuple,
    fun_kwargs: dict,
) -> None:
    """Runs a function from a child process.

    Args:
        proc_rank: Process rank.
        world_size: Total number of process.
        host: Hostname.
        port: Port number.
        backend: Distribution backend.
        error_queue: ErrorQueue.
        fun: Function that will be executed.
        fun_args: Positional arguments to the function call.
        fun_kwargs: Keyword arguments to the function call.

    """
    try:
        # Initialize the process group
        init_process_group(proc_rank, world_size, host, port, backend)
        # Run the function
        fun(*fun_args, **fun_kwargs)
    except KeyboardInterrupt:
        # Killed by the parent process
        pass
    except Exception:
        # Propagate exception to the parent process
        error_queue.put(traceback.format_exc())
    finally:
        # Destroy the process group
        destroy_process_group()


def multi_proc_run(
    num_proc: int,
    backend: str,
    host: str,
    port_range: List[int],
    fun: Callable,
    fun_args: tuple = (),
    fun_kwargs: Optional[dict] = None,
):
    """Runs a function in a multi-proc setting (unless num_proc == 1).

    Args:
        num_proc: Total number of process.
        backend: Distribution backend.
        host: Hostname.
        port_range: Range of port numbers that will be used to randomly
            generate port number.
        fun: Function that will be executeda.
        fun_args: Positional arguments to the function call.
        fun_kwargs: Keyword arguments to the function call.

    """
    # There is no need for multi-proc in the single-proc case
    fun_kwargs = fun_kwargs if fun_kwargs else {}
    if num_proc == 1:
        fun(*fun_args, **fun_kwargs)
        return
    # Handle errors from training subprocesses
    error_queue = multiprocessing.SimpleQueue()
    error_handler = ErrorHandler(error_queue)
    # Get a random port to use (without using global random number generator)
    port = random.Random().randint(port_range[0], port_range[1])
    # Run each training subprocess
    processes = []
    for i in range(num_proc):
        p_i = multiprocessing.Process(
            target=run, args=(i, num_proc, host, port, backend, error_queue, fun, fun_args, fun_kwargs)
        )
        processes.append(p_i)
        p_i.start()
        error_handler.add_child(p_i.pid)
    # Wait for each subprocess to finish
    for process in processes:
        process.join()


################################################################################
# Model
################################################################################


class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 32, 3, 1)
        self.conv2 = nn.Conv2d(32, 64, 3, 1)
        self.dropout1 = nn.Dropout(0.25)
        self.dropout2 = nn.Dropout(0.5)
        self.fc1 = nn.Linear(9216, 128)
        self.fc2 = nn.Linear(128, 10)

    def forward(self, x):
        x = self.conv1(x)
        x = F.relu(x)
        x = self.conv2(x)
        x = F.relu(x)
        x = F.max_pool2d(x, 2)
        x = self.dropout1(x)
        x = torch.flatten(x, 1)
        x = self.fc1(x)
        x = F.relu(x)
        x = self.dropout2(x)
        x = self.fc2(x)
        output = F.log_softmax(x, dim=1)
        return output


################################################################################
# Training
################################################################################
def train_model():

    # dataloaders
    dataset_train = datasets.MNIST(
        root="data/",
        train=True,
        download=True,
        transform=transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))]),
    )
    loader_train = torch.utils.data.DataLoader(
        dataset=dataset_train,
        batch_size=16,
        shuffle=False,
        sampler=torch.utils.data.distributed.DistributedSampler(dataset_train),
        num_workers=32,
        pin_memory=True,
        drop_last=False,
    )
    dataset_val = datasets.MNIST(
        root="data/",
        train=True,
        download=True,
        transform=transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))]),
    )
    loader_val = torch.utils.data.DataLoader(
        dataset=dataset_val,
        batch_size=250,
        shuffle=False,
        sampler=torch.utils.data.distributed.DistributedSampler(dataset_val),
        num_workers=8,
        pin_memory=True,
        drop_last=False,
    )

    # model, loss, optimizer, scheduler
    cur_device = torch.cuda.current_device()
    model = Net().cuda(device=cur_device)
    model = torch.nn.parallel.DistributedDataParallel(module=model, device_ids=[cur_device], output_device=cur_device)
    loss_fun = nn.NLLLoss()
    optimizer = optim.Adadelta(model.parameters(), lr=1.0)
    scheduler = StepLR(optimizer, step_size=1, gamma=0.7)
    scaler = torch.cuda.amp.GradScaler()

    start_epoch = 0

    # training loop
    for cur_epoch in range(start_epoch, num_epochs):

        # train
        model.train()
        start_time = time.time()
        for cur_iter, (inputs, labels) in enumerate(loader_train):
            inputs, labels = inputs.cuda(), labels.cuda(non_blocking=True)
            optimizer.zero_grad()
            with torch.cuda.amp.autocast(enabled=True):
                outputs = model(inputs)
                loss = loss_fun(outputs, labels)
            scaler.scale(loss).backward()
            scaler.step(optimizer)
            scaler.update()
            loss = scaled_all_reduce([loss])[0]
        elapsed_time = time.time() - start_time

        # validation
        test_correct = 0
        test_total = 0

        model.eval()
        for inputs, labels in loader_val:
            # transfer the data to the current GPU device
            if next(model.parameters()).device.type == "cuda":
                inputs, labels = inputs.cuda(), labels.cuda(non_blocking=True)

            with torch.cuda.amp.autocast():
                outputs = model(inputs).detach()

            # count correct predictions
            _, predictions = torch.max(outputs.data, 1)
            batch_total = torch.Tensor([labels.shape[0]]).cuda()
            batch_correct = (predictions == labels).sum()
            batch_total, batch_correct = all_reduce([batch_total, batch_correct])
            batch_total, batch_correct = batch_total.item(), batch_correct.item()

            # update meters
            test_total += batch_total
            test_correct += batch_correct

        accuracy = test_correct / test_total

        # report
        if is_master_proc():
            logger.info(
                "Epoch %3d time %5.1f lr = %.3f avg_loss = %.4f acc = %.1f",
                cur_epoch + 1,
                elapsed_time,
                scheduler.get_last_lr()[0],
                loss,                
                accuracy * 100,
            )
        scheduler.step()


if __name__ == "__main__":
    multi_proc_run(4, "nccl", "localhost", [10000, 65000], train_model)
```

Steps to reproduce the behavior:

1. `conda activate pt1.6`
1. `python train.py`
1. `conda activate pt1.7`
1. `python train.py`

<!-- If you have a code sample, error messages, stack traces, please provide it here as well -->

The results on AWS `p3.8xlarge` are shown below 
```
❯ conda activate pt1.6
(pt1.6)
❯ python example/train.py
INFO:__main__:Epoch   1 time   7.9 lr = 1.000 avg_loss = 0.0538 acc = 98.5
INFO:__main__:Epoch   2 time   7.5 lr = 0.700 avg_loss = 0.0160 acc = 99.2
INFO:__main__:Epoch   3 time   7.6 lr = 0.490 avg_loss = 0.0509 acc = 99.3
INFO:__main__:Epoch   4 time   7.5 lr = 0.343 avg_loss = 0.0133 acc = 99.5
INFO:__main__:Epoch   5 time   7.6 lr = 0.240 avg_loss = 0.0077 acc = 99.6
(pt1.6)
❯ conda activate pt1.7
(pt1.7)
❯ python example/train.py
INFO:root:Reducer buckets have been rebuilt in this iteration.
INFO:root:Reducer buckets have been rebuilt in this iteration.
INFO:root:Reducer buckets have been rebuilt in this iteration.
INFO:root:Reducer buckets have been rebuilt in this iteration.
INFO:__main__:Epoch   1 time   9.7 lr = 1.000 avg_loss = 0.0544 acc = 98.6
INFO:__main__:Epoch   2 time   9.2 lr = 0.700 avg_loss = 0.0109 acc = 99.2
INFO:__main__:Epoch   3 time   9.2 lr = 0.490 avg_loss = 0.0855 acc = 99.4
INFO:__main__:Epoch   4 time   9.2 lr = 0.343 avg_loss = 0.0031 acc = 99.5
INFO:__main__:Epoch   5 time   9.2 lr = 0.240 avg_loss = 0.0070 acc = 99.6

```

## Expected behavior

<!-- A clear and concise description of what you expected to happen. -->
The difference in training times between PyTorch 1.6 and 1.7 should not be this much. In fact, for larger trainings that we run over multiple days, PyTorch 1.7 takes around 10% more time to finish.

## Environment

As described in the reproduction section. We are comparing two versions of PyTorch, thus two environments are created using conda.



cc @mcarilli @VitalyFedyunin @ngimel @pietern @mrshenli @pritamdamania87 @zhaojuanmao @satgera @rohan-varma @gqchen @aazzolini @osalpekar @jiayisuse @agolynski @SciPioneer @H-Huang @mrzzd