feat: add custom logger

Author: samsja

open_diloco/train_fsdp.py

@@ -15,7 +15,6 @@
 
 from pydantic import model_validator
 import torch
-import wandb
 from pydantic_config import parse_argv, BaseConfig
 from datasets import load_dataset
 from datasets.distributed import split_dataset_by_node
@@ -35,7 +34,6 @@
     MixedPrecision,
 )
 from torch.distributed.device_mesh import DeviceMesh
-from torch.distributed import broadcast_object_list
 from open_diloco.ckpt_utils import (
     CKPT_PREFIX,
     CkptConfig,
@@ -46,7 +44,7 @@
     save_checkpoint,
 )
 from open_diloco.hivemind_diloco import AllReduceStrategy, DiLoCoOptimizer
-
+from open_diloco.utils import WandbLogger, DummyLogger
 
 from hivemind.dht.dht import DHT
 from hivemind.utils.networking import log_visible_maddrs
@@ -120,6 +118,7 @@ class Config(BaseConfig):
     precision: Literal["fp16-mixed", "bf16-mixed", "32-true"] = "fp16-mixed"
     # Checkpointing and logging
     project: str = "hivemind_debug"
+    metric_logger_type: Literal["wandb", "dummy"] = "wandb"
     log_activations_steps: int | None = None
     ckpt: CkptConfig = CkptConfig()
     # Hivemind
@@ -192,14 +191,9 @@ def train(config: Config):
         sharding_strategy = ShardingStrategy.NO_SHARD
         log("Hivemind is used, ShardingStrategy.NO_SHARD is used")
 
-    run_id = None
     if rank == 0:
-        wandb.init(project=config.project, config=config.model_dump())
-        run_id = wandb.run.id
-
-    run_id_list = [run_id]
-    broadcast_object_list(run_id_list, src=0)
-    run_id = run_id_list[0]
+        logger_cls = WandbLogger if config.metric_logger_type == "wandb" else DummyLogger
+        metric_logger = logger_cls(project=config.project, config=config.model_dump())
 
     if config.hv is not None:
         log("hivemind diloco enabled")
@@ -459,7 +453,7 @@ def scheduler_fn(opt):
 
                 current_time = time.time()
 
-                wandb.log(metrics)
+                metric_logger.log(metrics)
 
                 if config.hv is None:
                     log(
@@ -512,7 +506,7 @@ def scheduler_fn(opt):
             if config.max_steps is not None and real_step >= config.max_steps:
                 break
     log("Training completed.")
-    wandb.finish()
+    metric_logger.finish()
 
 
 if __name__ == "__main__":

open_diloco/utils.py

@@ -1,11 +1,13 @@
 import hashlib
 from functools import partial
-from typing import Any, Generator
+import json
+from typing import Any, Generator, Protocol
 
 import torch
 from torch.utils.hooks import RemovableHandle
 from torch.distributed.fsdp import ShardingStrategy
 from torch.utils.data import IterableDataset
+import wandb
 
 
 _FSDP_WRAPPED_MODULE = ["_forward_module.", "_fsdp_wrapped_module."]
@@ -175,3 +177,39 @@ def __iter__(self) -> Generator[dict[str, Any], Any, None]:
             input_ids = torch.randint(3, self.vocab_size, (self.seq_len,)).tolist()
             attention_mask = [1] * self.seq_len
             yield {"input_ids": input_ids, "attention_mask": attention_mask}
+
+
+class Logger(Protocol):
+    def __init__(self, project, config): ...
+
+    def log(self, metrics: dict[str, Any]): ...
+
+    def finish(self): ...
+
+
+class WandbLogger:
+    def __init__(self, project, config):
+        wandb.init(project=project, config=config)
+
+    def log(self, metrics: dict[str, Any]):
+        wandb.log(metrics)
+
+    def finish(self):
+        wandb.finish()
+
+
+class DummyLogger:
+    def __init__(self, project, config):
+        self.project = project
+        self.config = config
+        open(project, "a").close()  # Create an empty file at the project path
+
+        self.data = []
+
+    def log(self, metrics: dict[str, Any]):
+        self.data.append(metrics)
+
+    def finish(self):
+        with open(self.project, "a") as f:
+            for d in self.data:
+                f.write(json.dumps(d) + "\n")