added systemMetricsMonitor callback

mosaicml · Jun 6, 2024 · 38d1db6 · 38d1db6
1 parent ff92f3c
commit 38d1db6
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/llmfoundry/callbacks/__init__.py b/llmfoundry/callbacks/__init__.py
@@ -11,6 +11,7 @@
     OptimizerMonitor,
     RuntimeEstimator,
     SpeedMonitor,
+    SystemMetricsMonitor,
 )
 
 from llmfoundry.callbacks.async_eval_callback import AsyncEval
@@ -35,6 +36,7 @@
 from llmfoundry.callbacks.scheduled_gc_callback import ScheduledGarbageCollector
 from llmfoundry.registry import callbacks, callbacks_with_config
 
+callbacks.register('system_metrics_monitor', func=SystemMetricsMonitor)
 callbacks.register('lr_monitor', func=LRMonitor)
 callbacks.register('memory_monitor', func=MemoryMonitor)
 callbacks.register('memory_snapshot', func=MemorySnapshot)
@@ -71,4 +73,4 @@
     'AsyncEval',
     'CurriculumLearning',
     'LossPerpVsContextLengthLogger',
-]
+]