[tab] benchmark_batch_sizes_dynamics.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 genie import Genie
from utils.dataloader import get_dataloader
from utils.parameter_utils import count_parameters_by_component


@dataclass
class BenchmarkArgs:
    # Model parameters (same as train_dynamics.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"
    
    # Tokenizer parameters
    tokenizer_dim: int = 512
    latent_patch_dim: int = 32
    num_patch_latents: int = 1024
    patch_size: int = 4
    tokenizer_num_blocks: int = 8
    tokenizer_num_heads: int = 8
    
    # LAM parameters
    lam_dim: int = 512
    latent_action_dim: int = 32
    num_latent_actions: int = 6
    lam_patch_size: int = 16
    lam_num_blocks: int = 8
    lam_num_heads: int = 8
    
    # Dynamics parameters
    dyna_dim: int = 512
    dyna_num_blocks: int = 12
    dyna_num_heads: int = 8
    dropout: float = 0.0
    mask_limit: float = 0.5
    
    # 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_dynamics_co_train.json"
    
    # Training parameters
    lam_co_train: bool = True  # Set to True if you want to co-train LAM with dynamics


args = tyro.cli(BenchmarkArgs)


def dynamics_loss_fn(params, state, inputs):
    """Compute masked dynamics loss"""
    outputs = state.apply_fn(
        params,
        inputs,
        training=True,
        rngs={"params": inputs["rng"], "dropout": inputs["dropout_rng"]},
    )
    mask = outputs["mask"]
    ce_loss = optax.softmax_cross_entropy_with_integer_labels(
        outputs["token_logits"], outputs["video_tokens"]
    )
    ce_loss = (mask * ce_loss).sum() / mask.sum()
    return ce_loss, outputs["recon"]


@jax.jit
def train_step(state, inputs):
    """Update state and compute metrics"""
    grad_fn = jax.value_and_grad(dynamics_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, genie, 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, _rng_dropout, _rng_mask = jax.random.split(rng, 4)
            inputs = dict(
                videos=videos,
                rng=_rng,
                dropout_rng=_rng_dropout,
                mask_rng=_rng_mask,
            )
            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, _rng_mask = jax.random.split(rng, 4)
            inputs = dict(
                videos=videos,
                rng=_rng,
                dropout_rng=_rng_dropout,
                mask_rng=_rng_mask,
            )
            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 ---
    genie = Genie(
        # Tokenizer
        in_dim=args.image_channels,
        tokenizer_dim=args.tokenizer_dim,
        latent_patch_dim=args.latent_patch_dim,
        num_patch_latents=args.num_patch_latents,
        patch_size=args.patch_size,
        tokenizer_num_blocks=args.tokenizer_num_blocks,
        tokenizer_num_heads=args.tokenizer_num_heads,
        # LAM
        lam_dim=args.lam_dim,
        latent_action_dim=args.latent_action_dim,
        num_latent_actions=args.num_latent_actions,
        lam_patch_size=args.lam_patch_size,
        lam_num_blocks=args.lam_num_blocks,
        lam_num_heads=args.lam_num_heads,
        lam_co_train=args.lam_co_train,
        # Dynamics
        dyna_dim=args.dyna_dim,
        dyna_num_blocks=args.dyna_num_blocks,
        dyna_num_heads=args.dyna_num_heads,
        dropout=args.dropout,
        mask_limit=args.mask_limit,
    )
    
    # 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)
    dummy_inputs = dict(
        videos=jnp.zeros(
            (per_device_batch_size_for_init, args.seq_len, *image_shape),
            dtype=jnp.float32,
        ),
        action=jnp.zeros(
            (per_device_batch_size_for_init, args.seq_len), dtype=jnp.float32
        ),
        mask_rng=_rng,
    )
    rng, _rng = jax.random.split(rng)
    init_params = genie.init(_rng, dummy_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=genie.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, genie, 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()