[tab] benchmark_batch_sizes_tokenizer.py (python)
from dataclasses import dataclass, field
import os
import time
import json
from typing import List, Dict, Any, Optional

import einops
from flax.training import orbax_utils
from flax.training.train_state import TrainState
from jax.sharding import Mesh, PartitionSpec, NamedSharding
from jax.experimental.mesh_utils import create_device_mesh
import optax
import orbax.checkpoint as ocp
import numpy as np
import dm_pix as pix
import jax
import jax.numpy as jnp
import tyro
import grain

from models.tokenizer import TokenizerVQVAE
from utils.dataloader import get_dataloader
from utils.parameter_utils import count_parameters_by_component


def create_mock_dataloader(batch_size: int, seq_len: int, image_height: int, image_width: int, image_channels: int, num_devices: int):
    """
    Create a mock dataloader that generates instant dummy data.
    This eliminates IO bottlenecks to test if you're memory-bound.
    """
    per_device_batch_size = batch_size // num_devices
    
    def mock_dataloader():
        """Generator that yields dummy video data instantly."""
        while True:
            # Generate random dummy data with the correct shape
            # Using jax.random for consistency with JAX
            dummy_videos = jax.random.randint(
                jax.random.PRNGKey(0),  # Fixed key for deterministic but random data
                (per_device_batch_size, seq_len, image_height, image_width, image_channels),
                minval=0,
                maxval=255,
                dtype=jnp.uint8
            )
            yield dummy_videos
    
    return mock_dataloader()


@dataclass
class BenchmarkArgs:
    # Model parameters (same as train_tokenizer.py)
    seed: int = 0
    seq_len: int = 16
    image_channels: int = 3
    image_height: int = 90
    image_width: int = 160
    data_dir: str = "/home/franz.srambical/jafar/data_arrayrecord/dummy"
    
    # Tokenizer parameters
    model_dim: int = 512
    latent_dim: int = 32
    num_latents: int = 1024
    patch_size: int = 4
    num_blocks: int = 8
    num_heads: int = 8
    dropout: float = 0.0
    codebook_dropout: float = 0.01
    dtype: jnp.dtype = jnp.bfloat16
    param_dtype: jnp.dtype = jnp.float32
    use_flash_attention: bool = True
    
    # Benchmark parameters
    batch_sizes: List[int] = field(default_factory=lambda: [8*1, 9*1, 10*1, 11*1, 12*1, 13*1, 14*1, 15*1])
    warmup_steps: int = 10  # Steps to warm up before measuring
    measure_steps: int = 30  # Steps to measure
    num_iterations: int = 10  # Number of iterations to average over
    output_file: str = "batch_size_benchmark_tokenizer_flash_attention.json"
    use_mock_dataloader: bool = False  # Use mock dataloader instead of real one
    
    # VQ parameters
    vq_beta: float = 0.25


args = tyro.cli(BenchmarkArgs)


def tokenizer_loss_fn(params, state, inputs):
    # --- Compute loss ---
    inputs["videos"] = inputs["videos"].astype(args.dtype) / 255.0
    outputs = state.apply_fn(
        params,
        inputs,
        training=True,
        rngs={"params": inputs["rng"], "dropout": inputs["dropout_rng"]},
    )
    mse = jnp.square(inputs["videos"] - outputs["recon"]).mean()
    q_loss = jnp.square(jax.lax.stop_gradient(outputs["emb"]) - outputs["z"]).mean()
    commitment_loss = jnp.square(
        outputs["emb"] - jax.lax.stop_gradient(outputs["z"])
    ).mean()
    loss = mse + q_loss + args.vq_beta * commitment_loss
    return loss, outputs["recon"]


@jax.jit
def train_step(state, inputs):
    grad_fn = jax.value_and_grad(tokenizer_loss_fn, has_aux=True, allow_int=True)
    (loss, recon), grads = grad_fn(state.params, state, inputs)
    state = state.apply_gradients(grads=grads)
    return state, loss, recon


