[WIP][ROCm] PJRT: Defer device buffer free until the compute stream drains#979
[WIP][ROCm] PJRT: Defer device buffer free until the compute stream drains#979magaonka-amd wants to merge 1 commit into
Conversation
Under the kComputeSynchronized allocation model the host runs ahead of the device, so an AllocatedRawSEDeviceMemory's last compute-stream use may still be in flight when its host-side refcount drops to zero. Freeing the chunk immediately returns it to the allocator, where a subsequent allocation (possibly on a different stream) can reuse and overwrite memory that a live kernel still owns -- a use-after-free. This was observed on ROCm as an intermittent all-zero backward-filter conv gradient (tests/batching_test.py::testConvGeneralDilated): the freed buffer was re-served to the conv autotuner, whose candidate zero-initialized scratch (SetTensor(dW, 0)) aliased the still-live gradient. The early free is platform-agnostic (the contract is shared with CUDA). Defer the deallocation via LocalDeviceState::ThenExecuteCallback on a callback stream that waits on the compute stream, so the chunk is not returned to the allocator until the producing kernel has finished. The compute stream only records a lightweight event (it is not stalled) and the host never blocks, so the allocator remains asynchronous. Falls back to an immediate free if the callback cannot be enqueued.
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
| if (s.ok()) { | ||
| return; | ||
| } | ||
| // Fall through to an immediate free if the callback could not be | ||
| // enqueued. | ||
| } |
There was a problem hiding this comment.
nit: When ThenExecuteCallback fails, the error status s is silently discarded before falling through to the immediate free. Consider logging it so operators can diagnose why the deferred path was skipped:
| if (s.ok()) { | |
| return; | |
| } | |
| // Fall through to an immediate free if the callback could not be | |
| // enqueued. | |
| } | |
| if (s.ok()) { | |
| return; | |
| } | |
| LOG(WARNING) | |
| << "Could not enqueue deferred deallocation, falling back to " | |
| "immediate free: " | |
| << s; | |
| // Fall through to an immediate free if the callback could not be | |
| // enqueued. |
| [allocator, device_ordinal, memory]() { | ||
| absl::Status status = allocator->Deallocate(device_ordinal, memory); | ||
| if (!status.ok()) { | ||
| LOG(ERROR) << "Deferred buffer deallocation failed: " << status; |
There was a problem hiding this comment.
observation: The lambda captures allocator (a raw pointer to DeviceAddressAllocator) by value. This is safe because LocalDeviceState::~LocalDeviceState() calls SynchronizeAllActivity() which drains all callback streams before the allocator (owned by PjRtStreamExecutorClient::owned_allocator_) is destroyed. However, the correctness depends on the member declaration order in PjRtStreamExecutorClient (owned_devices_ before owned_allocator_), which controls destruction order. Consider adding a comment on this line noting the lifetime dependency, so a future reordering doesn't silently introduce a use-after-free in the callback.
| // different stream) reuse and overwrite memory that a live kernel still | ||
| // owns. Defer the free until the compute stream drains past this point. | ||
| // ThenExecuteCallback runs the deallocation from a callback stream that | ||
| // waits on the compute stream, so the compute stream only records a | ||
| // lightweight event (it is not stalled) and the host never blocks. | ||
| if (local_device_->allocation_model() == | ||
| LocalDeviceState::kComputeSynchronized) { | ||
| absl::Status s = local_device_->ThenExecuteCallback( |
There was a problem hiding this comment.
question: Under kComputeSynchronized, the original contract (per local_device_state.h) allows freeing a buffer once its last compute-stream use is enqueued (not completed). This change strengthens that by deferring the free until the compute stream actually drains past the deallocation point. While this is the right fix for the race, it does extend effective memory lifetimes — under heavy allocation churn, freed buffers are withheld from the pool until the compute stream catches up, which could increase peak device memory. The compute_semaphore_ bounds the host-ahead-of-device window and limits this effect, but it would be useful to call out this trade-off (perhaps in the comment block above) for future readers.
Review SummaryClean, well-motivated fix for a real use-after-free race under 3 inline comments posted:
No correctness bugs found. The change is small, well-scoped, and the inline comments are all suggestions for hardening, not blockers. |
github-actions
Bot
removed
the
claude-review
magaonka-amd
changed the title
1 participant
📝 Summary of Changes
Defer device-buffer deallocation under the
kComputeSynchronizedallocation model until the compute stream drains past the buffer's last use, instead of returning the chunk to the allocator the moment its host-side refcount hits zero.This closes a use-after-free where a freed chunk is re-served to a new allocation (possibly on another stream) and overwritten while a kernel that still uses it is in flight on the GPU.
🎯 Justification
On ROCm the eager backward-filter conv gradient in
tests/batching_test.py::testConvGeneralDilated(JAX) intermittently came back all-zero (42/450 elements). Root cause:~AllocatedRawSEDeviceMemoryfrees the buffer at host-enqueue time, but under async dispatch the producing kernel is still running. The conv autotuner then drew scratch from the same shared BFC pool on a non-compute stream, and a candidate'sSetTensor(dW, 0)zeroed the still-live gradient. The early free is the prerequisite; deferring it removes the hazard.In my opinion dilated conv operation being bit heavy and slow with MIOpen exposes this rare race condition. In theory bug exist in CUDA too but none of the Unit tests from JAX side expose this problem.
One more thing, dilated conv is not the only test that affected by this bug, there are many other conv related tests run into this flakyness , but among those dilated conv has had more frequency.
For more info on what the test is , math behind it and bit more detail on this bug : https://amd.atlassian.net/wiki/spaces/MLSE/pages/1747138095/Conv+flake+test+in+JAX+CI
🚀 Kind of Contribution
🐛 Bug Fix
Tests to highlight the problem
JAX test that exposes this bug is flaky and extremly rare. It only comes up in our CI nodes.
So I had to mimic a standalone, deterministic reproducer is included below as an example. It uses the real
MultiDeviceAdapter/tsl::BFCAllocator, and a longMemset32"delay" keeps a producer kernel in flight while a buffer is freed and re-served:Immediate free: the allocator re-serves the freed buffer's exact address
while the producer's event is still
kPending, and a cross-streamMemZeroclobbers it → the original buffer reads back
0.Deferred free (this fix): the chunk is withheld during the in-flight
window, so the new allocation gets a different address and the data survives.
This example exaggerates situation a bit to highlight the bug.
Sample output:
expand the section below to see XLA level test that shows the problem.
Standalone reproducer (async_dealloc_uaf_test.cc)
🧪 Execution Tests
testConvGeneralDilatedon a 4-GPU gfx950 runner (ROCm 7.2.0) no longer producesthe all-zero gradient with this fix applied.