@aws-rhsoln @amithrm FYI, this might affect you guys, since now all-gather will become all_reduce based when layout is pin. We need to do this for TPU.

Thanks @JackCaoG for letting us know.

Hi @JackCaoG, is the pinning behavior because of the way compiler handles CC ops in TPUs? Or is this a requirement of FSDP? Checking because do we need to this as a compulsory fix and hence test out all our use cases

when layout is unpin, and each process generate slightly different graph(which we do in pytorch/xla), there is a chance data corruption will happen(because of some XLA optimization assume all graph will be the same). I don't know if this will be the case for you guys through(whether this is a TPU specified behavior), maybe check with your compiler team?