[content] /hkfs/home/project/hk-project-p0023960/tum_dbd0378/jafar/benchmark_batch_sizes_tokenizer.py (python)
    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*4, 9*4, 10*4, 11*4, 12*4, 13*4, 14*4, 15*4])
    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,