Fix Multi-GPU Seed Problem (#220)

* fix multigpu

* fix multi-gpu seeding

---------

Co-authored-by: Gordon Guocheng Qian 钱国成 <guocheng.qian@outlook.com>

Author: bennyguo

launch.py

@@ -79,6 +79,7 @@ def main(args, extras) -> None:
         ProgressCallback,
     )
     from threestudio.utils.config import ExperimentConfig, load_config
+    from threestudio.utils.misc import get_rank
     from threestudio.utils.typing import Optional
 
     logger = logging.getLogger("pytorch_lightning")
@@ -97,7 +98,8 @@ def main(args, extras) -> None:
     cfg: ExperimentConfig
     cfg = load_config(args.config, cli_args=extras, n_gpus=n_gpus)
 
-    pl.seed_everything(cfg.seed)
+    # set a different seed for each device
+    pl.seed_everything(cfg.seed + get_rank(), workers=True)
 
     dm = threestudio.find(cfg.data_type)(cfg.data)
     system: BaseSystem = threestudio.find(cfg.system_type)(

threestudio/models/geometry/implicit_sdf.py

@@ -10,7 +10,7 @@
 from threestudio.models.geometry.base import BaseImplicitGeometry, contract_to_unisphere
 from threestudio.models.mesh import Mesh
 from threestudio.models.networks import get_encoding, get_mlp
-from threestudio.utils.misc import get_rank
+from threestudio.utils.misc import broadcast, get_rank
 from threestudio.utils.typing import *
 
 
@@ -209,6 +209,10 @@ def func(points_rand: Float[Tensor, "N 3"]) -> Float[Tensor, "N 1"]:
             loss.backward()
             optim.step()
 
+        # explicit broadcast to ensure param consistency across ranks
+        for param in self.parameters():
+            broadcast(param, src=0)
+
     def get_shifted_sdf(
         self, points: Float[Tensor, "*N Di"], sdf: Float[Tensor, "*N 1"]
     ) -> Float[Tensor, "*N 1"]:

threestudio/utils/misc.py

@@ -110,3 +110,11 @@ def barrier():
         return
     else:
         torch.distributed.barrier()
+
+
+def broadcast(tensor, src=0):
+    if not _distributed_available():
+        return tensor
+    else:
+        torch.distributed.broadcast(tensor, src=src)
+        return tensor