def test_checkpointing(self):
global_mesh = jtu.create_global_mesh((4, 2), ('x', 'y'))
global_input_shape = (8, 2)
mesh_axes = P('x', 'y')
num = util.prod(global_input_shape)
# First GDA
global_input_data1 = np.arange(num).reshape(global_input_shape)
def cb1(index):
return global_input_data1[index]
gda1 = GlobalDeviceArray.from_callback(global_input_shape, global_mesh,
mesh_axes, cb1)
ckpt_dir1 = pathlib.Path(self.create_tempdir('first').full_path)
# Second GDA
global_input_data2 = np.arange(num, num + num).reshape(global_input_shape)
def cb2(index):
return global_input_data2[index]
gda2 = GlobalDeviceArray.from_callback(global_input_shape, global_mesh,
mesh_axes, cb2)
ckpt_dir2 = pathlib.Path(self.create_tempdir('second').full_path)
# Third GDA
def cb3(index):
return np.array([])
global_mesh1d = jtu.create_global_mesh((8,), ('x',))
gda3 = GlobalDeviceArray.from_callback((0,), global_mesh1d, P(None), cb3)
ckpt_dir3 = pathlib.Path(self.create_tempdir('third').full_path)
ckpt_paths = [str(ckpt_dir1), str(ckpt_dir2), str(ckpt_dir3)]
tspecs = jax.tree_map(serialization.get_tensorstore_spec, ckpt_paths)
serialization.run_serialization([gda1, gda2, gda3], tspecs)
m1, m2, m3 = serialization.run_deserialization(
[global_mesh, global_mesh, global_mesh1d],
[mesh_axes, P('x'), P(None)],
tspecs)
self.assertArraysEqual(m1.local_shards[0].data.to_py(),
np.array([[0], [2]]))
self.assertArraysEqual(m1.local_shards[1].data.to_py(),
np.array([[1], [3]]))
self.assertEqual(m1.local_shards[0].data.shape, (2, 1))
self.assertEqual(m1.dtype, np.int32)
self.assertArraysEqual(m2.local_shards[0].data.to_py(),
np.array([[16, 17], [18, 19]]))
self.assertArraysEqual(m2.local_shards[1].data.to_py(),
np.array([[16, 17], [18, 19]]))
self.assertEqual(m2.local_shards[0].data.shape, (2, 2))
self.assertEqual(m2.dtype, np.int32)
for i, s in enumerate(m3.local_shards):
self.assertEqual(s.index, (slice(None),))
self.assertEqual(s.replica_id, i)
self.assertArraysEqual(s.data.to_py(), np.array([]))
self.assertEqual(m3.dtype, np.float32)
def test_checkpointing_with_bigger_shape(self):
global_mesh = create_global_mesh((2, 2), ('x', 'y'))
global_input_shape = (8, 2)
num = util.prod(global_input_shape)
# First GDA
global_input_data1 = np.arange(num).reshape(global_input_shape)
def cb1(index):
return global_input_data1[index]
gda1 = GlobalDeviceArray.from_callback(global_input_shape, global_mesh,
P('x', 'y'), cb1)
ckpt_dir1 = pathlib.Path(self.create_tempdir('first').full_path)
ckpt_paths = [str(ckpt_dir1)]
tspecs = jax.tree_map(serialization.get_tensorstore_spec, ckpt_paths)
serialization.run_serialization([gda1], tspecs)
m1, = serialization.run_deserialization(
[create_global_mesh((4, 2), ('x', 'y'))],
[P('x', 'y')],
tspecs,
[(12, 2)],
)
expected_data = {
0: np.array([[0], [2], [4]]),
1: np.array([[1], [3], [5]]),
2: np.array([[6], [8], [10]]),
3: np.array([[7], [9], [11]]),
4: np.array([[12], [14], [0]]),
5: np.array([[13], [15], [0]]),
6: np.array([[0], [0], [0]]),
7: np.array([[0], [0], [0]]),
}
for l in m1.local_shards:
self.assertArraysEqual(l.data.to_py(), expected_data[l.device.id])
def _save_checkpoint_gda(train_state: train_states.TrainState,
checkpoint_dir: str, overwrite: bool,
step: int) -> None:
"""Saves a checkpoint using JAX GDA serialization mechanism.
Note that all JAX processes must call _save_checkpoint_gda in sync because
each process may only have a slice of the global data.
Args:
train_state: A partitioned train_state that is a Pytree of
GlobalDeviceArray.
checkpoint_dir: Full path to parent checkpoint_dir.
overwrite: Whether to allow overwriting an existing target directory.
step: Step to save checkpoint for.
"""
if not overwrite:
# Does not contain directory path, only dirname is returned.
checkpoint_dirnames = tf.io.gfile.listdir(checkpoint_dir)
# Delete tmp directories if any.
if jax.process_index() == 0:
tmp_checkpoint_dirnames = [
x for x in checkpoint_dirnames if _is_tmp_checkpoint_asset(x)
]
if tmp_checkpoint_dirnames:
logging.warn('Found incompletely saved checkpoints %s; deleting them',
tmp_checkpoint_dirnames)
for x in tmp_checkpoint_dirnames:
tf.io.gfile.rmtree(os.path.join(checkpoint_dir, x))
# Note we must barrier across all processes after the tmp directory delete.
py_utils.sync_global_devices('Wait for checkpoint tmp dir deletions to '
'finish.')
sorted_dirnames = sorted(
[x for x in checkpoint_dirnames if _is_checkpoint_asset(x)])
if sorted_dirnames:
latest_checkpoint_dirname = sorted_dirnames[-1]
previous_step = get_step_from_checkpoint_asset(latest_checkpoint_dirname)
if previous_step >= step:
logging.warning(
'A more recent checkpoint `%d` has already been saved compared '
'to the current timestep `%d`. Skip saving a checkpoint.',
previous_step, step)
return
checkpoint_step_dir = _make_checkpoint_step_dir(checkpoint_dir, step)
checkpoint_step_tmp_dir = _make_tmp_checkpoint_dir(
checkpoint_dir, step, sync_timestamp=True)
logging.info('Saving to a tmp checkpoint dir %s', checkpoint_step_tmp_dir)
nested_names = _extract_nested_prefix_names(train_state)
flattened_nested_names, _ = jax.tree_util.tree_flatten(nested_names)
if jax.process_index() == 0:
# Create the tmp parent dir.
tf.io.gfile.makedirs(checkpoint_step_tmp_dir)
with futures.ThreadPoolExecutor() as executor:
ckpt_paths = list(
executor.map(_mkdir_path, flattened_nested_names,
[checkpoint_step_tmp_dir] * len(flattened_nested_names)))
py_utils.sync_global_devices('Wait for checkpoint tmp dir and subdirs '
f'creation {checkpoint_step_tmp_dir} to finish.')
tspecs = jax.tree_map(gda_serialization.get_tensorstore_spec, ckpt_paths)
leaves, _ = jax.tree_util.tree_flatten(train_state)
gda_serialization.run_serialization(leaves, tspecs)
# Note we must barrier across all processes before the directory rename.
py_utils.sync_global_devices('Wait for checkpoint chunk writes to '
f'{checkpoint_step_tmp_dir} to finish.')
if jax.process_index() == 0:
# Rename temporary checkpoint directory to its final location.
logging.info('Renaming %s to %s', checkpoint_step_tmp_dir,
checkpoint_step_dir)
tf.io.gfile.rename(checkpoint_step_tmp_dir, checkpoint_step_dir)
logging.info('Finished saving GDA checkpoint for step `%s` to `%s`.', step,
checkpoint_step_dir)