add memorysnapshot to callbacks (#810)

Co-authored-by: Daniel King <[email protected]>
mosaicml · Feb 9, 2024 · 491fa7a · 491fa7a
1 parent 12d1ca7
commit 491fa7a
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/llmfoundry/utils/builders.py b/llmfoundry/utils/builders.py
@@ -12,7 +12,7 @@
 import torch
 from composer import algorithms
 from composer.callbacks import (EarlyStopper, Generate, LRMonitor,
-                                MemoryMonitor, OptimizerMonitor,
+                                MemoryMonitor, MemorySnapshot, OptimizerMonitor,
                                 RuntimeEstimator, SpeedMonitor)
 from composer.core import Algorithm, Callback, Evaluator
 from composer.datasets.in_context_learning_evaluation import \
@@ -164,6 +164,8 @@ def build_callback(
         return LRMonitor()
     elif name == 'memory_monitor':
         return MemoryMonitor()
+    elif name == 'memory_snapshot':
+        return MemorySnapshot(**kwargs)
     elif name == 'speed_monitor':
         return SpeedMonitor(window_size=kwargs.get('window_size', 1),
                             gpu_flops_available=kwargs.get(