exo icon indicating copy to clipboard operation
exo copied to clipboard

Add `--gpu` CLI option so Linux users can choose a single GPU ID for tinygrad

Open freerainboxbox opened this issue 3 weeks ago • 0 comments

Original Behavior

  1. User has to set VISIBLE_DEVICES in shell for tinygrad to accept choice of GPUs.
  2. Even if set to something other than 0, linux_device_capabilities() will fetch only index 0 (this was hardcoded in).

New Behavior

  1. The user may specify --gpu followed by an integer. If parsed incorrectly or omitted, it will default to 0. 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.*
  2. 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.

freerainboxbox avatar Jan 30 '25 04:01 freerainboxbox