Migrating training scripts to torchrun

references/detection/README.md

@@ -27,24 +27,32 @@ python references/detection/train_pytorch.py db_resnet50 --train_path path/to/yo
 
 ### Multi-GPU support (PyTorch only)
 
-Multi-GPU support on Detection task with PyTorch has been added.
-Arguments are the same than the ones from single GPU, except:
-
-- `--devices`: **by default, if you do not pass `--devices`, it will use all GPUs on your computer**.
-You can use specific GPUs by passing a list of ids (ex: `0 1 2`). To find them, you can use the following snippet:
-
-```python
-import torch
-devices = [torch.cuda.device(i) for i in range(torch.cuda.device_count())]
-device_names = [torch.cuda.get_device_name(d) for d in devices]
-```
+We now use the built-in [`torchrun`](https://pytorch.org/docs/stable/elastic/run.html) launcher to spawn your DDP workers. `torchrun` will set all the necessary environment variables (`LOCAL_RANK`, `RANK`, etc.) for you. Arguments are the same than the ones from single GPU, except:
 
 - `--backend`: you can specify another `backend` for `DistribuedDataParallel` if the default one is not available on
 your operating system. Fastest one is `nccl` according to [PyTorch Documentation](https://pytorch.org/docs/stable/generated/torch.nn.parallel.DistributedDataParallel.html).
 
+#### Key `torchrun` parameters:
+- `--nproc_per_node=<N>`
+  Spawn `<N>` processes on the local machine (typically equal to the number of GPUs you want to use).
+- `--nnodes=<M>`
+  (Optional) Total number of nodes in your job. Default is 1.
+- `--rdzv_backend`, `--rdzv_endpoint`, `--rdzv_id`
+  (Optional) Rendezvous settings for multi-node jobs. See the [torchrun docs](https://pytorch.org/docs/stable/elastic/run.html) for details.
+
+#### GPU selection:
+By default all visible GPUs will be used. To limit which GPUs participate, set the `CUDA_VISIBLE_DEVICES` environment variable **before** running `torchrun`. For example, to use only CUDA devices 0 and 2:
+
 ```shell
-python references/detection/train_pytorch_ddp.py db_resnet50 --train_path path/to/your/train_set --val_path path/to/your/val_set --epochs 5 --devices 0 1 --backend nccl
-```
+CUDA_VISIBLE_DEVICES=0,2 \
+torchrun --nproc_per_node=2 references/detection/train_pytorch_ddp.py \
+  db_resnet50 \
+  --train_path path/to/train \
+  --val_path   path/to/val \
+  --epochs 5 \
+  --backend nccl
+  ```
+
 
 ## Data format
 

references/detection/train_pytorch_ddp.py

@@ -18,7 +18,6 @@
 
 # The following import is required for DDP
 import torch.distributed as dist
-import torch.multiprocessing as mp
 import wandb
 from torch.nn.parallel import DistributedDataParallel as DDP
 from torch.optim.lr_scheduler import CosineAnnealingLR, MultiplicativeLR, OneCycleLR, PolynomialLR
@@ -187,13 +186,18 @@
     return val_loss, recall, precision, mean_iou
 
 
-def main(rank: int, world_size: int, args):
+def main(args):
     """
     Args:
         rank (int): device id to put the model on
         world_size (int): number of processes participating in the job
         args: other arguments passed through the CLI
     """
+    # Setup device and distributed
+    rank = int(os.environ.get("LOCAL_RANK", 0))
+    torch.cuda.set_device(rank)
+    dist.init_process_group(backend=args.backend)
+
     slack_token = os.getenv("TQDM_SLACK_TOKEN")
     slack_channel = os.getenv("TQDM_SLACK_CHANNEL")
 
@@ -270,12 +274,11 @@
         model.load_state_dict(checkpoint)
 
     # create default process group
-    device = torch.device("cuda", args.devices[rank])
-    dist.init_process_group(args.backend, rank=rank, world_size=world_size)
+    device = torch.device("cuda", rank)
     # create local model
     model = model.to(device)
     # construct the DDP model
-    model = DDP(model, device_ids=[device])
+    model = DDP(model, device_ids=[rank])
 
     if rank == 0:
         # Metrics
@@ -349,7 +352,7 @@
         batch_size=args.batch_size,
         drop_last=True,
         num_workers=args.workers,
-        sampler=DistributedSampler(train_set, num_replicas=world_size, rank=rank, shuffle=False, drop_last=True),
+        sampler=DistributedSampler(train_set, rank=rank, shuffle=False, drop_last=True),
         pin_memory=torch.cuda.is_available(),
         collate_fn=train_set.collate_fn,
     )
@@ -509,7 +512,7 @@
     )
 
     # DDP related args
-    parser.add_argument("--backend", default="nccl", type=str, help="backend to use for torch DDP")
+    parser.add_argument("--backend", default="nccl", type=str, help="Backend to use for torch.distributed")
 
     parser.add_argument("arch", type=str, help="text-detection model to train")
     parser.add_argument("--output_dir", type=str, default=".", help="path to save checkpoints and final model")
@@ -518,7 +521,6 @@
     parser.add_argument("--name", type=str, default=None, help="Name of your training experiment")
     parser.add_argument("--epochs", type=int, default=10, help="number of epochs to train the model on")
     parser.add_argument("-b", "--batch_size", type=int, default=2, help="batch size for training")
-    parser.add_argument("--devices", default=None, nargs="+", type=int, help="GPU devices to use for training")
     parser.add_argument(
         "--save-interval-epoch", dest="save_interval_epoch", action="store_true", help="Save model every epoch"
     )
@@ -567,14 +569,4 @@
     args = parse_args()
     if not torch.cuda.is_available():
         raise AssertionError("PyTorch cannot access your GPUs. Please investigate!")
-
-    if not isinstance(args.devices, list):
-        args.devices = list(range(torch.cuda.device_count()))
-    # no of process per gpu
-    nprocs = len(args.devices)
-    # Environment variables which need to be
-    # set when using c10d's default "env"
-    # initialization mode.
-    os.environ["MASTER_ADDR"] = "localhost"
-    os.environ["MASTER_PORT"] = "29500"
-    mp.spawn(main, args=(nprocs, args), nprocs=nprocs, join=True)
+    main(args)

references/recognition/README.md

@@ -25,27 +25,36 @@ or PyTorch:
 python references/recognition/train_pytorch.py crnn_vgg16_bn --train_path path/to/your/train_set --val_path path/to/your/val_set --epochs 5
 ```
 
-### Multi-GPU support (PyTorch only - Experimental)
+### Multi-GPU support (PyTorch only )
 
-Multi-GPU support on recognition task with PyTorch has been added. It'll be probably added for other tasks.
-Arguments are the same than the ones from single GPU, except:
-
-- `--devices`: **by default, if you do not pass `--devices`, it will use all GPUs on your computer**.
-You can use specific GPUs by passing a list of ids (ex: `0 1 2`). To find them, you can use the following snippet:
-
-```python
-import torch
-devices = [torch.cuda.device(i) for i in range(torch.cuda.device_count())]
-device_names = [torch.cuda.get_device_name(d) for d in devices]
-```
+We now use the built-in [`torchrun`](https://pytorch.org/docs/stable/elastic/run.html) launcher to spawn your DDP workers. `torchrun` will set all the necessary environment variables (`LOCAL_RANK`, `RANK`, etc.) for you. Arguments are the same than the ones from single GPU, except:
 
 - `--backend`: you can specify another `backend` for `DistribuedDataParallel` if the default one is not available on
 your operating system. Fastest one is `nccl` according to [PyTorch Documentation](https://pytorch.org/docs/stable/generated/torch.nn.parallel.DistributedDataParallel.html).
 
+#### Key `torchrun` parameters:
+- `--nproc_per_node=<N>`
+  Spawn `<N>` processes on the local machine (typically equal to the number of GPUs you want to use).
+- `--nnodes=<M>`
+  (Optional) Total number of nodes in your job. Default is 1.
+- `--rdzv_backend`, `--rdzv_endpoint`, `--rdzv_id`
+  (Optional) Rendezvous settings for multi-node jobs. See the [torchrun docs](https://pytorch.org/docs/stable/elastic/run.html) for details.
+
+#### GPU selection:
+By default all visible GPUs will be used. To limit which GPUs participate, set the `CUDA_VISIBLE_DEVICES` environment variable **before** running `torchrun`. For example, to use only CUDA devices 0 and 2:
+
 ```shell
-python references/recognition/train_pytorch_ddp.py crnn_vgg16_bn --train_path path/to/your/train_set --val_path path/to/your/val_set --epochs 5 --devices 0 1 --backend nccl
+CUDA_VISIBLE_DEVICES=0,2 \
+torchrun --nproc_per_node=2 references/recognition/train_pytorch_ddp.py \
+  crnn_vgg16_bn \
+  --train_path path/to/train \
+  --val_path   path/to/val \
+  --epochs 5 \
+  --backend nccl
 ```
 
+
+
 ## Data format
 
 You need to provide both `train_path` and `val_path` arguments to start training.