exo
exo copied to clipboard
Add `--gpu` CLI option so Linux users can choose a single GPU ID for tinygrad
Original Behavior
- User has to set
VISIBLE_DEVICES
in shell for tinygrad to accept choice of GPUs. - Even if set to something other than
0
,linux_device_capabilities()
will fetch only index 0 (this was hardcoded in).
New Behavior
- The user may specify
--gpu
followed by an integer. If parsed incorrectly or omitted, it will default to0
.VISIBLE_DEVICES
in the parent environment is effectively abstracted away from the user, so if they want to pick a device, they can pick only one for the instance.* - Device capabilities will fetched according to this index, and the user will see the correct GPU as confirmation in the TUI (see behavior of
topology.py
).
*Users can run multiple GPUs at once by running multiple instances of exo. This workaround I wrote may be obsolete if we implement multi-GPU in a single instance of exo, but for now, this will allow users to take advantage of multiple GPUs in one device.