Adds torch.cuda.set_device calls to DDP examples (pytorch#1142)

Add set_device calls to DDP examples
mpc001 · May 15, 2023 · 79ef786 · 79ef786
1 parent 6a64939
commit 79ef786
Show file tree

Hide file tree

Showing 4 changed files with 4 additions and 0 deletions.
diff --git a/distributed/ddp-tutorial-series/multigpu.py b/distributed/ddp-tutorial-series/multigpu.py
@@ -19,6 +19,7 @@ def ddp_setup(rank, world_size):
     os.environ["MASTER_ADDR"] = "localhost"
     os.environ["MASTER_PORT"] = "12355"
     init_process_group(backend="nccl", rank=rank, world_size=world_size)
+    torch.cuda.set_device(rank)
 
 class Trainer:
     def __init__(

diff --git a/distributed/ddp-tutorial-series/multigpu_torchrun.py b/distributed/ddp-tutorial-series/multigpu_torchrun.py
@@ -12,6 +12,7 @@
 
 def ddp_setup():
     init_process_group(backend="nccl")
+    torch.cuda.set_device(int(os.environ["LOCAL_RANK"]))
 
 class Trainer:
     def __init__(

diff --git a/distributed/ddp-tutorial-series/multinode.py b/distributed/ddp-tutorial-series/multinode.py
@@ -12,6 +12,7 @@
 
 def ddp_setup():
     init_process_group(backend="nccl")
+    torch.cuda.set_device(int(os.environ["LOCAL_RANK"]))
 
 class Trainer:
     def __init__(

diff --git a/distributed/minGPT-ddp/mingpt/main.py b/distributed/minGPT-ddp/mingpt/main.py
@@ -8,6 +8,7 @@
 
 def ddp_setup():
     init_process_group(backend="nccl")
+    torch.cuda.set_device(int(os.environ["LOCAL_RANK"]))
 
 def get_train_objs(gpt_cfg: GPTConfig, opt_cfg: OptimizerConfig, data_cfg: DataConfig):
     dataset = CharDataset(data_cfg)