Merge pull request #276 from allenai/petew/sys-metrics

log sys metrics less often
allenai · Sep 20, 2023 · 2df922b · 2df922b
2 parents 921c254 + a76b3bf
commit 2df922b
Showing 1 changed file with 4 additions and 3 deletions.
diff --git a/olmo/train.py b/olmo/train.py
@@ -711,9 +711,10 @@ def split_batch(self, batch: Dict[str, Any]) -> List[Dict[str, Any]]:
 
     def system_metrics(self) -> Dict[str, float]:
         metrics = {}
-        peak_gpu_mb = peak_gpu_memory()
-        if peak_gpu_mb is not None:
-            metrics["System/Peak GPU Memory (MB)"] = peak_gpu_mb
+        if self.global_step < 3 or self.global_step % 10 == 0:
+            peak_gpu_mb = peak_gpu_memory()
+            if peak_gpu_mb is not None:
+                metrics["System/Peak GPU Memory (MB)"] = peak_gpu_mb
         return metrics
 
     def log_metrics_to_console(self, prefix: str, metrics: Dict[str, float]):