[Fix] Fix compile error by HAOCHENYE 路 Pull Request #4357 路 hpcaitech/ColossalAI (original) (raw)

馃搶 Checklist before creating the PR

馃毃 Issue number

Link this PR to your issue with words like fixed to automatically close the linked issue upon merge

e.g. fixed #1234, closed #1234, resolved #1234

馃摑 What does this PR do?

torch.cuda.get_arch_list() here:

for arch in torch.cuda.get_arch_list():

will get the available GPU arch for the current PyTorch lib. For example, the PyTorch 2.0 is built by the hopper arch GPU, it will return a list:

['sm_37', 'sm_50', 'sm_60', 'sm_70', 'sm_75', 'sm_80', 'sm_86', 'sm_90']

However, my GPU is an ampere arch GPU, return a list includes sm_90 here will cause an compile error like:

nvcc fatal : Unsupported gpu architecture 'compute_90'.

Therefore, I think the valid arch returned here should be >= sm60 but <= sm86, which is the max available arch for my GPU

馃挜 Checklist before requesting a review

猸愶笍 Do you enjoy contributing to Colossal-AI?

Tell us more if you don't enjoy contributing to Colossal-AI.