Repository: huggingface/accelerate
Pull request: Support for multi-GPU in notebook_launcher
This PR adds support for multi-GPU training in the `notebook_launcher` and cleans a bit the API.

The same constraints as for TPU are applicable: no defining `accelerator` outside of the training function **and** CUDA has to be uninitiliazed until inside the trainnig function or the forking process is not happy.

Fixes #36

Problem statement:
Multi GPU Training Not Working
While using Accelerate, it is only utilizing 1 out of the 2 GPUs present. I am training using the general instructions in the repository. The architecture is AutoEncoder. 

```
dataloader = DataLoader(dataset, batch_size = 2048, shuffle=True, pin_memory=False, num_workers=20)
encoder = Encoder(bottleneck_size = 2, embedding_size = 40, vocab = dataset.vocab).to(device)
decoder = Decoder(bottleneck_size = 2, embedding_size = 40, vocab = dataset.vocab).to(device)
model = AutoEncoder(encoder, decoder).to(device)
loss = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters())
model, optimizer, dataloader = accelerator.prepare(model, optimizer, dataloader)
```
I am transferring the samples in the batch to the device using the code below:
```
    for x in batch:
        batch[x] = batch[x].to(device)
```
The device is being determined by using:

```device = accelerator.device```

Both devices are visible which can be confirmed by using ```torch.cuda.device_count()``` which returns 2.

Devices are RTX 2080 with CUDA Version 11.2. Driver version is 460.67.
Distro is PopOS!.

Hints from issue discussion:
Did you set "num_processes" in your config to 2?

I noticed, that when this parameter is set to 1 only one gpu will be used.

My working configuration is this:

```
{
  "distributed_type": "MULTI_GPU",
  "fp16": true,
  "machine_rank": 0,
  "main_process_ip": null,
  "main_process_port": null,
  "main_training_function": "main",
  "num_machines": 1,
  "num_processes": 2
}
```
Could you please share the command you are using to launch your script? That would help debug your problem. Thanks!
I am using this with JupyterLab, is it only working while training with scripts or while using the CLI tool?
Yes, the accelerate library currently only supports launching training scripts. Notebook launchers are on the roadmap, but not implemented yet.
Hi- Thanks for the great library, Sylvain!
Not to hijack the thread, but I am having the same problem. (Happy to create a new issue, though).

The config file looks as follows:
```
compute_environment: LOCAL_MACHINE
distributed_type: MULTI_GPU
fp16: true
machine_rank: 0
main_process_ip: null
main_process_port: null
main_training_function: main
num_machines: 1
num_processes: 2
````

The relevant part of the code is as follows:
```
    accelerator = Accelerator(fp16=config['fp16'], cpu=config['cpu'])
    print(accelerator.device)

    # Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
    lr = config["lr"]
    num_epochs = int(config["num_epochs"])
    seed = int(config["seed"])
    batch_size = int(config["batch_size"])

    # If the batch size is too big we use gradient accumulation
    gradient_accumulation_steps = 1
    if batch_size > MAX_GPU_BATCH_SIZE:
        gradient_accumulation_steps = batch_size // MAX_GPU_BATCH_SIZE
        batch_size = MAX_GPU_BATCH_SIZE

    # Instantiate dataloaders.
    train_dataloader = DataLoader(
        train_dataset, shuffle=True, collate_fn=collate_fn, batch_size=batch_size
    )
    valid_dataloader = DataLoader(
        validation_dataset, shuffle=False, collate_fn=collate_fn, batch_size=EVAL_BATCH_SIZE
    )
    test_dataloader = DataLoader(
        test_dataset, shuffle=False, collate_fn=collate_fn, batch_size=EVAL_BATCH_SIZE
    )

    # Instantiate the model (we build the model here so that the seed also control new weights initialization)
    model = Wav2Vec2ForCTC.from_pretrained("facebook/wav2vec2-base-960h")


    # Instantiate optimizer
    optimizer = AdamW(params=model.parameters(), lr=lr)

    # Prepare everything
    # There is no specific order to remember, we just need to unpack the objects in the same order we gave them to the
    # prepare method.
    prepared = accelerator.prepare(
        model, optimizer, train_dataloader, valid_dataloader, test_dataloader
    )
    model, optimizer, train_dataloader, valid_dataloader, test_dataloader = prepared


    # Now we train the model
    for epoch in range(num_epochs):
        model.train()
        for step, batch in enumerate(train_dataloader):
            # We could avoid this line since we set the accelerator with `device_placement=True`.
            #batch.to(accelerator.device)
            outputs = model(**batch)
            loss = outputs.loss
            loss = loss / gradient_accumulation_steps
            accelerator.backward(loss)
            if step % gradient_accumulation_steps == 0:
                optimizer.step()
                lr_scheduler.step()
                optimizer.zero_grad()

