-
Notifications
You must be signed in to change notification settings - Fork 362
feat: Add support for device compilation setting #2190
New issue
Have a question about this project? # for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “#”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? # to your account
feat: Add support for device compilation setting #2190
Conversation
fa81096
to
9e1e58e
Compare
9e1e58e
to
8e32c07
Compare
8e32c07
to
776839a
Compare
776839a
to
89764ce
Compare
- Add updated Device utilities and automatic context-aware device detection for torch compile - Add testing for new utilities
89764ce
to
ba18185
Compare
@@ -54,3 +56,4 @@ class CompilationSettings: | |||
truncate_long_and_double: bool = TRUNCATE_LONG_AND_DOUBLE | |||
use_fast_partitioner: bool = USE_FAST_PARTITIONER | |||
enable_experimental_decompositions: bool = ENABLE_EXPERIMENTAL_DECOMPOSITIONS | |||
device: Device = field(default_factory=default_device) |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Device is now populated with a default factory once the CompilationSettings
object gets instantiated.
@@ -38,7 +43,7 @@ def compile( | |||
gm: Any, | |||
inputs: Any, | |||
*, | |||
device: Device = Device._current_device(), | |||
device: Optional[Union[Device, torch.device, str]] = DEVICE, |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Device can now be None
, which will take the Torch default device.
elif device is None: | ||
return Device(gpu_id=torch.cuda.current_device()) |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
If the device is None
, we use the default Torch context device
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
LGTM
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
LGTM
Description
Fixes #2172
Type of change
Checklist: