[Distributed] Remove sea-of-devices mode (#4905)
Summary:
This change removes sea-of-devices mode which is no longer valid and simply the code in all_reduce a lot. This is the 1st step to adopt pytorch/pytorch#93173.
Test Plan:
PJRT_DEVICE=TPU python test/test_mp_sync_batch_norm.py
PJRT_DEVICE=TPU python test/test_mp_replication.py
and CI...