```

The script utilizes a single GPU, though there are 2 GPUS.

```
>>> torch.cuda.device_count()
2
```
Launching the scipt in the command line:
```
accelerate launch training.py
```

Any help is appreciated. Thank you!
Yes, a new issue would be cleaner as this one is about accelerate not working out of thee box in a Jupyter environment :-) 

Could you report in your issue what gets printed (since you have a `print(accelerator.device)` in your script)? Thanks!
Should I close this issue then?
I'll work on a notebook launcher in the coming days. You can close the issue now or when it's ready, as you prefer :-)

Commit message: Address review comment

Patch to review:
diff --git a/README.md b/README.md
index 280473ee565..6e2211143f5 100644
--- a/README.md
+++ b/README.md
@@ -163,8 +163,6 @@ notebook_launcher(training_function)
 
 An example can be found in [this notebook](https://github.com/huggingface/notebooks/blob/master/examples/accelerate/simple_nlp_example.ipynb). [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/huggingface/notebooks/blob/master/examples/accelerate/simple_nlp_example.ipynb)
 
-Note that this launcher does not work in Jupyter Notebook on a machine with multiple GPUs (yet). This feature will be released in a later version.
-
 ## Why should I use 🤗 Accelerate?
 
 You should use 🤗 Accelerate when you want to easily run your training scripts in a distributed environment without having to renounce full control over your training loop. This is not a high-level framework above PyTorch, just a thin wrapper so you don't have to learn a new library, In fact the whole API of 🤗 Accelerate is in one class, the `Accelerator` object.
diff --git a/docs/source/launcher.rst b/docs/source/launcher.rst
index da28ad9201a..1c227994626 100644
--- a/docs/source/launcher.rst
+++ b/docs/source/launcher.rst
@@ -15,16 +15,15 @@ Notebook Launcher
 =======================================================================================================================
 
 Launch your training function inside a notebook. Currently supports launching a training with TPUs on [Google
-Colab](https://colab.research.google.com/) and [Kaggle kernels](https://www.kaggle.com/code), or training on one GPU,
-but support for training on several GPUs (if the machine on which you are running your notebook has them) is planned
-for a future release.
+Colab](https://colab.research.google.com/) and [Kaggle kernels](https://www.kaggle.com/code), as well as training on
+several GPUs (if the machine on which you are running your notebook has them).
 
 An example can be found in `this notebook
 <https://github.com/huggingface/notebooks/blob/master/examples/accelerate/simple_nlp_example.ipynb>`__.
 
 .. warning::
 
-    If you are training on Colab or a Kaggle kernel with TPUs, your :obj:`Accelerator` object should only be defined
-    inside the training function. This is because the initialization should be done inside the launcher only.
+    Your :obj:`Accelerator` object should only be defined inside the training function. This is because the
+    initialization should be done inside the launcher only.
 
 .. autofunction:: accelerate.notebook_launcher
diff --git a/docs/source/quicktour.rst b/docs/source/quicktour.rst
index d9846b77725..1e28fed85cc 100644
--- a/docs/source/quicktour.rst
+++ b/docs/source/quicktour.rst
@@ -199,9 +199,8 @@ Launching training from a notebook
 -----------------------------------------------------------------------------------------------------------------------
 
 In Accelerate 0.3.0, a new :class:`~accelerate.notebook_launcher` has been introduced to help you launch your training
-function from a notebook. Currently supports launching a training with TPUs on Colab, or training on one GPU, but
-support for training on several GPUs (if the machine on which you are running your notebook has them) is planned for a
-future release.
+function from a notebook. This launcher supports launching a training with TPUs on Colab or Kaggle, as well as training
+on several GPUs (if the machine on which you are running your notebook has them).
 
 Just define a function responsible for your whole training and/or evaluation in a cell of the notebook, then execute a
 cell with the following code:
@@ -214,8 +213,8 @@ cell with the following code:
 
 .. warning::
 
