callbacks.py

from .async_eval_callback import AsyncEval
from .curriculum_learning_callback import CurriculumLearning
from .eval_gauntlet_callback import EvalGauntlet
from .fdiff_callback import FDiffMetrics
from .hf_checkpointer import HuggingFaceCheckpointer
from .monolithic_ckpt_callback import MonolithicCheckpointSaver
from .resumption_callbacks import GlobalLRScaling, LayerFreezing
from .scheduled_gc_callback import ScheduledGarbageCollector
from .registry import callbacks, callbacks_with_config
callbacks.register('lr_monitor', func=LRMonitor)
callbacks.register('memory_monitor', func=MemoryMonitor)
callbacks.register('memory_snapshot', func=MemorySnapshot)
callbacks.register('speed_monitor', func=SpeedMonitor)
callbacks.register('runtime_estimator', func=RuntimeEstimator)
callbacks.register('optimizer_monitor', func=OptimizerMonitor)
callbacks.register('generate_callback', func=Generate)
callbacks.register('early_stopper', func=EarlyStopper)
callbacks.register('fdiff_metrics', func=FDiffMetrics)
callbacks.register('hf_checkpointer', func=HuggingFaceCheckpointer)
callbacks.register('global_lr_scaling', func=GlobalLRScaling)
callbacks.register('layer_freezing', func=LayerFreezing)
callbacks.register('mono_checkpoint_saver', func=MonolithicCheckpointSaver)
callbacks.register('scheduled_gc', func=ScheduledGarbageCollector)
callbacks.register('oom_observer', func=OOMObserver)
callbacks_with_config.register('async_eval', func=AsyncEval)
callbacks_with_config.register('curriculum_learning', func=CurriculumLearning)