Set cuda device before init_process_group (#56)

Author: yifuwang

generate.py

@@ -260,7 +260,6 @@ def main(
     rank = maybe_init_dist()
     use_tp = rank is not None
     if use_tp:
-        torch.cuda.set_device(rank)
         if rank != 0:
             # only print on rank 0
             print = lambda *args, **kwargs: None

tp.py

@@ -42,6 +42,7 @@ def maybe_init_dist() -> Optional[int]:
         # not run via torchrun, no-op
         return None
 
+    torch.cuda.set_device(rank)
     dist.init_process_group(backend="nccl", rank=rank, world_size=world_size)
     return rank