Repository: pytorch/benchmark
Issue: Discussion: Batch size automatically
Root cause: Unimplemented Feature Gap

```python 
import torch

def test_batch_size(model, batch_size, input_shape):
    try:
        input_data = torch.rand((batch_size,) + input_shape).cuda()
        output_data = model(input_data)
        del input_data, output_data
        torch.cuda.empty_cache()
        return True
    except RuntimeError as e:
        if 'out of memory' in str(e):
            return False
        else:
            raise e

def find_max_batch_size(model, input_shape):
    if not torch.cuda.is_available():
        raise ValueError("No GPUs available.")

    min_batch_size = 1
    max_batch_size = 1024
    while min_batch_size < max_batch_size:
        mid_batch_size = (min_batch_size + max_batch_size) // 2
        if test_batch_size(model, mid_batch_size, input_shape):
            min_batch_size = mid_batch_size + 1
        else:
            max_batch_size = mid_batch_size
    return max_batch_size - 1

# Let's say we have a simple model:
model = torch.nn.Sequential(
    torch.nn.Linear(100, 64),
    torch.nn.ReLU(),
    torch.nn.Linear(64, 10)
).cuda()

max_batch_size = find_max_batch_size(model, input_shape=(100,))
print("Max batch size: ", max_batch_size)
``` 
This is a silly idea, that increments batch size.  The idea is to change the random data with original datasets.
Another idea is taking the memory of gpu and calculate it




We have an old script that detects proper inference batch size at https://github.com/pytorch/benchmark/blob/main/scripts/proper_bs.py. It requires NVIDIA DCGM, but I think this can be optional.

The idea is to keep increasing the batch size until GPU OOM, for each batch size, measure the GPU hardware flops. And pick the smallest batch size which GPU reaches the highest flops.

> We have an old script that detects proper inference batch size at https://github.com/pytorch/benchmark/blob/main/scripts/proper_bs.py. It requires NVIDIA DCGM, but I think this can be optional.
> 
> The idea is to keep increasing the batch size until GPU OOM, for each batch size, measure the GPU hardware flops. And pick the smallest batch size which GPU reaches the highest flops.


