The increase in GPU memory usage with an increase in batch size is primarily due to the computational and memory requirements of processing larger batches of data in parallel. Here’s a breakdown of the key factors:

  1. Parallelism and Concurrency:
    • Modern deep learning frameworks and GPUs are designed to exploit parallelism. When you increase the batch size, you are essentially processing more data points simultaneously.
    • GPUs excel at parallelizing matrix operations and can perform computations on multiple data points in parallel. This is achieved through the use of parallel processing units (CUDA cores in the case of NVIDIA GPUs).
  2. Memory for Intermediate Values:
    • During the forward and backward passes in training a neural network, intermediate values such as activations, gradients, and other tensors need to be stored in memory.
    • Larger batch sizes result in more intermediate values being stored simultaneously, leading to increased memory requirements.
  3. Optimization for Efficiency:
    • GPUs are optimized for processing larger batches efficiently. Processing a small batch might not fully utilize the parallel processing capabilities of the GPU, leading to underutilization of the available resources.
    • Larger batch sizes allow the GPU to efficiently handle a larger number of parallel tasks, making better use of its computational power.
  4. Memory Footprint of Model Parameters:
    • The model’s parameters (weights and biases) also contribute to GPU memory usage. With a larger batch size, the model processes more parameters at once, requiring additional memory to store them during training.
  5. Memory Bandwidth:
    • The speed at which data can be transferred between the GPU’s memory and the GPU cores is limited by memory bandwidth. Larger batch sizes can help to better utilize this bandwidth.

It’s important to note that while increasing the batch size can improve training efficiency, there is a trade-off. Very large batch sizes may lead to GPU memory exhaustion, limiting the size of models or the overall amount of data that can be processed in parallel. Finding an optimal batch size often involves experimentation and depends on the specific model, dataset, and available hardware.

Links - gpu utilization for training LLMs - fast-ai forum - torch forum on gpu and batch size