From 617dc90ad946bafc9589f052168b21e455180700 Mon Sep 17 00:00:00 2001 From: Kewen Zhao Date: Thu, 25 Apr 2024 13:55:50 -0400 Subject: [PATCH] refactor scheduler --- logix/logix.py | 4 ++++ logix/scheduler.py | 7 ++++--- 2 files changed, 8 insertions(+), 3 deletions(-) diff --git a/logix/logix.py b/logix/logix.py index aab5cc5d..27f5a134 100644 --- a/logix/logix.py +++ b/logix/logix.py @@ -81,6 +81,10 @@ def __init__( self.type_filter = None self.name_filter = None + # Deferred Imports to avoid circular imports + from logix.scheduler import LogIXScheduler + self.scheduler = LogIXScheduler(self) + def watch( self, model: nn.Module, diff --git a/logix/scheduler.py b/logix/scheduler.py index cb94a3bd..b4c466ce 100644 --- a/logix/scheduler.py +++ b/logix/scheduler.py @@ -7,9 +7,6 @@ class LogIXScheduler: def __init__( self, logix: LogIX, - lora: str = "none", - hessian: str = "none", - save: str = "none", ): self.logix = logix @@ -17,6 +14,10 @@ def __init__( self._lora_epoch = -1 self._logix_state_schedule = [] + lora = logix.config.scheduler.lora + hessian = logix.config.scheduler.hessian + save = logix.config.scheduler.save + self.sanity_check(lora, hessian, save) self.configure_lora_epoch(lora) self.configure_schedule(lora, hessian, save)