max_pool2d_with_indices_backward does not zero grad_input, corrupting gradients for any Conv2d->MaxPool2d trainable graph

[anurag2796]

🐛 Describe the bug

max_pool2d_with_indices_backward in the portable CPU kernels scatter-adds into grad_input without
zeroing it first, so the gradient accumulates onto whatever the memory planner left in the arena.

Any trainable graph containing Conv2d -> MaxPool2d produces corrupted weight gradients. On a small CNN
this explodes to NaN within 3 steps; on a realistically-sized model it does not produce NaN at all
and instead degrades silently into non-convergence, which is the more dangerous failure mode.

Root cause

kernels/portable/cpu/op_max_pool2d_with_indices_backward.cpp — the accumulation loop touches only the
argmax positions:

// max_pool_backward_impl
grad_input_ptr[maxindex] += grad_output_ptr[index];   // only argmax positions are written

and the entry point only resizes the output, never clears it:

ET_KERNEL_CHECK(
    ctx,
    resize_tensor(grad_input, input.sizes()) == Error::Ok,   // does not zero the buffer
    InvalidArgument,
    grad_input);

ET_SWITCH_FLOATHBF16_TYPES(input.scalar_type(), ctx, name, CTYPE, [&]() {
  max_pool_backward_impl<CTYPE, false>(grad_input, grad_output, indices);
});

Every non-argmax element therefore retains stale bytes, and ExecuTorch recycles arena buffers across ops
and across iterations.

ATen does the same accumulation but zeroes first —
aten/src/ATen/native/DilatedMaxPool2d.cpp:

TORCH_IMPL_FUNC(max_pool2d_with_indices_backward_out_cpu)(...) {
  gradInput.zero_();                                   // <-- this
  max_pool2d_backward_kernel(kCPU, const_cast<Tensor&>(gradInput), gradOutput, indices);
}

and aten/src/ATen/native/cpu/MaxPoolKernel.cpp#L524
is the character-for-character same += loop. The loop was ported; the zero_() was not.

Why the symptom pattern looks strange (and is fully explained by this)

observation	explanation
step-0 loss is bit-exact vs a PyTorch oracle	forward path is untouched
Conv2d -> MaxPool2d diverges	polluted grad_input feeds convolution_backward -> corrupt conv weight grads
MaxPool2d -> Linear alone trains fine	maxpool's grad_input is the grad w.r.t. the network input; it is discarded and never reaches a parameter
Conv2d alone trains fine	no max_pool2d backward in the graph
strided-conv downsampling trains fine	no pooling op at all
diverges even at lr=0.001 (50x below a stable lr)	polluted gradient, not a step-size problem
XNNPACK kernels do not help	XNNPACK delegates the forward; backward maxpool still dispatches to this portable kernel

Measured impact

Identical binary flags, identical fixed batch, SGD momentum=0, 300 steps, fp32. Only the kernel .cpp
differs between columns.

Small CNN (Conv2d -> MaxPool2d -> Linear, 67,642 params, B=32, 3x32x32, 10 classes)

	before	after zeroing
loss	2.309544 -> NaN	2.309544 -> 0.003952
NaN steps	291 / 300	0 / 300

Five other trainable models (strided-conv CNN, conv-only, pool-only, MLP, a second strided variant)
produce byte-identical loss curves before and after, so the change is a strict fix rather than a
behavioural shift.

Correctness of the recovered gradient: against an independently-built LiteRT/TF graph on the same
task and init, the patched curve is bit-identical for the first 5 steps (max abs deviation
0.000e+00) and tracks within <1% from step 100 to 300.

ResNet-18 (11,177,538 trainable params, BatchNorm swapped to GroupNorm so the graph is exportable,
2x3x224x224, SGD lr=0.01, 30 steps) — this is the dangerous case:

	first loss	final loss	min loss	NaN steps
before	0.693080	0.814618	0.693080	0
after	0.693080	0.018742	0.018742	0

No NaN is produced. The loss simply wanders and ends above where it started, with min_loss equal
to the initial loss — the model never improves once. Anyone benchmarking on-device training at a
realistic model size would reasonably conclude that on-device fine-tuning does not work, with a
plausible-looking loss curve to support it.

On real hardware: cross-compiled for arm64-v8a (NDK 27.1.12297006) and run on a Snapdragon 845
handset (Android 10), pinned to the big cluster. Same result — 2.309544 -> -nan (291/300 NaN steps)
becomes 2.309544 -> 0.003952 (0 NaN), with all five previously-working models still byte-identical. The
fixed curve agrees with the macOS arm64 host to 3.07e-04 max deviation and within 1 ULP on the final
loss.

Worth noting the broken run is bit-identical across platforms for its first six steps
(2.30954, 12211.5, 1.32748e+24, 1.33672e+12, 28223.4, 8.40619) — the stale arena contents are
deterministic, so this never presents as flaky memory.

Proposed fix

  static constexpr auto name = "max_pool2d_with_indices_backward.grad_input";

+ // max_pool_backward_impl scatter-adds into grad_input, touching only argmax positions, so every
+ // other element keeps whatever the recycled arena buffer held. ATen zeroes gradInput before the
+ // identical loop (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<CTYPE, false>(grad_input, grad_output, indices);
  });

All-zero is the correct bit pattern for every dtype ET_SWITCH_FLOATHBF16_TYPES covers. Measured cost on
the handset: +0.06% median step latency (239.94 ms -> 240.09 ms), peak RSS 10.38 -> 10.48 MB.

Happy to send this as a PR, and I can share the minimal C++ repro harness (fixed batch + per-step loss log against a PyTorch oracle) if that helps.

Related, not fixed here

• avg_pool2d cannot be used in a trainable graph at all: aten.avg_pool2d_backward has no portable
  kernel, so the model fails to load (TrainingModule::named_parameters error 20). AdaptiveAvgPool2d
  is fine — it decomposes and needs no extra backward op.
• BatchNorm2d fails earlier still, at export: aten._native_batch_norm_legit_functional is not in the
  Core ATen opset, so a stock torchvision ResNet cannot be exported as a trainable graph. GroupNorm works.

Together these mean convolution_backward and max_pool2d_with_indices_backward are the only two
backward kernels portable ships, and one of them was silently wrong.

Versions

• ExecuTorch v1.3.1 (verified), and the defect is still present on main and in v1.4.0
  (released 2026-08-07) — checked today.
• Hosts: macOS arm64 (Apple M4 Max), portable kernels, EXECUTORCH_BUILD_KERNELS_OPTIMIZED=OFF
• Device: Snapdragon 845, Android 10, arm64-v8a, NDK 27.1.12297006
• Export path: torch.export.export -> _export_forward_backward -> to_edge().to_executorch()
• Training: extension/training TrainingModule + SGD

cc @larryliu0820 @manuelcandales @JakeStevens