[tab] benchmark_batch_sizes_lam.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.lam import LatentActionModel
from utils.dataloader import get_dataloader
from utils.parameter_utils import count_parameters_by_component


@dataclass
class BenchmarkArgs:
    # Model parameters (same as train_lam.py)
    seed: int = 0
    seq_len: int = 16
    image_channels: int = 3
    image_height: int = 90
    image_width: int = 160
    data_dir: str = "/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/data_new/open_ai_minecraft_arrayrecords/10fps_160x90"
    
    # LAM parameters
    model_dim: int = 512
    latent_dim: int = 32
    num_latents: int = 6
    patch_size: int = 16
    num_blocks: int = 8
    num_heads: int = 8
    dropout: float = 0.0
    codebook_dropout: float = 0.0
    
    # 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_lam.json"
    
    # VQ parameters
    vq_beta: float = 0.25


args = tyro.cli(BenchmarkArgs)


def lam_loss_fn(params, state, inputs):
    # --- Compute loss ---
    outputs = state.apply_fn(
        params, inputs, training=True, rngs={"dropout": inputs["rng"]}
    )
    gt_future_frames = inputs["videos"][:, 1:]
    mse = jnp.square(gt_future_frames - 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(lam_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, lam, train_state, mesh, videos_sharding, grain_iterator) -> 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
    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 = jax.random.split(rng)
            inputs = dict(videos=videos, rng=_rng)
            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 = jax.random.split(rng)
            inputs = dict(videos=videos, rng=_rng)
            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)
    }


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 ---
    lam = LatentActionModel(
        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,
    )
    
    # 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=jnp.float32,
        ),
        rng=_rng,
    )
    init_params = lam.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)
    train_state = TrainState.create(apply_fn=lam.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)

    # Initialize dataloader (will be recreated for each batch size)
    array_record_files = [
        os.path.join(args.data_dir, x)
        for x in os.listdir(args.data_dir)
        if x.endswith(".array_record")
    ]
    grain_dataloader = get_dataloader(
        array_record_files,
        args.seq_len,
        min_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)

    # Run benchmarks
    results = []
    for batch_size in args.batch_sizes:
        result = benchmark_batch_size(batch_size, lam, train_state, mesh, videos_sharding, grain_iterator)
        if result is not None:
            results.append(result)
    
    # Save results
    benchmark_data = {
        "args": vars(args),
        "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()