Save training checkpoint Arguments: save_dir: Required. Directory for saving the checkpoint tag: Optional. Checkpoint tag used as a unique identifier for the checkpoint, global step is used if not provided. Tag name must be the same across all ranks.
(self, save_dir, tag=None, client_state={}, save_latest=True, exclude_frozen_parameters=False)
| 4555 | logger.warning(msg) |
| 4556 | |
| 4557 | def save_checkpoint(self, save_dir, tag=None, client_state={}, save_latest=True, exclude_frozen_parameters=False): |
| 4558 | """Save training checkpoint |
| 4559 | |
| 4560 | Arguments: |
| 4561 | save_dir: Required. Directory for saving the checkpoint |
| 4562 | tag: Optional. Checkpoint tag used as a unique identifier for the checkpoint, global step is |
| 4563 | used if not provided. Tag name must be the same across all ranks. |
| 4564 | client_state: Optional. State dictionary used for saving required training states in the client code. |
| 4565 | save_latest: Optional. Save a file 'latest' pointing to the latest saved checkpoint. |
| 4566 | exclude_frozen_parameters: Optional. Exclude frozen parameters from checkpointed state. |
| 4567 | Important: all processes must call this method and not just the process with rank 0. It is |
| 4568 | because each process needs to save its master weights and scheduler+optimizer states. This |
| 4569 | method will hang waiting to synchronize with other processes if it's called just for the |
| 4570 | process with rank 0. |
| 4571 | |
| 4572 | """ |
| 4573 | if not save_dir: |
| 4574 | raise ValueError(f"save_dir must be a non-empty string, got {save_dir!r}") |
| 4575 | |
| 4576 | if self._optimizer_has_ckpt_event_prologue(): |
| 4577 | # Custom preparation for checkpoint save, if applicable |
| 4578 | self.optimizer.checkpoint_event_prologue() |
| 4579 | |
| 4580 | rank = self.local_rank if self.use_node_local_storage() else self.global_rank |
| 4581 | |
| 4582 | # This is to make sure the checkpoint names are created without collision |
| 4583 | # There seems to be issue creating them in parallel |
| 4584 | |
| 4585 | # Ensure save_dir directory exists |
| 4586 | if rank == 0: |
| 4587 | self.checkpoint_engine.makedirs(save_dir, exist_ok=True) |
| 4588 | dist.barrier() |
| 4589 | |
| 4590 | if tag is None: |
| 4591 | tag = f"global_step{self.global_steps}" |
| 4592 | |
| 4593 | # Ensure tag is a string |
| 4594 | tag = str(tag) |
| 4595 | commit_info = CheckpointCommitInfo(tag=tag, save_dir=save_dir, save_latest=save_latest) |
| 4596 | |
| 4597 | self.checkpoint_engine.create(commit_info) |
| 4598 | |
| 4599 | # Ensure checkpoint tag is consistent across ranks |
| 4600 | self._checkpoint_tag_validation(tag) |
| 4601 | |
| 4602 | if self.has_moe_layers: |
| 4603 | self.save_non_zero_checkpoint = False |
| 4604 | self._create_checkpoint_file(save_dir, tag, False) |
| 4605 | self._save_moe_checkpoint(save_dir, |
| 4606 | tag, |
| 4607 | client_state=client_state, |
| 4608 | exclude_frozen_parameters=exclude_frozen_parameters) |
| 4609 | |
| 4610 | # We distribute the task of saving layer checkpoint files among |
| 4611 | # data parallel instances, so all procs should call _save_checkpoint. |
| 4612 | # All procs then call module_state_dict(), but only procs of data |
| 4613 | # parallel rank 0 save the general model params. |
| 4614 | if not self.has_moe_layers: |