Method save

axlearn/common/checkpointer.py:1108–1130

See `BaseCheckpointer.save` for details. In addition to behavior in `BaseCheckpointer`, saving only happens if the configured checkpoint policy returns True for the given step and evaler summaries.

(
        self, *, step: int, state: NestedTensor, evaler_summaries: Optional[dict[str, Any]] = None
    )

        return build_step_dir(cfg.dir, step=step)

    def save(
        self, *, step: int, state: NestedTensor, evaler_summaries: Optional[dict[str, Any]] = None
    ):
        """See `BaseCheckpointer.save` for details.

        In addition to behavior in `BaseCheckpointer`, saving only happens if the configured
        checkpoint policy returns True for the given step and evaler summaries.
        """
        if not self._save_policy(step=step, evaler_summaries=(evaler_summaries or {})):
            return
        if step < 0 or step >= 10**8:
            raise ValueError(f"Out-of-range: {step}")
        ckpt_dir = self.ckpt_dir(step)
        self._storage.save_to_dir(
            step=step, state=state, ckpt_dir=ckpt_dir, on_commit_callback=self._index_writer
        )
        if "summary_writer" in self.children:
            self.summary_writer.log_checkpoint(
                step=step,
                state=state,
                ckpt_dir=ckpt_dir,
                action=CheckpointerAction.SAVE,
            )

    def _run_garbage_collection(self):
        """Runs one round of garbage collection of past checkpoints.