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 24 \
    --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] train_tokenizer.py (python)
    videos_sharding = NamedSharding(mesh, PartitionSpec("data", None, None, None, None))
    train_state = jax.device_put(train_state, replicated_sharding)

    # --- Load checkpoint ---
    step = 0
    if args.checkpoint:
        restore_target = {"model": train_state}
        restore_args = orbax_utils.restore_args_from_target(restore_target)
        train_state.params["params"].update(
            PyTreeCheckpointer()
            .restore(args.checkpoint, item=restore_target, restore_args=restore_args)[
                "model"
            ]
            .params["params"]
        )
        # Assume checkpoint is of the form tokenizer_<timestamp>_<step>
        step += int(args.checkpoint.split("_")[-1])

    # --- TRAIN LOOP ---
    tfrecord_files = [
        os.path.join(args.data_dir, x)
        for x in os.listdir(args.data_dir)
        if x.endswith(".tfrecord")
    ]
    dataloader = get_dataloader(
        # NOTE: We deliberately pass the global batch size
        # The dataloader shards the dataset across all processes
        tfrecord_files,
        args.seq_len,
        args.batch_size,
        *image_shape,
        seed=args.seed,
    )
    dataloader = (jax.make_array_from_process_local_data(videos_sharding, elem) for elem in dataloader)  # type: ignore
[tab] tests/data/generate_dummy_data.py (python)
import tyro
import tensorflow as tf
import numpy as np
from pathlib import Path
from dataclasses import dataclass


@dataclass
class Args:
    data_dir: str = "data_tfrecords/dummy"
    num_episodes: int = 5
    episode_length: int = 16


def _bytes_feature(value):
    """Returns a bytes_list from a string / byte."""
    if isinstance(value, type(tf.constant(0))):
        value = value.numpy()  # BytesList won't unpack a string from an EagerTensor.
    return tf.train.Feature(bytes_list=tf.train.BytesList(value=[value]))


def _int64_feature(value):
    """Returns an int64_list from a bool / enum / int / uint."""
    return tf.train.Feature(int64_list=tf.train.Int64List(value=[value]))


def create_tfrecord_example(episode_numpy_array):
    """Creates a TFRecord example from a numpy array video."""
    feature = {
        "height": _int64_feature(episode_numpy_array.shape[1]),
        "width": _int64_feature(episode_numpy_array.shape[2]),
        "channels": _int64_feature(episode_numpy_array.shape[3]),
        "sequence_length": _int64_feature(episode_numpy_array.shape[0]),
        "raw_video": _bytes_feature(episode_numpy_array.tobytes()),
    }
    return tf.train.Example(features=tf.train.Features(feature=feature))


def generate_dummy_tfrecord(
    output_path, num_episodes=5, episode_length=16, height=90, width=160, channels=3
):
    """Generates a dummy TFRecord file with synthetic video data."""
    print(f"Generating dummy TFRecord file at {output_path}")
    with tf.io.TFRecordWriter(str(output_path)) as writer:
        for i in range(num_episodes):
            np.random.seed(i)  # Seed per episode for some variation, but deterministic
            dummy_video = np.random.randint(
                0, 256, size=(episode_length, height, width, channels), dtype=np.uint8
            )
            tf_example = create_tfrecord_example(dummy_video)
            writer.write(tf_example.SerializeToString())
    print("Dummy TFRecord generation complete.")


if __name__ == "__main__":
    args = tyro.cli(Args)
    temp_dir = Path(args.data_dir)
    temp_dir.mkdir(parents=True, exist_ok=True)
    dummy_file = temp_dir / "dummy_test_shard.tfrecord"
    generate_dummy_tfrecord(
        dummy_file, num_episodes=args.num_episodes, episode_length=args.episode_length
    )
    print(f"Generated dummy file: {dummy_file}")