Method restore

axlearn/common/checkpointer_orbax.py:531–608

See `BaseCheckpointer.restore` for details.

(
        self,
        *,
        step: Optional[int] = None,
        state: Union[Nested[Tensor], Nested[TensorSpec]],
    )

            # https://orbax.readthedocs.io/en/latest/preemption_checkpointing.html
            if self._manager.reached_preemption(step):
                self._manager.wait_until_finished()
                raise SystemExit(f"Exiting after saving checkpoint at {step=} due to pre-emption.")
        finally:
            self._eval_summaries = None

    def restore(
        self,
        *,
        step: Optional[int] = None,
        state: Union[Nested[Tensor], Nested[TensorSpec]],
    ) -> Tuple[Optional[int], Nested[Tensor]]:
        """See `BaseCheckpointer.restore` for details."""

        cfg: OrbaxCheckpointer.Config = self.config

        def _restore_args(x: Any) -> ocp.RestoreArgs:
            if isinstance(x, (Tensor, TensorSpec)):
                arg = ocp.checkpoint_utils.construct_restore_args(
                    jax.ShapeDtypeStruct(shape=x.shape, dtype=x.dtype, sharding=x.sharding)
                )
                if cfg.enable_single_replica_ckpt_restoring and isinstance(
                    arg, ocp.type_handlers.ArrayRestoreArgs
                ):
                    mesh = x.sharding.mesh
                    arg = ocp.type_handlers.SingleReplicaArrayRestoreArgs(
                        restore_type=arg.restore_type,
                        dtype=arg.dtype,
                        mesh=arg.mesh,
                        mesh_axes=arg.mesh_axes,
                        sharding=arg.sharding,
                        global_shape=arg.global_shape,
                        shape=arg.shape,
                        strict=arg.strict,
                        single_replica_sharding=jax.sharding.NamedSharding(
                            jax.sharding.Mesh(
                                _replica_devices(mesh.devices, cfg.replica_axis_index),
                                mesh.axis_names,
                            ),
                            x.sharding.spec,
                        ),
                    )
                return arg
            elif isinstance(x, tf.data.Iterator):
                return _TfIteratorHandler.RestoreArgs(item=x)
            elif _GRAIN_INSTALLED and isinstance(x, _GrainIterator):
                return _PythonSavableHandler.RestoreArgs(item=x)
            elif isinstance(x, PythonSavable):
                return _PythonSavableHandler.RestoreArgs(item=x)
            else:
                return None

        restore_args = jax.tree.map(_restore_args, state)

        try:
            composite_state = self._manager.restore(
                step,
                args=ocp.args.Composite(
                    index=ocp.args.JsonRestore(None),