Merge pull request #566 from JamesKunstle/post-save-barrier

adds barriers after checkpoint saving

Author: JamesKunstle

src/instructlab/training/main_ds.py

@@ -505,6 +505,8 @@ def train(
                     is_lora=bool(args.lora_r),
                     hf_format=True,
                 )
+                base_logger.debug("RANK (%d) waiting at post-save barrier.", local_rank)
+                torch.distributed.barrier()
 
             # if (
             #     args.save_samples_ds is not None
@@ -533,6 +535,8 @@ def train(
                 hf_format=True,
                 epoch=epoch,
             )
+            base_logger.debug("RANK (%d) waiting at post-save barrier.", local_rank)
+            torch.distributed.barrier()
 
     if args.save_last:
         save_hf_format_accelerate(