fix(rerank): pin torch to cu128 wheels for wider CUDA compatibility#152
Closed
shan-nvidia wants to merge 1 commit intooholworthy/rerank-recipe-v1from
Closed
fix(rerank): pin torch to cu128 wheels for wider CUDA compatibility#152shan-nvidia wants to merge 1 commit intooholworthy/rerank-recipe-v1from
shan-nvidia wants to merge 1 commit intooholworthy/rerank-recipe-v1from
Conversation
Signed-off-by: Steve Han <sthan@nvidia.com>
Contributor
Author
|
Close this since it's already fixed. |
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
The rerank eval stage (
stage3_eval) resolved PyTorch from PyPI, whichpulls wheels compiled against the latest CUDA toolkit (currently 13.0).
This breaks on any server whose NVIDIA driver does not yet support
CUDA 13.0, for example:
Fix
Pin torch to the cu128 (CUDA 12.8) wheel index, which is
forward-compatible with CUDA 12.8, 12.9, 13.0, and newer drivers.
torchas a direct dependency (previously only a transitive depvia
sentence-transformers).[[tool.uv.index]]pointing tohttps://download.pytorch.org/whl/cu128withexplicit = true, soonly
torchresolves from it — all other packages continue to usePyPI.
[tool.uv.sources]mappingtorchto the cu128 index.uv.lock.Files changed
src/nemotron/recipes/rerank/stage3_eval/pyproject.tomlsrc/nemotron/recipes/rerank/stage3_eval/uv.lockTest plan
nemotron rerank evalon a CUDA 12.x driver — verify torch initializes CUDA successfully