-    If you are training on Colab with TPUs, your :obj:`Accelerator` object should only be defined inside the training
-    function. This is because the initialization should be done inside the launcher only.
+    Your :obj:`Accelerator` object should only be defined inside the training function. This is because the
+    initialization should be done inside the launcher only.
 
 
 Training on TPU
diff --git a/src/accelerate/notebook_launcher.py b/src/accelerate/notebook_launcher.py
index e9dce7677e6..11ccd91b743 100644
--- a/src/accelerate/notebook_launcher.py
+++ b/src/accelerate/notebook_launcher.py
@@ -14,15 +14,15 @@
 
 import os
 import sys
-import warnings
 
 import torch
+from torch.multiprocessing import start_processes
 
 from .state import AcceleratorState
 from .utils import PrepareForLaunch
 
 
-def notebook_launcher(function, args=(), num_processes=None, start_method="fork", **kwargs):
+def notebook_launcher(function, args=(), num_processes=None, use_fp16=False, use_port="29500"):
     """
     Launches a training function, using several processes if it's possible in the current environment (TPU with
     multiple cores for instance).
@@ -32,17 +32,15 @@ def notebook_launcher(function, args=(), num_processes=None, start_method="fork"
             The training function to execute. If it accepts arguments, the first argument should be the index of the
             process run.
         args (:obj:`Tuple`):
-            Tuple of arguments to pass to the function (it will receive :obj:`(index, *args)`).
+            Tuple of arguments to pass to the function (it will receive :obj:`*args`).
         num_processes (:obj:`int`, `optional`):
             The number of processes to use for training. Will default to 8 in Colab/Kaggle if a TPU is available, to
             the number of GPUs available otherwise.
-
-    .. warning::
-
-        Multiple GPUs is not yet supported.
+        use_fp16 (:obj:`bool`, `optional`, defaults to :obj:`False`):
+            If :obj:`True`, will use mixed precision training on multi-GPU.
+        use_port (:obj:`str`, `optional`, defaults to :obj:`"29500"`):
+            The port to use to communicate between processes when launching a multi-GPU training.
     """
-    launcher = PrepareForLaunch(function)
-
     # Are we in a google colab or a Kaggle Kernel?
     if "IPython" in sys.modules:
         in_colab_or_kaggle = "google.colab" in str(sys.modules["IPython"].get_ipython())
