fix multi-gpu training

Author: yawenzzzz

rslp/lightning_cli.py

@@ -3,11 +3,12 @@
 import os
 
 import jsonargparse
-import wandb
 from lightning.pytorch.callbacks import Callback
+from lightning.pytorch.utilities import rank_zero_only
 from rslearn.main import RslearnLightningCLI
 from upath import UPath
 
+import wandb
 from rslp import launcher_lib
 
 CHECKPOINT_DIR = "gs://{rslp_bucket}/projects/{project_id}/{experiment_id}/checkpoints/"
@@ -26,6 +27,7 @@ def __init__(self, project_id: str, experiment_id: str):
         self.project_id = project_id
         self.experiment_id = experiment_id
 
+    @rank_zero_only
     def on_fit_start(self, trainer, pl_module):
         """Called just before fit starts I think.
 
@@ -102,6 +104,16 @@ def before_instantiate_classes(self):
             c.trainer.logger.init_args.project = c.rslp_project
             c.trainer.logger.init_args.name = c.rslp_experiment
 
+            # Configure DDP strategy with find_unused_parameters=True
+            c.trainer.strategy = jsonargparse.Namespace(
+                {
+                    "class_path": "lightning.pytorch.strategies.DDPStrategy",
+                    "init_args": jsonargparse.Namespace(
+                        {"find_unused_parameters": True}
+                    ),
+                }
+            )
+
             # Set the checkpoint directory to canonical GCS location.
             checkpoint_callback = None
             upload_wandb_callback = None