Add Ascend NPU Support #1521
Add Ascend NPU Support #1521
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/torchchat/1521
Note: Links to docs will display an error until the docs builds have been completed. This comment was automatically generated by Dr. CI and updates every 15 minutes. |
|
Hi @xuedinge233! Thank you for your pull request and welcome to our community. Action RequiredIn order to merge any pull request (code, docs, etc.), we require contributors to sign our Contributor License Agreement, and we don't seem to have one on file for you. ProcessIn order for us to review and merge your suggested changes, please sign at https://code.facebook.com/cla. If you are contributing on behalf of someone else (eg your employer), the individual CLA may not be sufficient and your employer may need to sign the corporate CLA. Once the CLA is signed, our tooling will perform checks and validations. Afterwards, the pull request will be tagged with If you have received this in error or have any questions, please contact us at cla@meta.com. Thanks! |
facebook-github-bot
added
the
CLA Signed
|
Thank you for signing our Contributor License Agreement. We can now accept your code for this (and any) Meta Open Source project. Thanks! |
xuedinge233
changed the title
torchchat/generate.py
Outdated
| @@ -1595,6 +1599,12 @@ def sample( | |||
|
|
|||
| return idx_next, probs | |||
|
|
|||
| def is_npu_available(): | |||
There was a problem hiding this comment.
Out of tree device support autoload, don't need import torch_npu.
torchchat/utils/build_utils.py
Outdated
| @@ -275,33 +277,51 @@ def is_mps_available() -> bool: | |||
| # MPS, is that you? | |||
| return True | |||
|
|
|||
| def is_npu_available(check_device=False): | |||
There was a problem hiding this comment.
Out of tree device support autoload, don't need import torch_npu.
xuedinge233
changed the title
torchchat/utils/build_utils.py
Outdated
| return "cuda" | ||
| elif is_mps_available(): | ||
| return "mps" | ||
| elif hasattr(torch, "npu") and torch.npu.is_available(): |
| @@ -1213,6 +1213,8 @@ def callback(x, *, done_generating=False): | |||
| print(prof.key_averages().table(sort_by="self_cpu_time_total")) | |||
| elif self.builder_args.device == "cuda": | |||
| print(prof.key_averages().table(sort_by="self_cuda_time_total")) | |||
| elif self.builder_args.device == "npu": | |||
| print(prof.key_averages().table(sort_by="self_npu_time_total")) | |||
There was a problem hiding this comment.
where is self_npu_time_total defined.
|
@Jack-Khuu @mikekgfb Hello，can you help review this PR? |
| @@ -71,6 +71,9 @@ then | |||
| elif [[ -x "$(command -v xpu-smi)" ]]; | |||
| then | |||
| TORCH_NIGHTLY_URL="https://download.pytorch.org/whl/nightly/xpu" | |||
| elif [[ -x "$(command -v npu-smi)" ]] | |||
| then | |||
| TORCH_NIGHTLY_URL="https://download.pytorch.org/whl/test/cpu" | |||
There was a problem hiding this comment.
Why are we using a test wheel?
There was a problem hiding this comment.
Hey @Jack-Khuu thanks for your review. IMO, we should use nightly pytorch wheels like other backends. What do you think? @hipudding @xuedinge233
There was a problem hiding this comment.
But, we usually use pytorch RC versions, not sure if everything works fine on nightly
torchchat/cli/builder.py
Outdated
| self.device = "cuda" | ||
| elif torch.xpu.is_available(): | ||
| self.device = "xpu" | ||
| elif hasattr(torch, "npu") and torch.npu.is_available(): | ||
| self.device = "npu" | ||
| else: | ||
| self.device = "cpu" |
There was a problem hiding this comment.
How about torch.accelerator.current_accelerator(check_available=True)? We've done a lot of work on device-generic APIs. Hope they get used.
See: https://dev-discuss.pytorch.org/t/python-c-api-rules-for-device-generic-apis/2511
There was a problem hiding this comment.
Got it，I will make modifications
This PR introduces native support for Huawei Ascend NPU acceleration in generative AI tasks and chat applications. The implementation leverages Ascend's CANN computing architecture to optimize tensor operations for large language models.
Compatibility Matrix:
Python 3.10.15
torch 2.7.0+cpu
torch_npu 2.7.0+git36b5d8c (nightly)
torchvision 0.22.0
torchtune 0.6.0
The following is the result generated using Llama3.1.
