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

#SBATCH --nodes=1
#SBATCH --ntasks-per-node=8
#SBATCH --time=12:00:00
#SBATCH --cpus-per-task=1
#SBATCH --gres=gpu:8
#SBATCH --mem=100GB

source .venv/bin/activate

data_dir="${PWD}/data_arrayrecord/dummy"
ckpt_dir="${PWD}/checkpoints/tokenizer_openai_grain_checkpointing"

mkdir -p $ckpt_dir

export XLA_FLAGS=--xla_gpu_autotune_level=0
export PYTHONUNBUFFERED=1
srun python train_tokenizer.py \
    --batch_size 24 \
    --restore_ckpt \
    --save_ckpt \
    --log_checkpoint_interval 5 \
    --ckpt_dir $ckpt_dir \
    --num_steps 300000 \
    --warmup_steps 10000 \
    --seed 0 \
    --init_lr=0.0000866 \
    --max_lr=0.0000866 \
    --data_dir $data_dir
[content] tests/data/generate_dummy_data.py (python)
    test_dir = Path("tests/data/dummy_arrayrecord")
    test_dir.mkdir(parents=True, exist_ok=True)
    dummy_file = test_dir / "dummy_test_shard.array_record"

    generate_dummy_arrayrecord(dummy_file, episode_length=1000)
[content] tests/data/generate_dummy_data.py (python)
    height: int = 64,
    width: int = 64,
    channels: int = 3,
    num_action_tokens: int = 11,
    vocab_size: int = 68,  # default vocab size from minerl
    seed: int = 42,
):
    """Generates a dummy ArrayRecord file with synthetic video data for testing."""
    print(f"Generating dummy ArrayRecord file at {output_path}")

    output_path.parent.mkdir(parents=True, exist_ok=True)

    writer = ArrayRecordWriter(str(output_path), "group_size:1")
    try:
        for i in range(num_videos):
            np.random.seed(seed + i)
            dummy_video = np.random.randint(
                0, 256, size=(episode_length, height, width, channels), dtype=np.uint8
            )

            actions = np.random.randint(
                0,
                vocab_size,
                size=(episode_length, num_action_tokens),
                dtype=np.uint8,
            )

            record = {
                "raw_video": dummy_video.tobytes(),
                "sequence_length": episode_length,
                "actions": actions,
            }

            writer.write(pickle.dumps(record))
    finally:
        writer.close()

    print("Dummy ArrayRecord generation complete.")


if __name__ == "__main__":
    test_dir = Path("tests/data/dummy_arrayrecord")
    test_dir.mkdir(parents=True, exist_ok=True)
    dummy_file = test_dir / "dummy_test_shard.array_record"

    generate_dummy_arrayrecord(
        dummy_file, num_videos=5, episode_length=16, height=64, width=64, channels=3
    )
[tab] experiments/tokenizer_grain_checkpointing.sh (shellscript)
#!/usr/bin/env bash

#SBATCH --nodes=1
#SBATCH --ntasks-per-node=8
#SBATCH --time=12:00:00
#SBATCH --cpus-per-task=1
#SBATCH --gres=gpu:8
#SBATCH --mem=100GB

source .venv/bin/activate

data_dir="${PWD}/data_arrayrecord/dummy"
ckpt_dir="${PWD}/checkpoints/tokenizer_openai_grain_checkpointing"

mkdir -p $ckpt_dir

export XLA_FLAGS=--xla_gpu_autotune_level=0
export PYTHONUNBUFFERED=1
srun python train_tokenizer.py \
    --batch_size 12 \
    --save_ckpt \
    --log_checkpoint_interval 5 \
    --ckpt_dir $ckpt_dir \
    --num_steps 300000 \
    --warmup_steps 10000 \
    --seed 0 \
    --init_lr=0.0000866 \
    --max_lr=0.0000866 \
    --data_dir $data_dir
[tab] experiments/dynamics_grain_tok_restore.sh (shellscript)
#!/usr/bin/env bash

#SBATCH --nodes=2
#SBATCH --ntasks-per-node=8
#SBATCH --time=12:00:00
#SBATCH --cpus-per-task=1
#SBATCH --gres=gpu:8
#SBATCH --mem=100GB

source .venv/bin/activate

data_dir="${PWD}/data_arrayrecord/dummy"
ckpt_dir="${PWD}/checkpoints/causal_dynamics_openai_grain_tok_restore"
tokenizer_ckpt_dir="${PWD}/checkpoints/tokenizer_openai_grain_checkpointing"

export XLA_FLAGS=--xla_gpu_autotune_level=0
export PYTHONUNBUFFERED=1
srun python train_dynamics.py \
    --dyna_type 'causal' \
    --save_ckpt \
    --log_checkpoint_interval 5 \
    --batch_size 12 \
    --tokenizer_checkpoint $tokenizer_ckpt_dir \
    --ckpt_dir $ckpt_dir \
    --num_steps 300000 \
    --warmup_steps 10000 \
    --seed 0 \
    --init_lr=0.0000866 \
    --max_lr=0.0000866 \
    --data_dir $data_dir
[tab] checkpoints/tokenizer_openai_grain_checkpointing/000015/_CHECKPOINT_METADATA (plaintext)
{"item_handlers": {"dataloader_state": "grain._src.python.checkpoint_handlers.CheckpointHandler", "model_state": "orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeCheckpointHandler"}, "metrics": {}, "performance_metrics": {}, "init_timestamp_nsecs": 1753871211190928220, "commit_timestamp_nsecs": 1753871214017484231, "custom_metadata": {}}