rather use inner steps

PrimeIntellect-ai · Oct 29, 2024 · 4927359 · 4927359
1 parent d9d68c8
commit 4927359
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 1 deletion.
diff --git a/src/zeroband/checkpoint.py b/src/zeroband/checkpoint.py
@@ -151,6 +151,7 @@ class RemoteConfig(BaseConfig):
 class CkptConfig(BaseConfig):
     path: str | None = None
     interval: int | None = None
+    interval_inner_step: int | None = None
     topk: int | None = None
 
     remote: RemoteConfig | None = None

diff --git a/src/zeroband/train.py b/src/zeroband/train.py
@@ -446,7 +446,8 @@ def train(config: Config):
                 config.ckpt is not None
                 and training_progress.step > 0
                 and training_progress.step % num_inner_steps != 0
-                and training_progress.step % config.ckpt.interval == 0
+                and config.ckpt.interval_inner_step is not None
+                and training_progress.step % config.ckpt.interval_inner_step == 0
             ):
                 logger.info(
                     f"Saving inner step ckpt at {training_progress.step}. This will only save the inner model and optimizer"