feat: add device id #779
feat: add device id #779
Conversation
|
Why not use the |
|
@OlivierDehaene good to see you. I actually want to add py03 bindings for the backend, and wiring that though would help. Env vars can't change, so its more about adding this capability into backend. |
There was a problem hiding this comment.
Thank you @michaelfeil, I've left some comments!
P.S. Agree with @OlivierDehaene that using CUDA_VISIBLE_DEVICES would be ideal, but given that this might help with your use-case it's probably the same with other users, so I'm happy to move forward, unless Olivier thinks otherwise, I don't have a strong opinion against this change 🤗
| #[clap(long, env)] | ||
| dense_path: Option<String>, | ||
|
|
||
| /// The device ID to use for CUDA/Metal devices. Defaults to 0. |
There was a problem hiding this comment.
| /// The device ID to use for CUDA/Metal devices. Defaults to 0. | |
| /// The CUDA device ID where the model will be loaded. Defaults to 0 i.e., the first available device. |
| Ok(Device::Cpu) | ||
| } else if candle::utils::metal_is_available() { | ||
| Device::new_metal(0) | ||
| Device::new_metal(device_id) |
There was a problem hiding this comment.
AFAIK there are no instances where you have more than one M-chip on MacOS, right?
| Device::new_metal(device_id) | |
| Device::new_metal(0) |
| #[cfg(feature = "cuda")] | ||
| match compatible_compute_cap() { | ||
| Ok(true) => Device::new_cuda(0), | ||
| Ok(true) => Device::new_cuda(device_id), |
There was a problem hiding this comment.
IMO we should try to add some sort of validation here, to ensure that the device_id is within the bounds of the available devices in the given instance, or is the candle error enough?
3 participants
What does this PR do?
This adds a option to the CLI to select the device ID.
Fixes # (issue)
Before submitting
instasnapshots?Who can review?
Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.