p-doom/crowd-code-dataset-0.2 huggingface dataset:
[tab] experiments/tokenizer_cross_node_checkpointing_test.sh (shellscript)
#!/usr/bin/env bash
source .venv/bin/activate


data_dir='data_tfrecords'

srun python train_tokenizer.py \
    --batch_size 12 \
    --ckpt_dir checkpoints/tokenizer_cross_node_checkpointing_test \
    --log_checkpoint_interval 10 \
    --num_steps 300000 \
    --warmup_steps 10000 \
    --seed 0 \
    --min_lr=0.0000866 \
    --max_lr=0.0000866 \
    --data_dir $data_dir
[content] tests/test_checkpointer.py (python)
import unittest
import tempfile
import os
import jax
import jax.numpy as jnp
from flax.training import orbax_utils
from orbax.checkpoint import PyTreeCheckpointer
from pathlib import Path

from models.tokenizer import TokenizerVQVAE
from flax.training.train_state import TrainState
import optax
from jax.sharding import Mesh, PartitionSpec, NamedSharding
from jax.experimental.mesh_utils import create_device_mesh

class DistributedCheckpointerTest(unittest.TestCase):
    def setUp(self):
        super().setUp()
        self._temp_dir_manager = tempfile.TemporaryDirectory()
        self.checkpoint_dir = Path(self._temp_dir_manager.name)
        self.addCleanup(self._temp_dir_manager.cleanup)

        # FIXME (f.srambical): If the tests pass, we should use the default model config instead
        self.model_kwargs = dict(
            in_dim=3,
            model_dim=8,
            latent_dim=4,
            num_latents=16,
            patch_size=2,
            num_blocks=1,
            num_heads=1,
            dropout=0.0,
            codebook_dropout=0.0,
        )
        self.image_shape = (8, 8, 3)
        self.seq_len = 2
        self.batch_size = 2
        self.seed = 0

    def test_distributed_checkpointing(self):
        jax.distributed.initialize()
        num_devices = jax.device_count()
        self.assertGreater(num_devices, 0)

        model = TokenizerVQVAE(**self.model_kwargs)
        rng = jax.random.PRNGKey(self.seed)
        dummy_inputs = dict(
            videos=jnp.zeros((self.batch_size, self.seq_len, *self.image_shape), dtype=jnp.float32)
        )
        params = model.init(rng, dummy_inputs)

        tx = optax.adam(1e-3)
        state = TrainState.create(apply_fn=model.apply, params=params, tx=tx)

        device_mesh_arr = create_device_mesh((num_devices,))
        mesh = Mesh(devices=device_mesh_arr, axis_names=("data",))
        replicated_sharding = NamedSharding(mesh, PartitionSpec())
        state = jax.device_put(state, replicated_sharding)

        ckpt = {"model": state}
        orbax_checkpointer = PyTreeCheckpointer()
        save_args = orbax_utils.save_args_from_target(ckpt)
        ckpt_path = str(self.checkpoint_dir / "test_ckpt")
        orbax_checkpointer.save(ckpt_path, ckpt, save_args=save_args)
        self.assertTrue(os.path.exists(ckpt_path))

        restore_target = {"model": state}
        restore_args = orbax_utils.restore_args_from_target(restore_target)
        restored = orbax_checkpointer.restore(ckpt_path, item=restore_target, restore_args=restore_args)
        for k in state.params.keys():
            self.assertTrue(jax.tree_util.tree_all(jnp.allclose(state.params[k], restored["model"].params[k])))

if __name__ == "__main__":
    unittest.main()
[terminal_command] TERMINAL ()
salloc --reservation=haicu_stefan      -p gpu_p  --time=05:00:00      --job-name=interactive_bash  --qos=gpu_normal --gres=gpu:2 --gpu-bind=single:1 -w gpusrv69,gpusrv70 --cpus-per-task=4 --ntasks-per-node=2
[terminal_command] TERMINAL ()
salloc --reservation=haicu_stefan      -p gpu_p  --time=05:00:00      --job-name=interactive_bash  --qos=gpu_normal --gpu-bind=single:1 --gpus-per-task=1 -w gpusrv69,gpusrv70 --cpus-per-task=4 --ntasks-per-node=2
[terminal_command] TERMINAL ()
salloc --reservation=haicu_stefan      -p gpu_p  --time=05:00:00      --job-name=interactive_bash  --qos=gpu_normal --gres=gpu:2 --gpu-bind=single:1 --gpus-per-task=1 -w gpusrv69,gpusrv70 --cpus-per-task=4 --ntasks-per-node=2
[terminal_command] TERMINAL ()
salloc --reservation=haicu_stefan      -p gpu_p  --time=05:00:00      --job-name=interactive_bash  --qos=gpu_normal --gpu-bind=single:1 --gpus-per-task=1 -w gpusrv69,gpusrv70 --cpus-per-task=4 --ntasks-per-node=2 --ntasks=4