adds barriers after checkpoint saving

non-rank-0 processes were continuing on and waiting
for rank-0 forward, while rank 0 was saving artifacts. This was leading
to collective timeouts. In testing, this reduced the incidence of
timeouts because non-rank-0 processes are waiting at a known wait-point.

Signed-off-by: James Kunstle <[email protected]>

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(