| 406 | # self.print_params(model, ["model.mllm_backbone"]) |
| 407 | |
| 408 | class SaveCallback(TrainerCallback): |
| 409 | def on_save(self, args, state, control, model, **kwargs): |
| 410 | if not args.save_part_checkpoints: |
| 411 | return |
| 412 | |
| 413 | torch.distributed.barrier() |
| 414 | if torch.distributed.get_rank() == 0: |
| 415 | step = state.global_step |
| 416 | epoch = f"{state.epoch:.2f}".replace(".", "_") |
| 417 | |
| 418 | checkpoints_dir = os.path.join(args.base_dir, args.save_dir, "whole_model", f"epoch{epoch}_step{step}") |
| 419 | os.makedirs(checkpoints_dir, exist_ok=True) |
| 420 | torch.save( |
| 421 | model.state_dict(), |
| 422 | os.path.join(checkpoints_dir, f"model.pt") |
| 423 | ) |
| 424 | |
| 425 | torch.distributed.barrier() |