Open texasmichelle opened 4 years ago
As examples of how these mappings are defined at the command line, here's how you would expose both the CPU and GPU as selectable devices (assuming a single CPU and GPU):
export XRT_DEVICE_MAP='CPU:0;/job:localservice/replica:0/task:0/device:XLA_CPU:0|GPU:0;/job:localservice/replica:0/task:0/device:XLA_GPU:0'
and here's how you would expose two GPUs (not exposing the CPU):
export XRT_DEVICE_MAP='GPU:0;/job:localservice/replica:0/task:0/device:XLA_GPU:0|GPU:1;/job:localservice/replica:0/task:0/device:XLA_GPU:1'
Currently, only one default device is found and exposed. If you want something other than the default, you need to manually specify the XLA -> S4TF mapping for all devices you want. The devices are parsed from the XRT_DEVICE_MAP environment variable within ParseEnvDevices here. That may be the place to add CPU support on GPU-default systems, because we can safely assume the CPU is present there.
On a machine with GPU or TPU, I get a segfault if I try to use
Device
with CPU type on XLA backend, e.g.:A workaround is to set the
XRT_DEVICE_MAP
environment variable, but all device and backend combinations should be accessible without this.See swift-models/#654.