Log ZeRO and parallelism settings to monitor backends

Record static training topology on WandB/Comet/TensorBoard/CSV at
engine init so runs are comparable without digging through logs.

Fixes #7494

Signed-off-by: YeonwooSung <neos960518@gmail.com>

Author: YeonwooSung

deepspeed/monitor/comet.py

@@ -70,6 +70,10 @@ def write_events(self, event_list: List[Event]) -> None:
                     step=engine_global_samples,
                 )
 
+    def update_config(self, config_dict):
+        if self._experiment is not None:
+            self._experiment.log_parameters(config_dict)
+
 
 class EventsLogScheduler:
 

deepspeed/monitor/csv_monitor.py

@@ -65,3 +65,14 @@ def write_events(self, event_list):
                         self.filenames.append(filename)
                         csv_monitor_writer.writerow(['step', header])
                     csv_monitor_writer.writerow([step, value])
+
+    def update_config(self, config_dict):
+        if not self.enabled or dist.get_rank() != 0 or self.log_dir is None:
+            return
+        import csv
+        fname = os.path.join(self.log_dir, 'deepspeed_config.csv')
+        with open(fname, 'w') as csv_monitor_file:
+            csv_monitor_writer = csv.writer(csv_monitor_file)
+            csv_monitor_writer.writerow(['key', 'value'])
+            for key, value in config_dict.items():
+                csv_monitor_writer.writerow([key, value])

deepspeed/monitor/monitor.py

@@ -57,3 +57,14 @@ def write_events(self, event_list):
                 self.csv_monitor.write_events(event_list)
             if self.comet_monitor is not None:
                 self.comet_monitor.write_events(event_list)
+
+    def update_config(self, config_dict):
+        if dist.get_rank() == 0:
+            if self.tb_monitor is not None:
+                self.tb_monitor.update_config(config_dict)
+            if self.wandb_monitor is not None:
+                self.wandb_monitor.update_config(config_dict)
+            if self.csv_monitor is not None:
+                self.csv_monitor.update_config(config_dict)
+            if self.comet_monitor is not None:
+                self.comet_monitor.update_config(config_dict)

deepspeed/monitor/tensorboard.py

@@ -5,6 +5,7 @@
 
 from .utils import check_tb_availability
 from .monitor import Monitor
+import json
 import os
 
 import deepspeed.comm as dist
@@ -54,3 +55,7 @@ def write_events(self, event_list, flush=True):
     def flush(self):
         if self.enabled and self.summary_writer is not None and dist.get_rank() == 0:
             self.summary_writer.flush()
+
+    def update_config(self, config_dict):
+        if self.summary_writer is not None:
+            self.summary_writer.add_text("DeepSpeed/config", json.dumps(config_dict, indent=2), global_step=0)

deepspeed/monitor/utils.py

@@ -35,3 +35,71 @@ def check_comet_availability():
     except ImportError:
         print('If you want to use comet logging, please `pip install "comet_ml>=3.41.0"`')
         raise
