| import os |
| import shutil |
|
|
|
|
| def save_checkpoint(accelerator, args, global_step, logger): |
| """Save model checkpoint and enforce checkpoints_total_limit.""" |
| 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}") |
|
|
| if args.checkpoints_total_limit is not None: |
| checkpoints = os.listdir(args.output_dir) |
| checkpoints = [d for d in checkpoints if d.startswith("checkpoint")] |
| checkpoints = sorted(checkpoints, key=lambda x: int(x.split("-")[1])) |
|
|
| if len(checkpoints) >= args.checkpoints_total_limit: |
| num_to_remove = len(checkpoints) - args.checkpoints_total_limit + 1 |
| removing_checkpoints = checkpoints[:num_to_remove] |
|
|
| logger.info( |
| f"{len(checkpoints)} checkpoints already exist, removing {len(removing_checkpoints)} checkpoints" |
| ) |
| logger.info(f"removing checkpoints: {', '.join(removing_checkpoints)}") |
|
|
| for removing_checkpoint in removing_checkpoints: |
| checkpoint_path = os.path.join(args.output_dir, removing_checkpoint) |
| if os.path.exists(checkpoint_path): |
| try: |
| shutil.rmtree(checkpoint_path) |
| logger.info(f'Directory "{checkpoint_path}" has been removed.') |
| except Exception as e: |
| logger.info(f'Error removing directory "{checkpoint_path}": {e}. Continuing with the next item.') |
| else: |
| logger.info(f'Directory "{checkpoint_path}" does not exist.') |
|
|
