fix problem of 'accelerator.is_main_process' to run in mutiple GPUs (#5340)

fix problem of 'accelerator.is_main_process' to run in mutiple GPUs or NPUs
Co-authored-by: jiaqiw <wangjiaqi50@huawei.com>

examples/unconditional_image_generation/train_unconditional.py

@@ -607,6 +607,7 @@ def main(args):
                 progress_bar.update(1)
                 global_step += 1
 
+                if accelerator.is_main_process:
                     if global_step % args.checkpointing_steps == 0:
                         # _before_ saving state, check if this save would set us over the `checkpoints_total_limit`
                         if args.checkpoints_total_limit is not None:
@@ -628,7 +629,6 @@ def main(args):
                                     removing_checkpoint = os.path.join(args.output_dir, removing_checkpoint)
                                     shutil.rmtree(removing_checkpoint)
 
-                    if accelerator.is_main_process:
                         save_path = os.path.join(args.output_dir, f"checkpoint-{global_step}")
                         accelerator.save_state(save_path)
                         logger.info(f"Saved state to {save_path}")