Skip to content

Set pytorch cuda arch list explicitly to get rid of warning. #44

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
hzhangxyz merged 1 commit into from
Jun 26, 2025
Merged

Set pytorch cuda arch list explicitly to get rid of warning. #44

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
19 changes: 19 additions & 0 deletions qmb/hamiltonian.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,8 +16,27 @@ class Hamiltonian:

_hamiltonian_module: dict[tuple[str, int, int], object] = {}

@classmethod
def _set_torch_cuda_arch_list(cls) -> None:
"""
Set the CUDA architecture list for PyTorch to use when compiling the PyTorch extensions.
"""
if not torch.cuda.is_available():
return
if "TORCH_CUDA_ARCH_LIST" in os.environ:
return
supported_sm = [int(arch[3:]) for arch in torch.cuda.get_arch_list() if arch.startswith("sm_")]

Copilot AI Jun 26, 2025

Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If supported_sm is empty, calling max() will raise a ValueError. Consider guarding against an empty list (e.g. return early) to avoid a crash when no arch entries are found.

Suggested change
supported_sm = [int(arch[3:]) for arch in torch.cuda.get_arch_list() if arch.startswith("sm_")]
supported_sm = [int(arch[3:]) for arch in torch.cuda.get_arch_list() if arch.startswith("sm_")]
if not supported_sm:
return # No supported architectures, exit early

Copilot uses AI. Check for mistakes.

@hzhangxyz hzhangxyz Jun 26, 2025

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

torch.cuda.is_available() 之前判断过，不会空

max_supported_sm = max((sm // 10, sm % 10) for sm in supported_sm)
arch_list = set()
for i in range(torch.cuda.device_count()):
capability = min(max_supported_sm, torch.cuda.get_device_capability(i))
arch = f'{capability[0]}.{capability[1]}'
arch_list.add(arch)
os.environ["TORCH_CUDA_ARCH_LIST"] = ";".join(sorted(arch_list))

Copilot AI Jun 26, 2025

Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Sorting the version strings lexicographically can misorder multi-digit majors (e.g. "10.2" vs "9.1"). Use a numeric key (e.g. sorted(arch_list, key=lambda v: tuple(map(int, v.split('.'))))) to ensure correct ascending order.

Suggested change
os.environ["TORCH_CUDA_ARCH_LIST"] = ";".join(sorted(arch_list))
os.environ["TORCH_CUDA_ARCH_LIST"] = ";".join(sorted(arch_list, key=lambda v: tuple(map(int, v.split('.')))))

Copilot uses AI. Check for mistakes.

@msg-bq msg-bq Jun 26, 2025

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

这里copilot给的提示是对的吧，不像上一个没必要。师兄的版本号是str，字典序不妥，还得是像(sm // 10, sm % 10)那种拆开的叭。虽然排序无妨

@hzhangxyz hzhangxyz Jun 26, 2025

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

确实，这里排序是为了每次运行他们的顺序是一致的，不然无法有效利用缓存。

另外，我这段代码抄的:

https://github.com/pytorch/pytorch/blob/0d01bafc34fc99a0b3e46cbf1ecfd8f97563efa6/torch/utils/cpp_extension.py#L2428-L2444


@classmethod
def _load_module(cls, device_type: str = "declaration", n_qubytes: int = 0, particle_cut: int = 0) -> object:
cls._set_torch_cuda_arch_list()
if device_type != "declaration":
cls._load_module("declaration", n_qubytes, particle_cut) # Ensure the declaration module is loaded first
key = (device_type, n_qubytes, particle_cut)
Expand Down
Loading