Fix torch_memory_saver LD_PRELOAD path for CUDA-suffixed binaries - #1785
Merged
Conversation
Contributor
|
Caution The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased. |
torch_memory_saver 0.0.9.post1 ships CUDA-major-suffixed preload binaries (torch_memory_saver_hook_mode_preload_cu13.abi3.so), but actor_factory hardcoded the old unsuffixed name and asserted it exists, breaking every offload_train megatron run. Use the package's get_binary_path_from_package helper to resolve the right variant.
yueming-yuan
force-pushed
the
fix-tms-preload-path
branch
from
July 24, 2026 06:39
94aaa29 to
acf71ab
Compare
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
2 participants
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
Since torch_memory_saver was upgraded to 0.0.9.post1 (#1773, #1774), the package ships CUDA-major-suffixed preload binaries (e.g.
torch_memory_saver_hook_mode_preload_cu13.abi3.so) instead of the old unsuffixedtorch_memory_saver_hook_mode_preload.abi3.so.miles/ray/train/actor_factory.pystill hardcodes the old unsuffixed filename and asserts it exists, so everyoffload_trainmegatron run on the current:devimage dies at startup:This is currently red on main's own offload_train megatron path and blocks unrelated PRs (surfaced while monitoring #1261's stage-c-8-gpu CI).
Fix
Use the package-provided
get_binary_path_from_package("torch_memory_saver_hook_mode_preload")helper, which resolves the correct CUDA-suffixed variant (and raises a clear error itself if none is found, so the hand-written assert is no longer needed).Validation
On a
radixark/miles:devB300 devbox (torch_memory_saver 0.0.9.post1):pre-commit passes.