Repository navigation
Add skip_h2d_for_constants to keep constants in device memory - #24144
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/24144
Note: Links to docs will display an error until the docs builds have been completed. ❌ 7 Unclassified FailuresAs of commit 0cf1823 with merge base 0549e1e ( UNCLASSIFIED FAILURES - DrCI could not classify the following jobs because the workflow did not run on the merge base. The failures may be pre-existing on trunk or introduced by this PR:
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
fb3f67e to
f1833e1
Compare
|
@claude review |
|
Claude finished @shoumikhin's task in 4m 51s —— View job The approach is sound, but please fix a few things before merging: Detailed reviewPR Review: #24144I reviewed only the last two commits ( SummaryWith
API Design
Code Quality
Testing
Performance
RecommendationRequest Changes: the main blockers are the missing opt-out (or documented cost) for keeping constants on the device, the duplicate storage of a shared weight, and making the CUDA test actually check the feature. |
f1833e1 to
fc3ce9c
Compare
|
Thanks for the review. Changed in the new push:
Not changed, and why:
|
|
@claude review |
|
Claude finished @shoumikhin's task in 3m 52s —— View job Before this can merge, please delete the Detailed reviewPR Review: #24144SummaryThis PR adds the opt-in option Infrastructure
API Design
Code Quality
Testing
RecommendationRequest Changes: the stray |
fc3ce9c to
cd8691c
Compare
|
Thanks. Changed in the new push: the stray Not changed, and why:
|
|
@claude review |
|
Claude finished @shoumikhin's task in 4m 17s —— View job I found no blockers in the last two commits; what is left is one point about the Detailed reviewPR Review: #24144SummaryThis PR adds the opt-in option API Design
Code Quality
Testing
RecommendationApprove: none of the findings blocks the merge. The |
|
Thanks. Answers, with no code change:
|
| it to `False` while asking for either skip raises a `ValueError`, because copy insertion happens | ||
| during device-aware memory planning. | ||
|
|
||
| **With `skip_h2d_for_method_inputs`, an input that a device delegate reads stays on that device.** |
There was a problem hiding this comment.
can you keep this .md file more succinct and readable by humans
There was a problem hiding this comment.
Shortened, see the new version. On top of the shorter text from #24088, this PR now adds only one short paragraph on skip_h2d_for_constants, what it does and what it costs, and a pointer to it from the list of copies that stay.
|
it's hard to review stacked PR. not clear what is in the previous and what is new Can you use ghstack? |
|
@claude review About the last review: I kept |
|
Claude finished @shoumikhin's task in 4m 55s —— View job Nothing in this PR's own six commits ( Detailed reviewPR Review: #24144SummaryThis PR adds the opt-in I reviewed I accept the author's reason for keeping I checked and dropped these findings:
API Design
RecommendationApprove: the one finding is a wording fix in the docs and does not block the merge. |
1f367f7 to
7f587a9
Compare
|
@claude review |
|
Claude finished @shoumikhin's task in 6m 21s —— View job I found nothing to fix in this PR's own six commits ( Detailed reviewPR Review: #24144SummaryThis PR adds the opt-in option I reviewed I left out points that earlier rounds already settled with the author:
I checked and dropped these new findings:
RecommendationApprove: no findings remain. This is still a stacked PR on #24088, so it needs the planned rebase onto |
7f587a9 to
be1bb26
Compare
The next commit copies the initial data of a tensor planned in device memory there when a method loads, through the device allocator. On CUDA that copy can still be running when copy_host_to_device() returns: it is asynchronous on a caller stream, and from pageable memory even a plain cudaMemcpy may return before the data lands. Two things can then go wrong. The host buffer the data came from may be freed while the copy still reads it. And the first run reads the tensor on its own stream; if that stream is not ordered after the copy, for example when the method loads with one caller stream and runs with another, or runs on a non-blocking stream such as a green context stream, it can read the tensor before its data is there. DeviceAllocator gets synchronize(), which waits for the copies that the allocator made on the calling thread, with the same per-thread settings. The default does nothing, for allocators whose copies finish before they return. CudaAllocator switches to the device and waits for the stream its copies use: the caller stream when one is set, otherwise the default stream of the device. It switches even with a caller stream, because a handle such as nullptr or cudaStreamPerThread names a stream of the current device. Without a caller stream the copy runs on the legacy default stream, so waiting for it can also wait for work that other threads queued on blocking streams of that device. test_cuda_allocator checks that synchronize() waits for the caller stream and for the default stream, and, with two GPUs, for the default stream of device 0 while device 1 is current and the caller stream is nullptr.
A memory-planned tensor can have initial data, such as a mutable buffer with a starting value. The runtime loads that data into the tensor's planned memory when it loads the method, with the data loader or the named data map. Both write with the CPU, so they cannot fill a planned buffer that lives in device memory, such as CUDA memory. The tensor parser now checks the device of the planned buffer, which the memory manager's planned memory records. For a device buffer it loads the data into host memory and copies it with the device allocator registered for that device, once, at method load. It then waits for the copy with the allocator's synchronize(), so the host buffer can be freed and the first run reads the right data. A host buffer is loaded as before, and the copies made on every call are unchanged. An empty tensor now loads nothing. When all the initial data of a program is empty, the program has no mutable segment at all, so loading such a method failed with "No mutable data segments in program", also for a host buffer. The emitter test now loads such a program. Program gets a private load_mutable_subsegment(), which loads a part of a mutable segment into a new buffer and checks the same bounds as load_mutable_subsegment_into(). Both share the code that finds the part. The program test loads the same weights both ways and expects the same bytes, and expects both to fail for bad indices and sizes.
With skip_h2d_for_method_inputs, a method input that a device delegate reads stays in device memory. A constant, such as a weight or a buffer the method does not change, was still copied to the device on every call, because its data is stored in the program, on the host. The new opt-in PropagateDeviceConfig option skip_h2d_for_constants keeps a constant that only delegates on one device read in device memory. PropagateDevicePass tags it with that device instead of adding a copy. Memory planning gives it planned device memory for the whole method, also when graph inputs are not planned, because the method reads it on every call. The emitter stores its data as the initial state of that memory, as it does for a mutable buffer with an initial value, and the runtime loads it there once, when it loads the method. Each method that reads the constant keeps its own copy, so this is opt-in. A constant that anything else reads, such as a host op or a delegate on another device, is still copied, with the same warning, and error_on_per_call_copies turns that copy into an error under this option too. A constant that is changed on the host before a delegate reads it, such as by a view, a transpose or a clone, is not read by the delegate directly, so it is still copied; the docs say so. They also say that share_mutable_buffers, which needs every planned tensor in host memory, does not work with this option. The CUDA guide now says that every skip flag needs enable_non_cpu_memory_planning, and that skip_h2d_for_constants works on its own, without skip_h2d_for_method_inputs or unplanned inputs. The new Python test checks that a buffer or parameter read only by delegates on one device has no copy and no warning, has planned device memory for the whole method with its data as the initial state, with and without planned graph inputs, and that a constant also read on the host or on a second device, or any constant without the option, is still copied. It fails without this change. export_program_with_device_info.py now also exports such a program, on cuda:1, once with the constant in the program and once in a separate .ptd file. Two new C++ tests parse the constant of each into planned memory on a mock CUDA device and check that its data was copied there, to device 1, and that the copy was waited for. tensor_parser_device_test now also builds and runs with CMake. A new CUDA test runs a program whose constant is kept in device memory twice, with different inputs, and compares with eager PyTorch, with error_on_per_call_copies set so it fails if the constant is copied.
The warning for a constant copied to the device on every call only said to move it into the delegate. Without skip_h2d_for_constants, it now also suggests the option when only delegates on that device read the constant, unless the constant is empty and costs nothing to copy. With the option, it names the other reader that keeps the constant off the device.
be1bb26 to
0cf1823
Compare
|
@pytorchbot cherry-pick --onto release/1.6 -c fixnewfeature |
Cherry picking #24144Command Details for Dev Infra teamRaised by workflow job |
|
@pytorchbot cherry-pick --onto release/1.6 -c fixnewfeature |
Cherry picking #24144The cherry pick PR is at #24274 and it is recommended to link a fixnewfeature cherry pick PR with an issue. Details for Dev Infra teamRaised by workflow job |
What is wrong today
With
skip_h2d_for_method_inputs, the inputs of a device delegate arrive in device memory. Constants (weights and read-only buffers) do not. Their data is stored on the host, so a constant that a delegate reads from outside is copied to the device on every call. The CUDA partitioner already moves constants into its delegate, so this hits constants left outside, for example with a custom partitioner.What this change does
A new opt-in option keeps such a constant in device memory. The runtime copies it there once, when the method loads.
PropagateDevicePasstags a constant that only delegates on one device read with that device.DeviceAllocator. NewDeviceAllocator::synchronize()waits for the copy, so the host buffer can be freed and the first run reads the right data. On CUDA it waits for the stream the copy used, on the right device.Limits:
share_mutable_buffersdoes not work with it.Cost: the loading code is in every runtime build. In CI, the stripped Arm bare-metal
size_testgrows by 560 bytes, to 117884 (limit 118000). On x86 its code grows by 727 bytes with GCC and 607 with Clang, and the stripped file size does not change.What was tested
test_propagate_device_pass.py: only such constants skip the copy and get device memory; warnings and errors.tensor_parser_device_test.cpp: such a constant loads from the program and from a.ptdfile into device 1 of a mock CUDA device, and is waited for.test_emit.py: a method whose initial data is empty loads.test_cuda_allocator.cpp:synchronize()waits for the right stream, also with another GPU current (needs two GPUs; ran on H100s).test_device_placement.py: such a program runs twice on an H100 and matches eager PyTorch.Each new test fails without its change, and each commit passes on its own.