Gradient Accumulation Memory Limits and Batch Normalization Parity

Gradient Accumulation Behavior

PyTorch supports simulating larger batch sizes by summing gradients over multiple forward and backward passes before executing an optimizer step. This technique is intended to keep memory consumption proportional to the micro-batch size rather than the effective total batch size.

While the .grad attribute accumulates values across calls to .backward(), certain layer types may not maintain mathematical parity with a true large-batch implementation. Specifically, Batch Normalization layers calculate running statistics based on the immediate micro-batch rather than the accumulated total.

This creates a discrepancy between the gradient magnitude and the normalization statistics when comparing a single large batch to an equivalent accumulation cycle.

Technical Constraints

• Memory overhead must remain constant across accumulation steps.
• Gradient magnitude must be scaled by the number of accumulation steps to maintain average consistency.

Does PyTorch provide a native mechanism to synchronize Batch Normalization statistics across accumulation steps to ensure parity with large-batch training? If not, what is the documented behavior for maintaining statistic stability in memory-constrained environments?