Micro-Batch Size Selection

What a micro-batch actually is

Fix the vocabulary first, because the whole topic collapses into confusion otherwise. The micro-batch size b_m is how many sequences go through one forward+backward on one device before any synchronization. The global (or effective) batch is what the optimizer sees per step. They are linked by two multipliers:

B_global = b_m × g × d

  b_m = micro-batch size   (fits in one device's memory)
  g   = grad-accum steps   (micro-batches summed before optimizer.step)
  d   = data-parallel replicas

The point of separating b_m out is that B_global is a statistical choice — it controls gradient noise and learning-rate scaling — whereas b_m is a physical one. You pick B_global for convergence, then choose the largest b_m that fits and let g and d make up the difference. This article is about that second, physical choice, which the siblings on gradient accumulation and effective batch deliberately leave open.

The memory it commands: activations scale linearly

Weights, gradients, and optimizer state do not depend on the micro-batch — they are fixed by the parameter count. What b_m controls is activation memory: the intermediate tensors saved on the forward pass so the backward pass can compute gradients. For a transformer, these scale as:

A ∝ b_m × s × h × L

  s = sequence length     h = hidden size     L = number of layers

The dependence on b_m is strictly linear: double the micro-batch, double the activation footprint. That linearity is why the micro-batch is the memory knob of first resort — it is the one term in the peak-memory budget you can turn down without touching the model, the sequence length, or the numerics. It is also why activation checkpointing (recomputing activations in the backward pass instead of storing them) is the natural partner: checkpointing attacks the same b_m × s × h × L term from the other side, trading compute for the memory a large micro-batch would otherwise demand.

Why decouple physical from statistical batch

If memory were infinite you would never think about micro-batches: you would run the whole global batch in one shot. Micro-batching exists precisely because it is not. By setting b_m to the largest slice that fits and accumulating g of them, you reproduce the gradient of a batch far larger than any single device could hold.

The consequence worth internalizing: micro-batch size is a systems decision, not a learning decision. Changing b_m while holding B_global fixed (by adjusting g) leaves the mathematics of training identical — same gradients, same convergence, same final model, up to floating-point summation order. So you are free to tune b_m purely for throughput and memory, which is exactly what the rest of this article does. The only caveat is batch-norm-style layers that compute statistics over the micro-batch, but transformers use layer norm, which is per-token and immune to this.

Micro-batch size versus kernel efficiency

Here is the tension that stops you from simply setting b_m = 1 to save memory. The dominant operations are large matrix multiplies (GEMMs) of shape roughly [b_m · s, h] × [h, h]. Their efficiency is governed by arithmetic intensity — FLOPs performed per byte of memory moved:

intensity ≈ (b_m · s · h) / (h + b_m · s)   FLOP per element loaded
small b_m  → memory-bound  (reloading weights dominates, cores idle)
large b_m  → compute-bound (weights amortized over many rows)

When b_m · s is tiny, the GEMM spends its time streaming the weight matrix from memory and the arithmetic units starve — you are on the sloped part of the roofline. As b_m grows, each loaded weight is reused across more rows, intensity rises, and the kernel climbs toward the hardware’s peak FLOP ceiling. This is why a micro-batch of 1 wastes an accelerator: not because the math is wrong, but because the hardware is designed to be fed wide tiles.

The throughput curve: rise then plateau

Put memory and efficiency together and the throughput-versus-b_m curve has a characteristic shape. It rises steeply at first — each extra row in the micro-batch is nearly free once the weights are resident — then flattens as the GEMM saturates the compute units. Past the knee, larger micro-batches buy little throughput while consuming linearly more memory.

A rough worked example. Suppose at b_m = 1 a step processes 1,000 tokens/s (badly memory-bound), at b_m = 4 it reaches 3,200, at b_m = 8 it reaches 5,000, and at b_m = 16 it reaches 5,600 — but b_m = 16 needs 4× the activation memory of b_m = 4. The efficient operating point is the knee: around b_m = 8 here, where you have captured most of the utilization win before the curve goes flat. The practical recipe is to grow b_m until throughput stops improving meaningfully or you run out of memory — whichever comes first.

The pipeline bubble: micro-batches as bubble filler

Micro-batching earns its keep most dramatically in pipeline parallelism, where the model’s layers are split across p stages on p devices. Run a single batch through and only one stage is busy at a time while the others wait — the idle time is the bubble. The fix is to chop the batch into m micro-batches and stream them through the pipe so stages overlap. Under the classic GPipe schedule the wasted fraction is:

bubble_fraction = (p - 1) / (m + p - 1)

Read this carefully, because it is the crux. The bubble depends on the number of micro-batches m, not their size. With p = 4 stages and only m = 4 micro-batches, the bubble is 3/7 ≈ 43% — nearly half the pipeline idle. Push to m = 32 and it falls to 3/35 ≈ 9%. More, smaller micro-batches fill the pipe.

The bind: smaller micro-batches shrink the bubble but hurt kernels

Now the two forces collide. For a fixed global batch, m × b_m (per replica) is roughly constant, so the only way to get more micro-batches m — and a smaller bubble — is to make each micro-batch smaller. But we just saw that smaller micro-batches run less efficient GEMMs. Pipeline parallelism therefore pulls b_m down while kernel efficiency pulls it up.

The resolution is that these two costs live on different curves. The bubble penalty is large and steep when m is small (below roughly 4×p), then flattens; the kernel-efficiency penalty is small until b_m gets genuinely tiny. So the usual sweet spot is a micro-batch small enough that m ≥ 4p or so — enough to bury the bubble — but not so small that individual GEMMs fall off the roofline. You are balancing an idle-hardware cost against an underfed-hardware cost, and both are forms of wasted FLOPs.

Interaction with activation memory in the pipeline

There is a third pressure that the bubble formula hides. In a pipeline, activations for every in-flight micro-batch must be stashed until its backward pass arrives. Under naive GPipe, a stage can have up to m micro-batches worth of activations resident at the peak — so cranking m up to kill the bubble balloons memory right back up.

This is what 1F1B (one-forward-one-backward, the PipeDream-Flush schedule) fixes. By interleaving a backward pass as soon as each forward completes, it caps the number of in-flight micro-batches per stage at roughly p rather than m, so peak activation memory becomes:

peak_activations ≈ (in-flight micro-batches) × (b_m · s · h · L_stage)
  GPipe : in-flight ≤ m      1F1B : in-flight ≤ p

1F1B keeps the same bubble fraction as GPipe but decouples it from memory, letting you raise m freely. Even so, b_m still multiplies that peak linearly — so the micro-batch size remains the memory knob, now weighted by how many micro-batches the schedule holds live.

A worked selection

Tie it together with numbers. Say a stage can hold 8 micro-batches of activations before OOM, you run p = 4 pipeline stages with 1F1B, and your target per-replica batch is 64 sequences. Because 1F1B keeps only about p = 4 micro-batches live, memory is comfortable, so you are free to choose b_m for the bubble-versus-kernel balance.

Try b_m = 2, giving m = 32: bubble is 3/35 ≈ 9% — excellent — but at b_m · s = 2s the GEMMs may be a touch memory-bound. Try b_m = 4, giving m = 16: bubble is 3/19 ≈ 16% — still fine — and the wider tiles run kernels near peak. The b_m = 4 point usually wins: the extra 7 points of bubble are cheaper than the kernel efficiency lost at b_m = 2. The lesson is that you rarely minimize either term alone — you find where their sum is smallest.

Micro-batching on CPUs and small models

For CPU inference and small language models the calculus shifts but the principles hold. A CPU has far less memory bandwidth relative to its compute than an accelerator, and no massive parallel array to feed, so the throughput knee arrives at a much smaller b_m. Batched decoding of several prompts at once still helps — it amortizes the cost of streaming the weight matrices from RAM across multiple sequences, which is the dominant cost of memory-bound CPU matmuls — but the win saturates quickly.

Cache behavior matters more here than on a GPU: a micro-batch whose activation working set overflows the last-level cache can actually slow things down, so the efficient b_m is often small (frequently 1–8 for a modest SLM). Pipeline bubbles are usually a non-issue on a single CPU socket, so the CPU story is almost entirely the memory-versus-kernel trade with the pipeline term dropped — measure the tokens/s curve and stop at its knee.

Common pitfalls

A few mistakes recur. Confusing micro-batch with global batch: changing b_m for speed and forgetting to adjust g silently changes your optimization, so hold B_global fixed unless you mean to retune. Chasing peak GEMM efficiency blindly: the largest micro-batch that fits is often past the throughput knee — you pay linear memory for sub-linear speed. Ignoring the bubble in pipelines: a large, efficient micro-batch with too few micro-batches can leave 30–40% of the pipeline idle, dwarfing any kernel gain.

And the subtle one: assuming more micro-batches are free. Under GPipe they cost activation memory linearly; only a 1F1B-style schedule breaks that link. Always name which schedule you are on before reasoning about how many micro-batches you can afford — the memory ceiling and the bubble floor are set by different mechanisms, and the micro-batch size is the one term that presses on both at once.