Method save_checkpoint

deepspeed/runtime/engine.py:4859–4956

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)

                logger.warning(msg)

    def save_checkpoint(self, save_dir, tag=None, client_state={}, save_latest=True, exclude_frozen_parameters=False):
        """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.
            client_state: Optional. State dictionary used for saving required training states in the client code.
            save_latest: Optional. Save a file 'latest' pointing to the latest saved checkpoint.
            exclude_frozen_parameters: Optional. Exclude frozen parameters from checkpointed state.
        Important: all processes must call this method and not just the process with rank 0. It is
        because each process needs to save its master weights and scheduler+optimizer states. This
        method will hang waiting to synchronize with other processes if it's called just for the
        process with rank 0.

        """
        if not save_dir:
            raise ValueError(f"save_dir must be a non-empty string, got {save_dir!r}")

        if self._optimizer_has_ckpt_event_prologue():
            # Custom preparation for checkpoint save, if applicable
            self.optimizer.checkpoint_event_prologue()

        rank = self.local_rank if self.use_node_local_storage() else self.global_rank

        # This is to make sure the checkpoint names are created without collision
        # There seems to be issue creating them in parallel

        # Ensure save_dir directory exists
        if rank == 0:
            self.checkpoint_engine.makedirs(save_dir, exist_ok=True)
        dist.barrier()

        if tag is None:
            tag = f"global_step{self.global_steps}"

        # Ensure tag is a string
        tag = str(tag)
        commit_info = CheckpointCommitInfo(tag=tag, save_dir=save_dir, save_latest=save_latest)

        self.checkpoint_engine.create(commit_info)

        # Ensure checkpoint tag is consistent across ranks
        self._checkpoint_tag_validation(tag)

        if self.has_moe_layers:
            self.save_non_zero_checkpoint = False
            self._create_checkpoint_file(save_dir, tag, False)
            self._save_moe_checkpoint(save_dir,
                                      tag,
                                      client_state=client_state,
                                      exclude_frozen_parameters=exclude_frozen_parameters)

        # We distribute the task of saving layer checkpoint files among
        # data parallel instances, so all procs should call _save_checkpoint.
        # All procs then call module_state_dict(), but only procs of data
        # parallel rank 0 save the general model params.
        if not self.has_moe_layers: