Skip to content

Commit

Permalink
add oom observer callback
Browse files Browse the repository at this point in the history
  • Loading branch information
cli99 committed Feb 1, 2024
1 parent 4b82876 commit 59a5a04
Showing 1 changed file with 3 additions and 1 deletion.
4 changes: 3 additions & 1 deletion llmfoundry/utils/builders.py
Original file line number Diff line number Diff line change
@@ -12,7 +12,7 @@
import torch
from composer import algorithms
from composer.callbacks import (EarlyStopper, Generate, LRMonitor,
MemoryMonitor, OptimizerMonitor,
MemoryMonitor, OOMObserver, OptimizerMonitor,
RuntimeEstimator, SpeedMonitor)
from composer.core import Algorithm, Callback, Evaluator
from composer.datasets.in_context_learning_evaluation import \
@@ -166,6 +166,8 @@ def build_callback(
return LRMonitor()
elif name == 'memory_monitor':
return MemoryMonitor()
elif name == 'oom_observer':
return OOMObserver(**kwargs)
elif name == 'speed_monitor':
return SpeedMonitor(window_size=kwargs.get('window_size', 1),
gpu_flops_available=kwargs.get(

0 comments on commit 59a5a04

Please sign in to comment.