Repository: huggingface/accelerate

Issue:
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!.


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 :-)

Resolving patch:
diff --git a/src/accelerate/notebook_launcher.py b/src/accelerate/notebook_launcher.py
--- 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"] = str(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
--- 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)


Test patch: