Method test_gda

axlearn/common/checkpointer_test.py:395–419

(self, checkpointer_cls, mesh_shape)

        mesh_shape=[(1, 1), (2, 2), (4, 2)],
    )
    def test_gda(self, checkpointer_cls, mesh_shape):
        if not test_utils.is_supported_mesh_shape(mesh_shape):
            return
        with _mesh(mesh_shape):
            cfg = _checkpointer_config(checkpointer_cls)
            ckpt: Checkpointer = cfg.instantiate(parent=None)
            state = dict(x=jnp.arange(16).reshape((4, 4)))
            ckpt.save(step=10, state=state)
            ckpt.wait_until_finished()

            state0 = dict(x=jnp.zeros(shape=(4, 4), dtype=jnp.int32))
            step, restored_state = ckpt.restore(
                step=None,
                state=state0,
            )
            self.assertEqual(10, step)
            self.assertNestedEqual(restored_state, state)

            # dtype mismatch.
            with self.assertRaisesRegex(ValueError, "checkpoint tree dtypes or shapes"):
                ckpt.restore(
                    step=None,
                    state=dict(x=jnp.zeros(shape=(4, 4), dtype=jnp.float32)),
                )
            ckpt.stop()

    @parameterized.product(
        checkpointer_cls=[Checkpointer, OrbaxCheckpointer],