[Feature] Add NPU Support for Single Agent #3229
Open
+59
−4
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.
Description
Adds support for Ascend NPUs in most single-agent sota-implementations. These changes are non-intrusive to existing code and have been thoroughly tested on Ascend NPU devices.
As the next step, we are working to extend TorchRL NPU support to multi-agent experiments and GRPO, and we plan to submit a separate PR for this work soon.
Motivation and Context
We aim to extend TorchRL to enable execution on NPUs, providing broader hardware compatibility.
And our issue is #3154
Types of changes
What types of changes does your code introduce? Remove all that do not apply:
Checklist
Go over all the following points, and put an
xin all the boxes that apply.If you are unsure about any of these, don't hesitate to ask. We are here to help!