add peak memory stats

src/zeroband/train.py

@@ -21,7 +21,7 @@
 from zeroband.diloco import Diloco, DilocoConfig
 from zeroband.comms import ElasticDeviceMesh
 
-from zeroband.utils import PerfCounter, get_module_signature, get_sharding_strategy
+from zeroband.utils import GPUMemoryMonitor, PerfCounter, get_module_signature, get_sharding_strategy
 from zeroband.utils.monitor import WandbMonitor, DummyMonitor
 from zeroband.data import TEST_VOCAB_SIZE, get_dataloader
 from zeroband.models.llama import get_model
@@ -192,6 +192,8 @@ def train(config: Config):
         logger_cls = WandbMonitor if config.metric_logger_type == "wandb" else DummyMonitor
         metric_logger = logger_cls(project=config.project, config=config.model_dump(), resume=False)
 
+    gpu_mem_monitor = GPUMemoryMonitor()
+
     train_dataloader_iterator = iter(train_dataloader)
 
     num_inner_steps = config.diloco.inner_steps if config.diloco is not None else 1
@@ -254,6 +256,10 @@ def train(config: Config):
                 "Perplexity": torch.exp(loss_batch).item(),
                 "total_tokens": training_progress.total_tokens,
             }
+
+            peak_gpu_stats = gpu_mem_monitor.get_peak_stats()
+            metrics.update(peak_gpu_stats)
+
             log = f"step: {training_progress.step}, loss: {loss_batch.item():.4f}"
 
             tokens_per_second = perf_counter.get_tokens_per_second()
@@ -303,6 +309,8 @@ def train(config: Config):
             mfu = 100 * num_flop_per_token * tokens_per_second / gpu_peak_flops / world_info.local_world_size
             logger.info(f"effective mfu: {mfu}")
 
+        logger.info(f"outer step peak gpu stats: {gpu_mem_monitor.format_peak_states()}")
+
         if training_progress.step >= config.optim.total_steps:
             # we only allow to break outisde of the inner loop.
             # This avoid ending the training in the middle of a the inner loop

src/zeroband/utils/__init__.py

@@ -1,8 +1,11 @@
 import hashlib
 import time
+from typing import Any
 import torch
 from torch.distributed.fsdp import ShardingStrategy
 
+from zeroband.utils.logging import get_logger
+
 
 __all__ = ["get_sharding_strategy", "get_peak_flops", "get_num_flop_per_token", "get_num_params"]
 
@@ -121,3 +124,50 @@ def get_module_signature(module: torch.nn.Module, compress: bool = True) -> str:
         return hashlib.md5(str(state_dict_sig).encode("utf-8")).hexdigest()
     else:
         return "\n".join(f"{name}: {sig}" for name, sig in state_dict_sig.items())
+
+
+class GPUMemoryMonitor:
+    # inspired from https://github.com/pytorch/torchtitan/blob/eef8bb2b1b6f0875ab0581079e1511d51654910e/torchtitan/metrics.py#L32
+    def __init__(self, device: str = "cuda"):
+        self.device = torch.device(device)  # device object
+        self.device_capacity = torch.cuda.get_device_properties(self.device).total_memory
+        self.device_capacity_gib = self._to_gib(self.device_capacity)
+        torch.cuda.reset_peak_memory_stats()
+        torch.cuda.empty_cache()
+
+        self._logger = get_logger()
+
+    def _to_gib(self, memory_in_bytes):
+        # NOTE: GiB (gibibyte) is 1024, vs GB is 1000
+        _gib_in_bytes = 1024 * 1024 * 1024
+        memory_in_gib = memory_in_bytes / _gib_in_bytes
+        return memory_in_gib
+
+    def _to_pct(self, memory):
+        return 100 * memory / self.device_capacity
+
+    def get_peak_stats(self) -> dict[str, Any]:
+        cuda_info = torch.cuda.memory_stats(self.device)
+
+        max_active = cuda_info["active_bytes.all.peak"]
+        max_active_gib = self._to_gib(max_active)
+        max_active_pct = self._to_pct(max_active)
+
+        max_reserved = cuda_info["reserved_bytes.all.peak"]
+        max_reserved_gib = self._to_gib(max_reserved)
+        max_reserved_pct = self._to_pct(max_reserved)
+
+        return {
+            "gpu_max_active_gib": max_active_gib,
+            "gpu_max_active_pct": max_active_pct,
+            "gpu_max_reserved_gib": max_reserved_gib,
+            "gpu_max_reserved_pct": max_reserved_pct,
+        }
+
+    def reset_peak_stats(self):
+        torch.cuda.reset_peak_memory_stats()
+
+    def format_peak_states(self, peak_stats: dict[str, Any] | None = None) -> str:
+        if peak_stats is None:
+            peak_stats = self.get_peak_stats()
+        return f"Active {peak_stats['gpu_max_active_gib']:.2f} GiB ({peak_stats['gpu_max_active_pct']:.2f}%), Reserved {peak_stats['gpu_max_reserved_gib']:.2f} GiB ({peak_stats['gpu_max_reserved_pct']:.2f}%)"