[TPU] Show / configure the number of TPU chips similar to how the number of GPUs are configured

[peterschmidt85]

Problem

Currently, if I want to specify TPU resources, I cannot specify the number of GPUs but I have to specify the number of cores.

Example:

# ...

resources:
  gpu: v6e-4

This is shown as a single GPU of total 128GB instead of 4 GPU each 32GB.

Therefore, there are two issues:

1. As a user, for me, it's easier to think in terms of chips rather than core. It's easier to calculate GPU memory, etc.
2. As a user, I'd appreciate using the same syntax for specifying the number of GPUs as for other vendors even if there are fundamental differences between TPU and other chips.

Solution

Allow to specify the number of chips the same way w allow to specify the number of GPUs:

# ...

resources:
  gpu: v6e:4

And, in this case, show it as 4 chips each 32GB with dstack apply.

Note: We also need to check the maximum size of chips per node, and only allow to specify what is supported per node.

Workaround

No response

Would you like to help us implement this feature by sending a PR?

Yes