@@ -65,15 +63,60 @@ def notebook_launcher(function, args=(), num_processes=None, start_method="fork"
             if num_processes is None:
                 num_processes = 8
 
+            launcher = PrepareForLaunch(function, distributed_type="TPU")
+            print(f"Launching a training on {num_processes} TPU cores.")
             xmp.spawn(launcher, args=args, nprocs=num_processes, start_method="fork")
         else:
             # No need for a distributed launch otherwise as it's either CPU or one GPU.
-            launcher(0, *args)
+            if torch.cuda.is_available():
+                print("Launching training on one GPU.")
+            else:
+                print("Launching training on CPU.")
+            function(*args)
 
     else:
         if num_processes is None:
-            num_processes = torch.cuda.device_count() if torch.cuda.is_available() else 1
+            raise ValueError(
+                "You have to specify the number of GPUs you would like to use, add `num_process=...` to your call."
+            )
 
         if num_processes > 1:
-            warnings.warn("`notebook_launcher` does not support multiple GPUs yet, launching the training on one GPU.")
-        launcher(0, *args)
+            # Multi-GPU launch
+            if len(AcceleratorState._shared_state) > 0:
+                raise ValueError(
+                    "To launch a multi-GPU training from your notebook, the `Accelerator` should only be initialized "
+                    "inside your training function. Restart your notebook and make sure no cells initializes an "
+                    "`Accelerator`."
+                )
+
+            if torch.cuda.is_initialized():
+                raise ValueError(
+                    "To launch a multi-GPU training from your notebook, you need to avoid running any instruction "
+                    "using `torch.cuda` in any cell. Restart your notebook and make sure no cells use any CUDA "
+                    "function."
+                )
+
+            # torch.distributed will expect a few environment variable to be here. We set the ones common to each
+            # process here (the other ones will be set be the launcher).
+            os.environ["WORLD_SIZE"] = str(num_processes)
+            os.environ["MASTER_ADDR"] = "127.0.0.1"
+            os.environ["MASTER_PORT"] = use_port
+            os.environ["USE_FP16"] = str(use_fp16)
+
+            launcher = PrepareForLaunch(function, distributed_type="MULTI_GPU")
+            try:
+                print(f"Launching a training on {num_processes} GPUs.")
+                start_processes(launcher, nprocs=num_processes, start_method="fork")
+            finally:
+                # Clean up the environment variables set.
+                del os.environ["WORLD_SIZE"]
+                del os.environ["MASTER_ADDR"]
+                del os.environ["MASTER_PORT"]
+
+        else:
+            # No need for a distributed launch otherwise as it's either CPU or one GPU.
+            if torch.cuda.is_available():
+                print("Launching training on one GPU.")
+            else:
+                print("Launching training on CPU.")
+            function(*args)
diff --git a/src/accelerate/utils.py b/src/accelerate/utils.py
index 0c2282b2c00..5dfe85e5401 100644
--- a/src/accelerate/utils.py
+++ b/src/accelerate/utils.py
@@ -13,7 +13,7 @@
 # limitations under the License.
 
 import importlib
-import inspect
+import os
 import random
 from enum import Enum
 from typing import List, Optional, Union
@@ -252,15 +252,23 @@ def save(obj, f):
 
 class PrepareForLaunch:
     """
-    Prepare a function that launches a script.
+    Prepare a function that will launched in a distributed setup.
+
+    Args:
+        launcher (:obj:`Callable`):
+            The function to launch.
+        distributed_type (:class:`~accelerate.state.DistributedType`):
+            The distributed type to prepare for.
     """
 
-    def __init__(self, launcher):
+    def __init__(self, launcher, distributed_type="NO"):
         self.launcher = launcher
+        self.distributed_type = DistributedType(distributed_type)
+
+    def __call__(self, index, *args):
+        if self.distributed_type == DistributedType.MULTI_GPU:
+            # Prepare the environment for torch.distributed
+            os.environ["LOCAL_RANK"] = str(index)
+            os.environ["RANK"] = str(index)
 
-    def __call__(self, index):
-        launcher_sig = inspect.signature(self.launcher)
-        if len(launcher_sig.parameters) == 0:
-            self.launcher()
-        else:
-            self.launcher(index)
+        self.launcher(*args)


Code review comments:
diff_hunk: @@ -65,15 +63,60 @@ def notebook_launcher(function, args=(), num_processes=None, start_method="fork"
             if num_processes is None:
                 num_processes = 8
 
+            launcher = PrepareForLaunch(function, distributed_type="TPU")
+            print(f"Launching a training on {num_processes} TPU cores.")
             xmp.spawn(launcher, args=args, nprocs=num_processes, start_method="fork")
         else:
             # No need for a distributed launch otherwise as it's either CPU or one GPU.
-            launcher(0, *args)
+            if torch.cuda.is_available():
+                print("Launching training on one GPU.")
+            else:
+                print("Launching training on CPU.")
+            function(*args)
 
     else:
         if num_processes is None:
-            num_processes = torch.cuda.device_count() if torch.cuda.is_available() else 1
+            raise ValueError(
+                "You have to specify the number of GPUs you would like to use, add `num_process=...` to your call."
+            )
 
         if num_processes > 1:
-            warnings.warn("`notebook_launcher` does not support multiple GPUs yet, launching the training on one GPU.")
-        launcher(0, *args)
+            # Multi-GPU launch
+            if len(AcceleratorState._shared_state) > 0:
+                raise ValueError(
+                    "To launch a multi-GPU training from your notebook, the `Accelerator` should only be initialized "
+                    "inside your training function. Restart your notebook and make sure no cells initializes an "
+                    "`Accelerator`."
+                )
+
+            if torch.cuda.is_initialized():
+                raise ValueError(
+                    "To launch a multi-GPU training from your notebook, you need to avoid running any instruction "
+                    "using `torch.cuda` in any cell. Restart your notebook and make sure no cells use any CUDA "
+                    "function."
+                )
+
+            # torch.distributed will expect a few environment variable to be here. We set the ones common to each
+            # process here (the other ones will be set be the launcher).
+            os.environ["WORLD_SIZE"] = str(num_processes)
+            os.environ["MASTER_ADDR"] = "127.0.0.1"
+            os.environ["MASTER_PORT"] = use_port
path: src/accelerate/notebook_launcher.py
text: @user1:
```suggestion
            os.environ["MASTER_PORT"] = str(use_port)
```
Would add the extra conversion, might be else error prone