How to do multi-machine SPMD/FSDPv2 training with TPU？

[radna0]

❓ Questions and Help

I saw #6362 but there's no example training script found? For example, if I have multiple TPU v3-8 VMs, how would I achieve this with SPMD/FSDPv2?

I'm currently sending the commands to all TPU VMs this way:

python3.10 podrun --include-local -- hostname

[radna0]

*Update: I found https://github.com/pytorch/xla/blob/master/docs/source/perf/spmd_distributed_checkpoint.md#process-groups
Is there any example script for this?

[radna0]

Anybody can help? I'm still stuck on this