As detailed in the paper [Scaling Language Model Training to a Trillion Parameters Using Megatron](https://arxiv.org/pdf/2104.04473.pdf), the majority of floating-point operations in the model are performed in the matrix multiplications (GEMMs). If we consider only these GEMMs operations, the number of FLOPs per iteration is $F = 96 B s l h^2 (1 + \frac{s}{6h} + \frac{V}{16lh})$ where $B$ is the batch size. 

And, in case we have an estimate of the time spent per iteration `Time_per_iteration_second`, it is possible then to compute the theoretical peak FLOP per second and per GPUs and estimate the GPU usage by comparing. 

The following table shows the training performance of several GPT model sizes (from 1.7B to 1 trillion) pretrained using Megatron-LM library on a SuperPOD cluster with A100 GPUs.
<img src="https://github.com/NVIDIA/Megatron-LM/blob/main/images/cases_april2021.png?raw=true"/>

```python
# Theoretical peak FLOP per second per GPU - with activation checkpointing (2 forwards and 1 backward)
def calculate_Theoretical_peak_FLOP_s_GPU(B,s, l,h,number_GPUs,Time_per_iteration_second):
    # The number of FLOPs per iteration
    F = 96*B*s* l*h*h *(1 + s/ (6*h) + V/(16*l*h))/1e+12
    
    #Theoretical peak FLOP per second per GPU
    PF= (F/Time_per_iteration_second/number_GPUs)
    print("Theoretical peak FLOP/s/GPU: {}\n".format(PF))
    
    # Percentage of theoretical peak FLOP/s on a A100 FP16 (change according the hardware)
    GPU_usage= PF/ 312 *100
    print("Percentage of theoretical peak FLOP/s: {}%".format(GPU_usage))
    
    return PF, GPU_usage
    ```
    
    ```
    The percentage of theoretical peak FLOP/s in the previous function is based on **A100** hardware capabilities in **FP16/BF16 which is 312**. This needs to be updated according to the corresponding Tensor Core GPU performance specs. To learn more about the [Ampere architecture specifications](https://developer.nvidia.com/blog/nvidia-ampere-architecture-in-depth/).


If we consider our previous 18B parameters models pretrained on 16 GPUs with a global batch size of 512, they have a time per iteration of 32.09s. 
Let's compute the theoretical peak FLOP per second per GPU and the percentage GPU utilization:
```
```python
global_batch_size=512
number_GPUs=16
Time_per_iteration_second=32.09

# Considering the 18B parameters model
l=40
h=6144
s=1048
V=50257
    
PF,GPU_usage=calculate_Theoretical_peak_FLOP_s_GPU(global_batch_size,s, l,h,number_GPUs,Time_per_iteration_second)

```

Theoretical peak FLOP/s/GPU: 157.7296879710454

Percentage of theoretical peak FLOP/s: 50.55438717020686%

This idea is extracted from Nvidia

> > We have an old script that detects proper inference batch size at https://github.com/pytorch/benchmark/blob/main/scripts/proper_bs.py. It requires NVIDIA DCGM, but I think this can be optional.
> > The idea is to keep increasing the batch size until GPU OOM, for each batch size, measure the GPU hardware flops. And pick the smallest batch size which GPU reaches the highest flops.
> 
> As detailed in the paper [Scaling Language Model Training to a Trillion Parameters Using Megatron](https://arxiv.org/pdf/2104.04473.pdf), the majority of floating-point operations in the model are performed in the matrix multiplications (GEMMs). If we consider only these GEMMs operations, the number of FLOPs per iteration is F=96Bslh2(1+s6h+V16lh) where B is the batch size.
> 
> And, in case we have an estimate of the time spent per iteration `Time_per_iteration_second`, it is possible then to compute the theoretical peak FLOP per second and per GPUs and estimate the GPU usage by comparing.
> 
> The following table shows the training performance of several GPT model sizes (from 1.7B to 1 trillion) pretrained using Megatron-LM library on a SuperPOD cluster with A100 GPUs. ![](https://github.com/NVIDIA/Megatron-LM/blob/main/images/cases_april2021.png?raw=true)
> 
> ```shell
> # Theoretical peak FLOP per second per GPU - with activation checkpointing (2 forwards and 1 backward)
> def calculate_Theoretical_peak_FLOP_s_GPU(B,s, l,h,number_GPUs,Time_per_iteration_second):
>     # The number of FLOPs per iteration
>     F = 96*B*s* l*h*h *(1 + s/ (6*h) + V/(16*l*h))/1e+12
>     
>     #Theoretical peak FLOP per second per GPU
>     PF= (F/Time_per_iteration_second/number_GPUs)
>     print("Theoretical peak FLOP/s/GPU: {}\n".format(PF))
>     
>     # Percentage of theoretical peak FLOP/s on a A100 FP16 (change according the hardware)
>     GPU_usage= PF/ 312 *100
>     print("Percentage of theoretical peak FLOP/s: {}%".format(GPU_usage))
>     
>     return PF, GPU_usage
>     ```
>     
>     ```
>     The percentage of theoretical peak FLOP/s in the previous function is based on **A100** hardware capabilities in **FP16/BF16 which is 312**. This needs to be updated according to the corresponding Tensor Core GPU performance specs. To learn more about the [Ampere architecture specifications](https://developer.nvidia.com/blog/nvidia-ampere-architecture-in-depth/).
> 
> 
> If we consider our previous 18B parameters models pretrained on 16 GPUs with a global batch size of 512, they have a time per iteration of 32.09s. 
> Let's compute the theoretical peak FLOP per second per GPU and the percentage GPU utilization:
> ```
> 
> ```
> global_batch_size=512
> number_GPUs=16
> Time_per_iteration_second=32.09
> 
> # Considering the 18B parameters model
> l=40
> h=6144
> s=1048
> V=50257
>     
> PF,GPU_usage=calculate_Theoretical_peak_FLOP_s_GPU(global_batch_size,s, l,h,number_GPUs,Time_per_iteration_second)
> ```
> 
> Theoretical peak FLOP/s/GPU: 157.7296879710454
> 
> Percentage of theoretical peak FLOP/s: 50.55438717020686%
> 
> This idea is extracted from Nvidia

or maybe
```python
import torch
import time
import numpy as np
from torchvision.models import resnet50
from torch.profiler import profile, record_function, ProfilerActivity

def measure_time_and_flops(model, input, n_trials=50):
    with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA], record_shapes=True) as prof:
        with record_function("model_inference"):
            for _ in range(n_trials):
                output = model(input)
    torch.cuda.synchronize()  # Wait for all kernels to finish
    total_flops = prof.key_averages().total().cpu_memory_usage()
    return total_flops / n_trials  # Return the average FLOPS

def find_optimal_batch_size(model, input_shape):
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    model = model.to(device)
    model.eval()

    batch_size = 1
    max_flops = 0
    optimal_batch_size = 1

    while True:
        try:
            input_data = torch.rand((batch_size,) + input_shape).to(device)
            flops = measure_time_and_flops(model, input_data)
            if flops > max_flops:
                max_flops = flops
                optimal_batch_size = batch_size
            batch_size *= 2
        except RuntimeError as e:
            if 'out of memory' in str(e):
                torch.cuda.empty_cache()
                break
            else:
                raise e
    return optimal_batch_size, max_flops

model = resnet50()
optimal_batch_size, max_flops = find_optimal_batch_size(model, input_shape=(3,224,224))
print("Optimal batch size: ", optimal_batch_size)
print("Max FLOPS: ", max_flops)

```

This FLOPS algorithm only applies for resnet50 model. We need to use NVIDIA DCGM to count the hardware FLOPS.
Or use the theoretical FLOPS counter Horace proposed at https://dev-discuss.pytorch.org/t/the-ideal-pytorch-flop-counter-with-torch-dispatch/505

> This FLOPS algorithm only applies for resnet50 model. We need to use NVIDIA DCGM to count the hardware FLOPS. Or use the theoretical FLOPS counter Horace proposed at https://dev-discuss.pytorch.org/t/the-ideal-pytorch-flop-counter-with-torch-dispatch/505

The ideal case would be to make NVIDIA DCGM for CUDA and the Horace for CPU or MPS, no?

Hardware FLOPS and theoretical FLOPS are two different metrics.
For Hardware FLOPS, we need to use NVIDIA DCGM.
For theoretical FLOPS, I suggest we could use the `__torch_dispatch__` based FLOPS counter.