Merge pull request #78 from fabianlim/checkpoint

Another Fix for Checkpointing
instructlab · Jun 26, 2024 · d9237f8 · d9237f8
2 parents 7fa8e28 + a95abd8
commit d9237f8
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/src/instructlab/training/utils.py b/src/instructlab/training/utils.py
@@ -501,7 +501,8 @@ def ensure_loadable_granite_checkpoint(model_name_or_path: str):
         if dist.is_initialized():
             dist.barrier()
         yield tmpdir
-        shutil.rmtree(tmpdir, ignore_errors=True)
+        if not dist.is_initialized() or dist.get_rank() == 0:
+            shutil.rmtree(tmpdir, ignore_errors=True)
 
 
 # this function is for supporting gradient checkpointing for padding free