From efc82210ef7c2daa487ee43092ce928183b26bd0 Mon Sep 17 00:00:00 2001 From: anurag2796 <97580958+anurag2796@users.noreply.github.com> Date: Sat, 8 Aug 2026 01:48:54 -0400 Subject: [PATCH] Zero grad_input in max_pool2d_with_indices_backward max_pool_backward_impl scatter-adds into grad_input, writing only the argmax positions. The kernel only resizes grad_input and never clears it, and the memory planner recycles arena buffers across ops and iterations, so every non-argmax element accumulates onto stale data. ATen performs the same accumulation but zeroes first, in max_pool2d_with_indices_backward_out_cpu (aten/src/ATen/native/DilatedMaxPool2d.cpp); aten/src/ATen/native/cpu/MaxPoolKernel.cpp holds the identical += loop. The loop was ported here without the zeroing. Any trainable graph containing Conv2d -> MaxPool2d therefore gets corrupted weight gradients: a 67k-param CNN explodes to NaN within 3 steps, while an 11.2M-param ResNet-18 produces no NaN at all and silently fails to converge. Measured with identical flags and a fixed batch, only this file differing: Conv2d->MaxPool2d->Linear goes from 2.309544 -> NaN (291/300 NaN steps) to 2.309544 -> 0.003952 (0 NaN), and five other trainable models (strided-conv, conv-only, pool-only, MLP) produce byte-identical loss curves. Confirmed on macOS arm64 and on a Snapdragon 845 handset; step latency cost is +0.06%. Fixes #21686 --- .../portable/cpu/op_max_pool2d_with_indices_backward.cpp | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/kernels/portable/cpu/op_max_pool2d_with_indices_backward.cpp b/kernels/portable/cpu/op_max_pool2d_with_indices_backward.cpp index 99dc8a89293..46f59bba89f 100644 --- a/kernels/portable/cpu/op_max_pool2d_with_indices_backward.cpp +++ b/kernels/portable/cpu/op_max_pool2d_with_indices_backward.cpp @@ -9,6 +9,8 @@ #include #include +#include + namespace torch { namespace executor { namespace native { @@ -171,6 +173,13 @@ Tensor& max_pool2d_with_indices_backward_out( static constexpr auto name = "max_pool2d_with_indices_backward.grad_input"; + // max_pool_backward_impl scatter-adds into grad_input (`grad_input_ptr[maxindex] += ...`), writing + // only the argmax positions. resize_tensor does not clear the buffer and the memory planner recycles + // arena allocations across ops and iterations, so every other element would accumulate onto stale + // data. ATen zeroes gradInput before dispatching the identical loop + // (aten/src/ATen/native/DilatedMaxPool2d.cpp). + memset(grad_input.mutable_data_ptr(), 0, grad_input.nbytes()); + ET_SWITCH_FLOATHBF16_TYPES(input.scalar_type(), ctx, name, CTYPE, [&]() { max_pool_backward_impl(grad_input, grad_output, indices); });