jax.device_count() Results on TPU v4-32

[yixiaoer]

Each TPU v3 chip contains two TensorCores, and a TPU v3-8 includes 4 chips, totaling 8 TensorCores. Running jax.device_count() on a TPU v3-8 returns 8, which matches expectations.

However, given that each TPU v4 chip also contains two TensorCores, I would expect a TPU v4-32 setup (theoretically comprising 16 chips) to report 32 TensorCores when jax.device_count() is executed. Instead, it returns 16, why this is the case?

[hawkinsp]

Tensorcores on TPU v4 chips can actually be used in two modes: either each each core works independently, in which case you see 32 devices, or both cores work together, in which case you see 16, equal to the number of chips. If I remember right we default to the latter mode, since it's usually what you want.

(This "working together" capability is not available on TPUv3 or earlier.)

[yixiaoer]

Thank you so much for the reply about TPU v4-32 devices!
To further clarify, does this mean that, across different TPU versions, the capability of each TensorCore is similar, but the "working together" mode on TPU v4 is utilized to enhance their power through increased memory and bandwidth? Is this why the "working together" mode is generally preferred?

[hawkinsp]

• the individual tensorcores are more powerful on TPUv4 vs TPUv3, as you'd expect with a newer generation of a chip.
• In TPUv3 each core has its own private HBM memory not shared with the other core, whereas in TPUv4 the HBM memory is shared between cores.
• Since the memory (and memory bandwidth) is shared, it's better if the compiler schedules both cores together, and it's usually simpler to think about for the user: you don't have to shard your data quite as much!

[yixiaoer]

Thanks for the explanation! Really helped me get a better grasp on TPUs. Cheers!