def benchmark_batch_size(batch_size: int, tokenizer, train_state, mesh, videos_sharding) -> Optional[Dict[str, Any]]:
    """Benchmark a specific batch size and return throughput metrics."""
    print(f"
Benchmarking batch size: {batch_size}")
    
    num_devices = jax.device_count()
    per_device_batch_size = batch_size // num_devices
    
    # Create dataloader for this batch size
    if args.use_mock_dataloader:
        print("  Using mock dataloader (instant dummy data)")
        image_shape = (args.image_height, args.image_width, args.image_channels)
        mock_dataloader = create_mock_dataloader(
            batch_size, args.seq_len, args.image_height, args.image_width, args.image_channels, num_devices
        )
        dataloader = (jax.make_array_from_process_local_data(videos_sharding, np.array(elem)) for elem in mock_dataloader)
    else:
        print("  Using real dataloader")
        array_record_files = [
            os.path.join(args.data_dir, x)
            for x in os.listdir(args.data_dir)
            if x.endswith(".array_record")
        ]
        image_shape = (args.image_height, args.image_width, args.image_channels)
        grain_dataloader = get_dataloader(
            array_record_files,
            args.seq_len,
            batch_size,
            *image_shape,
            num_workers=8,
            prefetch_buffer_size=1,
            seed=args.seed,
        )
        initial_state = grain_dataloader._create_initial_state()
        grain_iterator = grain.DataLoaderIterator(grain_dataloader, initial_state)
        dataloader = (jax.make_array_from_process_local_data(videos_sharding, elem) for elem in grain_iterator)
    
    rng = jax.random.PRNGKey(args.seed)
    step_times = []
    
    # Warmup phase
    print(f"  Warming up for {args.warmup_steps} steps...")
    for step in range(args.warmup_steps):
        try:
            videos = next(dataloader)
            rng, _rng, _rng_dropout = jax.random.split(rng, 3)
            inputs = dict(videos=videos, rng=_rng, dropout_rng=_rng_dropout)
            train_state, loss, recon = train_step(train_state, inputs)
            jax.block_until_ready(train_state)
        except StopIteration:
            print(f"  Warning: Dataloader exhausted during warmup at step {step}")
            break
    
    # Measurement phase
    print(f"  Measuring for {args.measure_steps} steps...")
    start_time = time.time()
    for step in range(args.measure_steps):
        try:
            videos = next(dataloader)
            rng, _rng, _rng_dropout = jax.random.split(rng, 3)
            inputs = dict(videos=videos, rng=_rng, dropout_rng=_rng_dropout)
            train_state, loss, recon = train_step(train_state, inputs)
            jax.block_until_ready(train_state)
            
            if step > 0:  # Skip first step as it might include compilation time
                step_times.append(time.time() - start_time)
            start_time = time.time()
            
        except StopIteration:
            print(f"  Warning: Dataloader exhausted during measurement at step {step}")
            break
    
    if len(step_times) < 5:
        print(f"  Warning: Only {len(step_times)} step times measured, results may be unreliable")
        return None
    
    # Calculate metrics
    avg_step_time = np.mean(step_times)
    std_step_time = np.std(step_times)
    throughput = batch_size / avg_step_time
    
    print(f"  Results: avg_step_time={avg_step_time:.4f}s ± {std_step_time:.4f}s, throughput={throughput:.2f} samples/s")
    
    return {
        "batch_size": batch_size,
        "avg_step_time": avg_step_time,
        "std_step_time": std_step_time,
        "throughput": throughput,
        "num_measurements": len(step_times),
        "dataloader_type": "mock" if args.use_mock_dataloader else "real"
    }


def main():
    jax.distributed.initialize()
    num_devices = jax.device_count()
    if num_devices == 0:
        raise ValueError("No JAX devices found.")
    print(f"Running on {num_devices} devices.")

    # Check if all batch sizes are divisible by number of devices
    for batch_size in args.batch_sizes:
        if batch_size % num_devices != 0:
            raise ValueError(
                f"Batch size {batch_size} must be divisible by number of devices {num_devices}."
            )

    rng = jax.random.PRNGKey(args.seed)

    # --- Initialize model ---
    tokenizer = TokenizerVQVAE(
        in_dim=args.image_channels,
        model_dim=args.model_dim,
        latent_dim=args.latent_dim,
        num_latents=args.num_latents,
        patch_size=args.patch_size,
        num_blocks=args.num_blocks,
        num_heads=args.num_heads,
        dropout=args.dropout,
        codebook_dropout=args.codebook_dropout,
        param_dtype=args.param_dtype,
        dtype=args.dtype,
        use_flash_attention=args.use_flash_attention,
    )
    
    # Initialize with smallest batch size for model setup
    min_batch_size = min(args.batch_sizes)
    per_device_batch_size_for_init = min_batch_size // num_devices
    
    rng, _rng = jax.random.split(rng)
    image_shape = (args.image_height, args.image_width, args.image_channels)
    inputs = dict(
        videos=jnp.zeros(
            (per_device_batch_size_for_init, args.seq_len, *image_shape),
            dtype=args.dtype,
        ),
    )
    init_params = tokenizer.init(_rng, inputs)

    param_counts = count_parameters_by_component(init_params)
    print("Parameter counts:")
    print(param_counts)

    # --- Initialize optimizer ---
    lr_schedule = optax.warmup_cosine_decay_schedule(0.0, 3e-4, 1000, 10000)
    tx = optax.adamw(learning_rate=lr_schedule, b1=0.9, b2=0.9, weight_decay=1e-4, mu_dtype=args.dtype)
    train_state = TrainState.create(apply_fn=tokenizer.apply, params=init_params, tx=tx)

    # Setup sharding
    device_mesh_arr = create_device_mesh((num_devices,))
    mesh = Mesh(devices=device_mesh_arr, axis_names=("data",))
    replicated_sharding = NamedSharding(mesh, PartitionSpec())
    videos_sharding = NamedSharding(mesh, PartitionSpec("data", None, None, None, None))
    train_state = jax.device_put(train_state, replicated_sharding)

    # Run benchmarks
    results = []
    for batch_size in args.batch_sizes:
        result = benchmark_batch_size(batch_size, tokenizer, train_state, mesh, videos_sharding)
        if result is not None:
            results.append(result)
    
    # Save results
    benchmark_data = {
        "num_devices": num_devices,
        "param_counts": param_counts,
        "results": results
    }
    
    with open(args.output_file, 'w') as f:
        json.dump(benchmark_data, f, indent=2)
    
    print(f"
Benchmark results saved to {args.output_file}")
    
    # Print summary
    print("
Batch Size vs Throughput Summary:")
    print("Batch Size | Avg Step Time (s) | Throughput (samples/s)")
    print("-" * 55)
    for result in results:
        print(f"{result['batch_size']:9d} | {result['avg_step_time']:15.4f} | {result['throughput']:20.2f}")


if __name__ == "__main__":
    main()