frozenleaves/pytorch
1
0
Fork 0
mirror of https://github.com/pytorch/pytorch.git synced 2025-10-20 21:14:14 +08:00
Code Issues Packages Projects Releases Wiki Activity

[Easy] Add ROCm support to nightly pull tool (#141282)

Browse Source 
Pull Request resolved: https://github.com/pytorch/pytorch/pull/141282
Approved by: https://github.com/malfet
ghstack dependencies: #143263
This commit is contained in:
Xuehai Pan
2024-12-26 16:48:42 +08:00
committed by PyTorch MergeBot
parent 8059d56ec3
commit c4bff71854
4 changed files with 64 additions and 20 deletions
Download Patch File Download Diff File Expand all files Collapse all files

68
tools/nightly.py
View File

@ -20,6 +20,11 @@ To install the nightly binaries built with CUDA, you can pass in the flag --cuda
$ ./tools/nightly.py checkout -b my-nightly-branch --cuda
$ source venv/bin/activate # or `& .\venv\Scripts\Activate.ps1` on Windows
To install the nightly binaries built with ROCm, you can pass in the flag --rocm::
$ ./tools/nightly.py checkout -b my-nightly-branch --rocm
$ source venv/bin/activate # or `& .\venv\Scripts\Activate.ps1` on Windows
You can also use this tool to pull the nightly commits into the current branch as
well. This can be done with::
@ -134,6 +139,12 @@ PIP_SOURCES = {
supported_platforms={"Linux", "Windows"},
accelerator="cuda",
),
"rocm-6.2.4": PipSource(
name="rocm-6.2.4",
index_url=f"{PYTORCH_NIGHTLY_PIP_INDEX_URL}/rocm6.2.4",
supported_platforms={"Linux"},
accelerator="rocm",
),
}
@ -882,6 +893,17 @@ def make_parser() -> argparse.ArgumentParser:
default=argparse.SUPPRESS,
metavar="VERSION",
)
subparser.add_argument(
"--rocm",
help=(
"ROCm version to install "
"(defaults to the latest version available on the platform)"
),
dest="rocm",
nargs="?",
default=argparse.SUPPRESS,
metavar="VERSION",
)
return parser
@ -889,6 +911,8 @@ def parse_arguments() -> argparse.Namespace:
parser = make_parser()
args = parser.parse_args()
args.branch = getattr(args, "branch", None)
if hasattr(args, "cuda") and hasattr(args, "rocm"):
parser.error("Cannot specify both CUDA and ROCm versions.")
return args
@ -901,26 +925,32 @@ def main() -> None:
sys.exit(status)
pip_source = None
if hasattr(args, "cuda"):
available_sources = {
src.name[len("cuda-") :]: src
for src in PIP_SOURCES.values()
if src.name.startswith("cuda-") and PLATFORM in src.supported_platforms
}
if not available_sources:
print(f"No CUDA versions available on platform {PLATFORM}.")
sys.exit(1)
if args.cuda is not None:
pip_source = available_sources.get(args.cuda)
if pip_source is None:
print(
f"CUDA {args.cuda} is not available on platform {PLATFORM}. "
f"Available version(s): {', '.join(sorted(available_sources, key=Version))}"
)
for toolkit in ("CUDA", "ROCm"):
accel = toolkit.lower()
if hasattr(args, accel):
requested = getattr(args, accel)
available_sources = {
src.name[len(f"{accel}-") :]: src
for src in PIP_SOURCES.values()
if src.name.startswith(f"{accel}-")
and PLATFORM in src.supported_platforms
}
if not available_sources:
print(f"No {toolkit} versions available on platform {PLATFORM}.")
sys.exit(1)
else:
pip_source = available_sources[max(available_sources, key=Version)]
else:
if requested is not None:
pip_source = available_sources.get(requested)
if pip_source is None:
print(
f"{toolkit} {requested} is not available on platform {PLATFORM}. "
f"Available version(s): {', '.join(sorted(available_sources, key=Version))}"
)
sys.exit(1)
else:
pip_source = available_sources[max(available_sources, key=Version)]
if pip_source is None:
pip_source = PIP_SOURCES["cpu"] # always available
with logging_manager(debug=args.verbose) as logger: