SIGN IN SIGN UP

Preserve BatchNorm running state in composable quantizer

Summary:
Quantization-aware training with the composable `PatternQuantizer` incorrectly observed BatchNorm affine parameters and running-state buffers as activation inputs. The inserted fake-quant nodes broke the connection between `aten.batch_norm.default` and the registered `running_mean` and `running_var` buffers, so training used batch statistics while evaluation consumed stale running statistics.

Keep `aten.batch_norm.default` arguments 1-4 directly connected to registered state by excluding them from activation observation while retaining activation and output quantization. Add Arm and Sleep/ModAI regression tests that verify the state operands remain registered and the running statistics update during QAT.

Differential Revision: D112989631
D
Deniz Kilinc committed
6bd03553e7cf5bac24498940fafaae49ea33285f
Parent: cd380e7
Committed by Deniz Kilinc <denizkilincx@gmail.com> on 7/21/2026, 9:18:41 PM