Update monitor config setup

Author: manveerxyz

configs/150M/3090.toml

@@ -1,6 +1,5 @@
 name_model = "150M"
 project = "debug_150m_zero_band"
-metric_logger_type = "dummy"
 
 [train]
 micro_bs = 16 # change this base on the gpu

src/zeroband/train.py

@@ -56,7 +56,6 @@ class TrainConfig(BaseConfig):
 
 
 class MonitorConfig(BaseConfig):
-    enable_monitor: bool = False
     log_flush_interval: int = 10
     base_url: str | None = None
     auth_token: str | None = None
@@ -76,7 +75,7 @@ class Config(BaseConfig):
     data: DataConfig = DataConfig()
     optim: OptimConfig = OptimConfig()
     train: TrainConfig
-    monitor: MonitorConfig = MonitorConfig()
+    monitor: MonitorConfig | None = None
 
 
 def train(config: Config):
@@ -167,8 +166,9 @@ def train(config: Config):
         logger_cls = WandbMetricLogger if config.metric_logger_type == "wandb" else DummyMetricLogger
         metric_logger = logger_cls(project=config.project, config=config.model_dump(), resume=False)
 
-        monitor = HttpMonitor(config=config.model_dump(), resume=False)
-        monitor.set_stage("init")
+        if config.monitor is not None:
+            monitor = HttpMonitor(config=config.model_dump(), resume=False)
+            monitor.set_stage("init")
 
     train_dataloader_iterator = iter(train_dataloader)
 
@@ -182,7 +182,7 @@ def train(config: Config):
             # if we don't use diloco we don't print the outer step logs
             logger.info(f"outer_step step: {outer_step}")
 
-        if world_info.rank == 0:
+        if world_info.rank == 0 and config.monitor is not None:
             monitor.set_stage("inner_loop")
 
         for inner_step in range(num_inner_steps):
@@ -245,12 +245,13 @@ def train(config: Config):
 
             if world_info.rank == 0:
                 metric_logger.log(metrics)
-                monitor.log(metrics)
+                if config.monitor is not None:
+                    monitor.log(metrics)
 
             logger.info(log)
 
         if config.diloco is not None:
-            if world_info.rank == 0:
+            if world_info.rank == 0 and config.monitor is not None:
                 monitor.set_stage("outer_loop")
             diloco.step(model)
 
@@ -263,8 +264,9 @@ def train(config: Config):
             break
 
     if world_info.rank == 0:
-        monitor.finish()
         metric_logger.finish()
+        if config.monitor is not None:
+            monitor.finish()
 
 
 if __name__ == "__main__":

src/zeroband/utils/monitor.py

@@ -25,7 +25,6 @@ class HttpMonitor:
 
     def __init__(self, config, *args, **kwargs):
         self.data = []
-        self.enabled = config["monitor"]["enable_monitor"]
         self.log_flush_interval = config["monitor"]["log_flush_interval"]
         self.base_url = config["monitor"]["base_url"]
         self.auth_token = config["monitor"]["auth_token"]
@@ -50,19 +49,13 @@ def _remove_duplicates(self):
         self.data = unique_logs
 
     def set_stage(self, stage: str):
-        if not self.enabled:
-            return
-
         import time
 
         # add a new log entry with the stage name
         self.data.append({"stage": stage, "time": time.time()})
         self._handle_send_batch(flush=True)  # it's useful to have the most up-to-date stage broadcasted
 
     def log(self, data: dict[str, Any]):
-        if not self.enabled:
-            return
-    
         # Lowercase the keys in the data dictionary
         lowercased_data = {k.lower(): v for k, v in data.items()}
         self.data.append(lowercased_data)
@@ -131,9 +124,6 @@ def _finish(self):
             return False
 
     def finish(self):
-        if not self.enabled:
-            return
-        
         self.set_stage("finishing")
 
         # Send any remaining logs