fix(graph): deduplicate shared mmap weight files across component graphs - #794
Open
Moeez34 wants to merge 1 commit into
Open
fix(graph): deduplicate shared mmap weight files across component graphs#794Moeez34 wants to merge 1 commit into
Moeez34 wants to merge 1 commit into
Conversation
Signed-off-by: Moeez Ahmed <ahmedmoeez7860@gmail.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Problem
When loading models composed of multiple component graphs (e.g.
lm_encoder,lm_encoder_text_chunk,lm_encoder_stepfor models likegemma-4-e2b-it), shared weight files likeembed_tokens_per_layer.weightsare bound by multiple components.Previously,
weight_cache_was isolated to each individualCactusGraphinstance. As a result, each component graph created its ownMappedFileobject and issued an independentmmap()system call for shared weights.On operating systems with per-process virtual address space / mapped-memory ceilings (such as iOS's ~6.5 GiB limit), repeatedly mapping a 1.13 GiB weight file 3+ times alongside other weights exhausts the process address space budget. This caused
mmapto fail withENOMEM("Cannot map file") duringcactus_init.Solution
Process-Wide Thread-Safe Mmap Cache (
MappedFileRegistry):GraphFile::MappedFileRegistryincactus_graph.h/io.cpp.std::filesystem::canonical) tostd::weak_ptr<GraphFile::MappedFile>.MappedFileinstance instead of creating duplicate virtual memory mappings.Refactored
CactusGraphStorage:CactusGraph::mapped_files_fromstd::vector<std::unique_ptr<MappedFile>>tostd::vector<std::shared_ptr<MappedFile>>.mmap_embeddings,mmap_weights,bind_mmap_weights, andembeddingto retrieve sharedMappedFilepointers viaMappedFileRegistry::get_or_load().Automatic Cleanup (RAII):
std::shared_ptrreference count reaches zero, triggeringMappedFile::~MappedFile()(munmap), cleanly unmapping the region.Verification
test_shared_mmap_weights_dedup()unit test incactus-graph/tests/test_io.cpp.CactusGraphinstances binding the same weight file receive identical mmap buffer pointers (ptr1 == ptr2) and release the mapping safely upon graph teardown.