[pipelining] Fix backward hook accumulation on recv buffer activations - #185355

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/185355

• 📄 Preview Python docs built from this PR
• 📄 Preview C++ docs built from this PR
• 🤖 See Auto PR Triage runs for this PR
• ❓ Need help or want to give feedback on the CI? Visit the bot commands wiki

Note: Links to docs will display an error until the docs builds have been completed.

⚠️ 4 Awaiting Approval

As of commit a971006 with merge base 2812db3 ():

This comment was automatically generated by Dr. CI and updates every 15 minutes.

The following ciflow label(s) have been added but CI has not been triggered yet because the workflows are awaiting approval:

• ciflow/trunk
• ciflow/torchtitan

Once a maintainer approves the workflows (scroll to the bottom of the PR page), the corresponding CI jobs will be triggered automatically. Please ping one of the reviewers if you do not have access to approve and run workflows.

@pytorchbot label "topic: not user facing"

@wconstab @wanchaol @howard-huang could you please take a look when you get a chance? This fixes hook accumulation on recv buffer activations (issue #185331). Happy to address any feedback!

@claude can you see if there is an alternative that avoids a .clone on the recv but also prevents leaking registered hooks between microbatch executions?

Claude finished @wconstab's task in 3m 48s —— View job

Investigating alternatives to .clone() for recv buffer hook leakage

• Read the current PR changes and understand the fix
• Examine the _retrieve_recv_activations method and surrounding context
• Investigate how backward hooks work on views vs clones vs detached tensors
• Analyze safety of detach() approach given pipeline scheduling
• Propose alternative approaches

Analysis

The core question: can we avoid the .clone() data copy while still preventing hooks from accumulating on the recv buffer across steps?

Why .clone() works: It creates a completely independent tensor — new storage, new identity. Hooks on the clone never touch the buffer.

Alternative: info.buffer.detach().requires_grad_(effective_requires_grad)

detach() creates a new tensor object that shares storage with the buffer but is disconnected from the autograd graph. Each call to _retrieve_recv_activations would produce a fresh tensor identity, so hooks registered on one step's activation don't leak to the next step's activation. Key properties:

1. Fixes the bug — the returned tensor is a new Python object each step, so hooks registered on it don't accumulate on info.buffer.
2. Is a leaf tensor — detach() creates a tensor with no grad_fn, and .requires_grad_(True) makes it a leaf. Satisfies the activation.is_leaf check at line 643.
3. No data copy — shares storage with the buffer (zero-copy).
4. Gradient flow still works — stage_backward reads val.grad and clears it (line 427-434 of _backward.py), which operates on the detached activation tensor, not the buffer.
5. Per-chunk buffers are safe — each microbatch has its own buffer (line 1247 creates fresh recv info per chunk), so no cross-chunk storage aliasing within a step.
6. Cross-step safety — between steps, fwd_cache.clear() drops references to old activations. The next step's irecv fills the buffer before _retrieve_recv_activations runs, so the shared storage has fresh data.

Why not a view()? Since PyTorch 2.0, backward hooks on views propagate to the base tensor via the autograd view-tracking machinery. A view would NOT fix the bug — hooks would still effectively be on the base buffer.

Potential concern with detach() + storage sharing: If a user's hook captures a reference to the activation tensor and inspects its data after the next step's irecv fills the buffer, they'd see stale/overwritten data. However, this same issue existed in the original pre-fix code (where activation IS info.buffer), so detach() is no worse. Only .clone() protects against this edge case.

Concern @wconstab raised about pinned/nvshmem buffers: With detach(), the buffer itself remains the irecv target, and its allocation properties (pinned memory, special allocator, etc.) are fully preserved. The detached activation just references the same storage — it doesn't interfere with the buffer's lifecycle or allocation strategy.

Recommendation

info.buffer.detach().requires_grad_(effective_requires_grad) is a valid zero-copy alternative that fixes the hook accumulation bug. It's semantically equivalent to the original behavior (shared storage) but with a fresh tensor identity each step so hooks don't leak.

The tradeoff vs .clone():

Since the "user holds ref across steps" scenario is unusual and was already broken before this PR, detach() seems like the better default — it preserves the performance and allocation properties of the buffer while fixing the actual bug.

thanks for the suggestion @wconstab! Switched from .clone() to .detach() — you're right that clone would lose any pinned memory or custom allocator properties on the recv buffer, which is not ideal.

.detach() gives us the same fix (fresh tensor identity each step so hooks don't accumulate) but is zero-copy and fully preserves the buffer's underlying memory allocation. Updated the comments in both the fix and the test to reflect this.

@karthik-0306 Can you fix the lint?

Thanks @sanketpurandare! Fixed the lint and moved the test into test_stage.py reusing MultiMLP and existing test infrastructure.

@karthik-0306 Please check the test failure: trunk / linux-jammy-rocm-py3.10-mi355 / test (distributed, 1, 4, linux.rocm.gpu.gfx950.2) (gh)
test/distributed/pipelining/test_stage.py::StageTest::test_recv_buffer_hooks_do_not_accumulate

Your own test is failing

@sanketpurandare Root cause found — pipeline() was called without split_spec, so it kept the model as a single stage regardless of world_size, causing the "Pipeline group size N cannot be larger than number of stages 1" error on the 4-GPU runner. Fixed by passing MultiMLP's split_spec through, same pattern test_tracer already uses.

The other failures (torchtitan integration tests, win-vs2022 builds) look unrelated to this PR — pure Python diff, no build/dependency changes touched. The Windows build log showed a pip self-modification error, and torchtitan runs full training jobs that are known to be flaky independent of the change under test. Happy to dig further if you think otherwise, but wanted to flag in case these can just be re-triggered.

Pushed the split_spec fix, should be green on the real failure now.

Looks like this PR hasn't been updated in a while so we're going to go ahead and mark this as Stale.
Feel free to remove the Stale label if you feel this was a mistake.
If you are unable to remove the Stale label please contact a maintainer in order to do so.
If you want the bot to never mark this PR stale again, add the no-stale label.
Stale pull requests will automatically be closed after 30 days of inactivity.

Removing stale — this PR is still active and ready for review.

175/214 checks passing. The 8 remaining failures are all infrastructure issues unrelated to this change (torchtitan flakyintegration tests, Windows pip self-modification errors, and an aarch64 container infra failure). The distributed pipelining tests themselves pass cleanly including on ROCm.

Happy to address any further feedback from @wconstab or @sanketpurandare.

@pytorchbot label "no-stale"

@wconstab @sanketpurandare I've rebased this onto current main and updated the PR.

Since #196463 (just-in-time recv buffer allocation), each receive gets a fresh buffer that is released once consumed, so hooks can no longer accumulate on a persistent buffer. That makes the .detach() change in stage.py unnecessary, so I've dropped it.

The PR now only adds a regression test, test_recv_buffer_hooks_do_not_accumulate, so the behavior stays guarded. I haven't run it on GPU locally, so I'm relying on CI for that. If it fails on current main, the bug is still there and I'll look at it again.
If you'd rather not land a test-only change, I'm fine closing this PR.