Experimental DirectML support via torch-directml #1702
Draft
+35
−6
Conversation
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This commit introduces two improvements:
1. DirectML acceleration:
- Added support for running optimum commands on DirectML hardware (Windows only) using the --device dml flag.
- Automatically sets the device to torch_directml.device() when the flag is specified.
2. Improved device handling:
- Ensures the model is directly initialized in the device only when applicable.
This commit refines the device handling in optimum/exporters/tasks.py for the following improvements: - More precise device check: Instead of checking for not device.type, the condition is updated to device.type != "privateuseone". This ensures the initialization happens on the requested device only if it's not a private use device (e.g., DirectML). - Improved clarity: The code comments are updated to better explain the purpose of the device initialization and its benefits for large models.
- Extends device compatibility to "privateuseone" in export_pytorch for exporting models usable on specific hardware. This commit allows exporting PyTorch models compatible with the "privateuseone" device, potentially enabling inference on specialized hardware platforms.
This commit adds support for running PyTorch models on the DML device within the Optimum framework. - Dynamic DML device handling: Introduces dynamic import of torch_directml for improved maintainability. - Consistent device selection: Ensures consistent device selection across optimum/exporters/onnx/convert.py, optimum/exporters/tasks.py, and optimum/onnxruntime/io_binding/io_binding_helper.py. This change allows users to leverage DML capabilities for efficient PyTorch model inference with Optimum.
This commit removes unnecessary code for handling the DML device in optimum/commands/optimum_cli.py. - Redundant import: The code previously imported torch_directml conditionally, which is no longer needed as DML device support is handled in other parts of the codebase. This change simplifies the code and avoids potential conflicts.
This commit updates `setup.py` to include the following changes: - Introduces a new conditional section "exporters-directml" with dependencies required for exporting models for DML inference. - This section mirrors the existing "exporters" and "exporters-gpu" sections, adding `onnxruntime-directml` as a dependency. This update ensures users have the necessary libraries for working with DML devices when installing Optimum with DML support.
…into feature/directml
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
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.
What does this PR do?
This PR adds experimental DirectML support via torch-directml, which is still in preview and lacks several PyTorch functions such as microsoft/DirectML#449.
If you are interested in this PR, please leave comments below.
Before submitting