Revert "disable offloaded optimizer state for live reco bc of potenta…

…il memory leak" This reverts commit 56a3499.
PrimeIntellect-ai · Nov 11, 2024 · f6bccd3 · f6bccd3
1 parent 55714cf
commit f6bccd3
Showing 1 changed file with 1 addition and 2 deletions.
diff --git a/src/zeroband/train.py b/src/zeroband/train.py
@@ -468,8 +468,7 @@ def train(config: Config):
                 monitor.set_stage("outer_loop")
 
             # todo we could skip this is we don't have live recovery enabled
-            # disable because of potential memory leak
-            # ckpt_manager.cache_inner_optimizer()
+            ckpt_manager.cache_inner_optimizer()
 
             time_start_inner = time.perf_counter()
             diloco.step(model=model, flag=training_progress.outer_step, num_effective_peers=num_effective_peers)