Skip to content

Fix bias arithmetic in Conv2d BatchNorm fusion - #2528

Open
AaravMehta-07 wants to merge 1 commit into
huggingface:mainfrom
AaravMehta-07:codex/conv-batchnorm-bias
Open

AaravMehta-07 wants to merge 1 commit into
huggingface:mainfrom
AaravMehta-07:codex/conv-batchnorm-bias

Conversation

@AaravMehta-07

@AaravMehta-07 AaravMehta-07 commented Oct 13, 2026 •

Copy link
Copy Markdown

What does this PR do?

Preserve inference outputs when fusing an eval BatchNorm2d into Conv2d.

Two bias terms currently break computation preservation:

  • The convolution's original bias is added without the BatchNorm scale. Use (conv_bias - running_mean) * gamma / sqrt(running_var + eps) + beta, rather than leaving conv_bias unscaled.
  • For affine=False, BatchNorm has no additive beta term, so its replacement must be zeros, not ones. The existing fallback adds an unintended constant offset to every output.

Only those two production lines change. Linear/BatchNorm1d fusion, graph eligibility, training-mode behavior and dependency pins are untouched.

Validation

  • Actual upstream 0df84f23 replay: 6 failures, 2 passing controls in the new Conv2d tests.
  • Patched CPU tests: 8 passed, independently repeated by a GPT-6 Luna reviewer and root.
  • Tests compare the complete transformed Torch FX model's outputs to the unfused eval model, and assert BatchNorm removal. They cover convolution bias enabled/disabled, affine normalization enabled/disabled, and float32/float64 with nontrivial means, variances and affine parameters.
  • Repository-wide Black 23.12.1 and Ruff 0.1.5 checks pass; git diff --check passes.

Local environment: Windows, Python 3.11.9, Torch 2.13.0+cpu, Transformers 5.19.0. No pretrained model downloads, mocks of inference/fusion, GPU benchmark or speed claim. Full legacy FX tests were attempted but could not collect because they import the removed transformers.utils.fx module; the new tests use actual torch.fx.symbolic_trace and run offline. This is not a full-suite claim.

Additional isolated legacy compatibility run with Transformers 4.57.6, tokenizers 0.22.2 and Hub 0.36.2: the eight new cases plus the existing offline test_get_parent pass (9 passed, 16 deselected). This overlay is not the current main dependency set and does not change repository pins; no pretrained-model tests were run.

Live duplicate checks for FuseBatchNorm, BatchNorm2d, bn2d and affine fusion found no matching active fix. The original batchnorm-folding PR #348 is already merged; this corrects its bias arithmetic, not a second implementation.

Before submitting

  • Typo/docs-only change: not applicable.
  • Documentation change: existing computation-preservation behavior is restored; no new API.
  • Added necessary output-equivalence regression tests.

AI disclosure: investigated, implemented and tested with OpenAI Codex, independently reviewed/tested by a GPT-6 Luna worker and root. An initial wider linear-fusion candidate was dropped after review identified an unsupported 3D eligibility risk; this PR contains only the verified Conv2d correction. No personal human review/test attestation is claimed.

@AaravMehta-07

Copy link
Copy Markdown
Author

Hi maintainers, whenever you have a chance, I'd appreciate a review of this fix for preserving bias arithmetic in Conv2d/BatchNorm fusion. The eight focused numerical fusion tests passed locally, including float64 and non-affine cases.

Please let me know if you'd prefer a different approach or additional focused coverage. Thank you for your time.

AI-assisted follow-up.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant