feat: use wandb auto resuming feature

open_diloco/train_fsdp.py

@@ -193,9 +193,11 @@ def train(config: Config):
         sharding_strategy = ShardingStrategy.NO_SHARD
         log("Hivemind is used, ShardingStrategy.NO_SHARD is used")
 
+    resume_from_ckpt, resume_path = get_resume_info(config.ckpt)
+
     if rank == 0:
         logger_cls = WandbLogger if config.metric_logger_type == "wandb" else DummyLogger
-        metric_logger = logger_cls(project=config.project, config=config.model_dump())
+        metric_logger = logger_cls(project=config.project, config=config.model_dump(), resume=resume_from_ckpt)
 
     if config.hv is not None:
         log("hivemind diloco enabled")
@@ -257,8 +259,6 @@ def scheduler_fn(opt):
             num_training_steps=config.total_steps,
         )
 
-    resume_from_ckpt, resume_path = get_resume_info(config.ckpt)
-
     if config.hv is not None:
         if resume_from_ckpt:
             # We need to load with a fake optimizer to set the model parameters correctly before initializing the DiLoCoOptimizer
@@ -510,6 +510,7 @@ def scheduler_fn(opt):
 
             if config.max_steps is not None and real_step >= config.max_steps:
                 break
+
     log("Training completed.")
     if rank == 0:
         metric_logger.finish()

open_diloco/utils.py

@@ -176,8 +176,10 @@ def finish(self): ...
 
 
 class WandbLogger:
-    def __init__(self, project, config):
-        wandb.init(project=project, config=config)
+    def __init__(self, project, config, resume: bool):
+        wandb.init(
+            project=project, config=config, resume="auto" if resume else None
+        )  # make wandb reuse the same run id if possible
 
     def log(self, metrics: dict[str, Any]):
         wandb.log(metrics)
@@ -187,7 +189,7 @@ def finish(self):
 
 
 class DummyLogger:
-    def __init__(self, project, config):
+    def __init__(self, project, config, *args, **kwargs):
         self.project = project
         self.config = config
         open(project, "a").close()  # Create an empty file at the project path