+
+
+def _json_safe_int(value, default):
+    if value is None:
+        return default
+    return int(value)
+
+
+def _parallel_int(getter, default):
+    # TP/DP groups are only created when those parallelisms are configured.
+    try:
+        return _json_safe_int(getter(), default)
+    except (AssertionError, RuntimeError, ValueError, TypeError):
+        return default
+
+
+def _offload_device(offload_cfg):
+    if offload_cfg is None:
+        return None
+    device = getattr(offload_cfg, "device", None)
+    if device is None:
+        return None
+    device_name = device.value if hasattr(device, "value") else str(device)
+    if device_name == 'none':
+        return None
+    return device_name
+
+
+def _pipeline_parallel_rank(engine):
+    if not getattr(engine, "pipeline_parallelism", False) and getattr(engine, "mpu", None) is None:
+        return 0
+    mpu = getattr(engine, "mpu", None)
+    if mpu is None:
+        return 0
+    if hasattr(mpu, "get_pipeline_model_parallel_rank"):
+        return int(mpu.get_pipeline_model_parallel_rank())
+    if hasattr(mpu, "get_pipe_parallel_rank"):
+        return int(mpu.get_pipe_parallel_rank())
+    return 0
+
+
+def collect_monitor_config(engine):
+    """Build a JSON-safe snapshot of ZeRO / precision / parallelism settings."""
+    from deepspeed.utils import groups
+    from deepspeed.utils.bwc import bwc_pipeline_parallel_world_size
+    import deepspeed.comm as dist
+
+    world_size_default = dist.get_world_size() if dist.is_initialized() else 1
+    rank_default = dist.get_rank() if dist.is_initialized() else 0
+
+    return {
+        'zero_stage': int(engine.zero_optimization_stage()),
+        'offload_optimizer': _offload_device(engine.zero_offload_optimizer()),
+        'offload_param': _offload_device(engine.zero_offload_param()),
+        'fp16': bool(engine.fp16_enabled()),
+        'bf16': bool(engine.bfloat16_enabled()),
+        'train_batch_size': int(engine.train_batch_size()),
+        'train_micro_batch_size_per_gpu': int(engine.train_micro_batch_size_per_gpu()),
+        'gradient_accumulation_steps': int(engine.gradient_accumulation_steps()),
+        'data_parallel_world_size': _parallel_int(groups.get_data_parallel_world_size, world_size_default),
+        'tensor_parallel_world_size': _parallel_int(groups.get_tensor_model_parallel_world_size, 1),
+        'pipeline_parallel_world_size': int(bwc_pipeline_parallel_world_size(engine.mpu)),
+        'sequence_parallel_world_size': int(groups._get_sequence_parallel_world_size()),
+        'data_parallel_rank': _parallel_int(groups.get_data_parallel_rank, rank_default),
+        'tensor_parallel_rank': _parallel_int(groups.get_tensor_model_parallel_rank, 0),
+        'pipeline_parallel_rank': _pipeline_parallel_rank(engine),
+        'sequence_parallel_rank': int(groups._get_sequence_parallel_rank()),
+    }

deepspeed/monitor/wandb.py

@@ -36,3 +36,8 @@ def write_events(self, event_list):
                 value = event[1]
                 step = event[2]
                 self.log({label: value}, step=step)
+
+    def update_config(self, config_dict):
+        if self.enabled and dist.get_rank() == 0:
+            import wandb
+            wandb.config.update(config_dict, allow_val_change=True)

deepspeed/runtime/engine.py

@@ -113,6 +113,7 @@
     STEP_GLOBAL_TIMER
 from deepspeed.utils.debug import debug_extract_module_and_param_names, debug_clear_module_and_param_names
 from deepspeed.monitor.monitor import MonitorMaster
+from deepspeed.monitor.utils import collect_monitor_config
 from deepspeed.runtime.progressive_layer_drop import ProgressiveLayerDrop
 from deepspeed.runtime.utils import clip_grad_norm_, compare_tensors_in_structures, maybe_loss_for_backward
 from deepspeed.runtime.eigenvalue import Eigenvalue
@@ -613,6 +614,9 @@ def __init__(self,
         if self.dist_backend is None:
             self.enable_backward_allreduce = False
 
+        if self.monitor.enabled:
+            self.monitor.update_config(collect_monitor_config(self))
+
     def _optimized_linear_offload_setup(self):
         self.optimized_linear_base_weight_sharding = False
         self.optimized_linear_lora_enabled = False

docs/_tutorials/monitor.md

@@ -42,18 +42,18 @@ When using DeepSpeed for model training, the Monitor can be configured in the De
     "enabled": true,
     "output_path": "output/ds_logs/",
     "job_name": "train_bert"
-  }
+  },
   "wandb": {
     "enabled": true,
     "team": "my_team",
     "group": "my_group",
     "project": "my_project"
-  }
+  },
   "comet": {
     "enabled": true,
     "project": "my_project",
     "experiment_name": "my_experiment"
-  }
+  },
   "csv_monitor": {
     "enabled": true,
     "output_path": "output/ds_logs/",
@@ -62,7 +62,7 @@ When using DeepSpeed for model training, the Monitor can be configured in the De
 }
 ```
 
-DeepSpeed will automatically log to all available and enabled monitoring backends listed in the config, and will generate live monitoring views such as those listed above.
+DeepSpeed will automatically log to all available and enabled monitoring backends listed in the config, and will generate live monitoring views such as those listed above. When the engine is initialized, DeepSpeed also records static parallelism and ZeRO settings such as ZeRO stage, offload devices, precision, and DP/TP/PP/SP sizes on the enabled monitor backends.
 
 ### Custom Monitoring