[Feature]: Add TPU support

[PrasannaBhattarai]

Problem

TPUs are widely available, having TPU support would be really useful.

Solution

Collaborate with Google platform for adding TPUs.

Benefit

Large scale, energy efficient solutions for deep learning. An additional option for users to work with.

Alternatives

No.

Would you like to help contributing this feature?

Yes

[r4victor]

@PrasannaBhattarai, thanks for the issue and sorry for the slow response. We're certainly considering adding TPU support at some point but haven't researched it thoroughly yet.

If my understanding is correct, to add TPU support to dstack, we have to support TPU-specific VM and docker images (https://cloud.google.com/tpu/docs/run-in-container).

There is also a question of how to specify TPU when requesting compute. Maybe treat like another GPU, e.g. --gpu tpu-v5p?

@PrasannaBhattarai, will keep you updated on the progress. Feel free to share your vision of the implementation if you're willing to help with this feature!

[peterschmidt85]

This issue is stale because it has been open for 30 days with no activity.

[Bihan]

@PrasannaBhattarai #1229