-
Notifications
You must be signed in to change notification settings - Fork 364
feat: Add support for device compilation setting #2190
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in 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: