Add opt-in sharing of identical weight literals across programs on a device#5039
Add opt-in sharing of identical weight literals across programs on a device#5039aditya-dl wants to merge 3 commits into
Conversation
Add a process- and device-scoped content-addressed pool so multiple compiled programs (e.g. an LLM prefill and decode program) that hold byte-identical weights share a single VRAM copy. A hash selects a candidate and the bytes are always compared before aliasing, so a hash collision never causes an incorrect share; on miss or mismatch a private copy is uploaded. Gated by MIGRAPHX_SHARE_LITERALS (default off).
|
Thank you for your contribution! Since this is an external pull request, a maintainer must review PR and add the "ok-to-test" label if it is approved for testing. |
TedThemistokleous
left a comment
TedThemistokleous
left a comment
There was a problem hiding this comment.
One minor comment on the loop - use ranges we have them in MIGraphX
Another on hash collsion case
You need to add tests for this change. Generating two programs and toggling the flag on in tests. We should have the following cases. It'll give us a better idea on how this would work.
- Case where program A shares its weights with B. (flag on)
- Case where program A and B can share weights but don't because flag is off
- Case of the above with more than two programs (A,B,C,D)
| const auto* p = reinterpret_cast<const unsigned char*>(data.data()); | ||
| const std::size_t n = data.get_shape().bytes(); | ||
| std::uint64_t h = 1469598103934665603ULL; // FNV offset basis | ||
| for(std::size_t i = 0; i < n; ++i) |
There was a problem hiding this comment.
You can use migraphx ranges here with transform to avoid the raw loop
| } | ||
| // Miss, or a hash collision with different bytes: upload a private copy. | ||
| // On a genuine miss, register it so later identical weights can share. | ||
| gpu_data = to_gpu(data); |
There was a problem hiding this comment.
Just to be clear, hash collision should always upload then, even if the data is already on the GPU?
CharlieL7
left a comment
CharlieL7
left a comment
There was a problem hiding this comment.
In the future, please also make the PR motivation and technical detail sections more terse. It looks like it was LLM generated with the usage of hypens to split sentences.
Do you have an example script that uses this functionality?
|
|
||
| ### Added | ||
|
|
||
| * Added environment variable `MIGRAPHX_SHARE_LITERALS` which, when enabled, shares a single device copy of byte-identical weight literals across programs compiled in the same process (#PR_NUMBER). |
| // Process-lifetime, per-(device,content-hash) pool of already-uploaded weights. | ||
| // Each entry retains the host literal (for the collision byte-compare) and the | ||
| // device buffer; gpu_literals that opt in share the SAME device `argument` | ||
| // (shared_ptr-refcounted), so N identical weights => 1 VRAM copy. Refs are held | ||
| // for the process lifetime, which suits the co-residency use case (the weights | ||
| // stay live as long as the resident programs); a process that repeatedly loads | ||
| // and unloads distinct models would retain their weights until exit. Guarded by | ||
| // a mutex since finalize may run on multiple compile threads. |
There was a problem hiding this comment.
Make this comment terser
| // Shares a single VRAM copy of a weight literal across multiple compiled | ||
| // programs on the same device when their weight bytes are identical (for | ||
| // example, an LLM's prefill and decode programs). Default off: each program | ||
| // uploads its own copy as before. |
There was a problem hiding this comment.
Make this comment terser
| std::uint64_t literal_content_hash(const argument& data) | ||
| { | ||
| const auto* p = reinterpret_cast<const unsigned char*>(data.data()); | ||
| const std::size_t n = data.get_shape().bytes(); | ||
| std::uint64_t h = 1469598103934665603ULL; // FNV offset basis | ||
| for(std::size_t i = 0; i < n; ++i) | ||
| { | ||
| h ^= p[i]; | ||
| h *= 1099511628211ULL; // FNV prime | ||
| } | ||
| return h; | ||
| } |
There was a problem hiding this comment.
I'm pretty sure you could use migraphx::to_value's hash function instead of rolling another hash function here on the argument's data buffer.
| // Exact equality of two host literals (same byte length and identical bytes). | ||
| // Guards against FNV-1a hash collisions: two distinct weights that hash equal | ||
| // must not be aliased to the same device buffer. | ||
| bool same_bytes(const argument& a, const argument& b) | ||
| { | ||
| const std::size_t n = a.get_shape().bytes(); | ||
| if(n != b.get_shape().bytes()) | ||
| return false; | ||
| return std::memcmp(a.data(), b.data(), n) == 0; | ||
| } |
There was a problem hiding this comment.
I don't see why this would need std::memcmp over using operator== of argument from raw_data. Unless the shapes of the arguments are different but it's the same data.
| // FNV-1a 64-bit over the literal's raw host bytes, used as the content key. The | ||
| // hash only selects a candidate; the bytes are always compared before sharing | ||
| // (see same_bytes), so a hash collision never causes an incorrect share. | ||
| std::uint64_t literal_content_hash(const argument& data) |
There was a problem hiding this comment.
This should use md5 hash function as it can avoid more collisions and is twice the width.
| // Exact equality of two host literals (same byte length and identical bytes). | ||
| // Guards against FNV-1a hash collisions: two distinct weights that hash equal | ||
| // must not be aliased to the same device buffer. | ||
| bool same_bytes(const argument& a, const argument& b) |
| { | ||
| // Namespace the pool per-device via the public gfx-name (device_id is | ||
| // private on hip_device). Content-hash keys identical bytes together. | ||
| const std::string key = ctx.get_current_device().get_gfx_name() + ":" + |
There was a problem hiding this comment.
You cant key on the gfx name. Usually multiple gpus are the same gfx.
| auto& pool = shared_literal_pool::instance(); | ||
| std::lock_guard<std::mutex> lock(pool.mtx); | ||
| auto it = pool.buffers.find(key); | ||
| if(it != pool.buffers.end() and same_bytes(it->second.host, data)) |
There was a problem hiding this comment.
If you are going to check for equality then you should just make the argument a key as hash collisions will cause data to not be stored.
Although I think it would be better to just store the md5 sum and the tensor shape as the key and then we dont need to store the host data and compare as hash collisions should not really happen(it is sufficient for a binary cache so it should work here as well?)/
| } | ||
| // Miss, or a hash collision with different bytes: upload a private copy. | ||
| // On a genuine miss, register it so later identical weights can share. | ||
| gpu_data = to_gpu(data); |
There was a problem hiding this comment.
It doesnt seem good to lock while copying to the gpu. I dont know if there is a better way.
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## develop #5039 +/- ##
===========================================
+ Coverage 92.89% 92.90% +0.01%
===========================================
Files 603 603
Lines 32448 32526 +78
===========================================
+ Hits 30140 30217 +77
- Misses 2308 2309 +1 🚀 New features to boost your workflow:
|
4 participants
Motivation
When several programs are compiled and kept resident in the same process — for example the prefill and decode programs of a generative LLM, which differ only in input shape — each program uploads its own copy of the model's weight literals to device memory. The weights are byte-identical between those programs, so this is a redundant second (or Nth) VRAM copy of the largest tensors in the model. For large models that duplication can be the difference between fitting in device memory and not.
This change lets programs that hold identical weights share a single device copy, removing the duplication, with no change to numerical results.
See this for the companion change: onnxruntime/onnxruntime-ep-amdgpu#37
Technical Details
All changes are in
src/targets/gpu/write_literals.cpp(the GPU target), gated behind a new environment variableMIGRAPHX_SHARE_LITERALS(default off).get_gfx_name()) plus an FNV-1a 64-bit hash of the literal's host bytes.gpu_literal::finalize, when sharing is enabled and the literal is device-resident (not host-pinned), the pool is consulted. The hash only selects a candidate; the full bytes are then compared (same_bytes) before any sharing happens, so a hash collision can never alias two genuinely different weights. On a confirmed match the existing device buffer is aliased (argument::share(), a refcount bump — no new upload). On a miss, or a hash collision with differing bytes, a private copy is uploaded as before.std::mutexsince finalization can run on multiple compile threads.MIGRAPHX_SHARE_LITERALSis unset, execution falls through to the originalregister_on_gpu/to_gpupath unchanged.Scope note: the pool holds references for the lifetime of the process, which suits the resident-program use case (the weights stay live as long as the programs that use them). A process that repeatedly loads and unloads many distinct models would retain their weights until exit; eviction can be added later if that pattern needs it (dropping a pooled reference is safe — any program still using a buffer keeps it alive via the shared handle).
Changelog Category
Add a
CHANGELOG.mdentry for any option other thanNot Applicable