p-doom/crowd-code-dataset-0.2 huggingface dataset:
[tab] 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)
        # Compare parameters recursively, handling nested structure
        def compare_params(original, restored):
            if isinstance(original, dict):
                for k in original.keys():
                    compare_params(original[k], restored[k])
            else:
                self.assertTrue(jnp.allclose(original, restored))
        
        compare_params(state.params, restored["model"].params)

if __name__ == "__main__":
    unittest.main()
[terminal_command] TERMINAL ()
/usr/bin/python3 /ictstr01/home/aih/franz.srambical/.cursor-server/extensions/ms-python.python-2024.12.3-linux-x64/python_files/printEnvVariablesToFile.py /ictstr01/home/aih/franz.srambical/.cursor-server/extensions/ms-python.python-2024.12.3-linux-x64/python_files/deactivate/bash/envVars.txt