With accelerate and colab tpu all devices always xla:0 and none of them is_main_process

Here is code

So, seems like a bug, but I’m not sure

I’m not sure where in the notebook you see that. Could you post a minimal reproducer? Printing the accelerator.device in the training function shows 8 different devices.

Just launched notebook, I’m getting such output from my prints

    print(accelerator.is_main_process)
    print(accelerator.device)
False
xla:0
xla:0
False
xla:0
False
xla:0
False
xla:0
False
xla:0
False
xla:0

So, all is not main process and all xla 0 while I expect xla 0-7 and one main process.

Also, left outputs in notebook, so you can see output from train loop.

You are right, it looks like it associates the device 0 to all but one process that gets the device 1 (on my example). Not sure what’s going on, I’m investigating.