How to run t5x on multi-node GPUs?

I'm exploring how to use t5x in a multi-node GPU setting. I'm using SLURM with a singularity container to execute the training script.

```
#!/bin/bash
#SBATCH --partition=gpu
#SBATCH --job-name=lintang-t5x-multinode
#SBATCH --nodes=2
#SBATCH --gpus-per-node=8
#SBATCH --ntasks-per-node=1
#SBATCH --output=logs/%x_%j.out
#SBATCH --cpus-per-task=32
#SBATCH --exclusive
#SBATCH --requeue
#SBATCH --wait-all-nodes=1
#SBATCH --comment=ProjectName

# Cache Directories
export SINGULARITY_CACHEDIR=...
export BASE_DIR="..."

export PROJECT_DIR=${BASE_DIR}"..."
export MODEL_DIR="..."
export TFDS_DATA_DIR="..."

# directory where the T5X repo is cloned.
export T5X_DIR=${BASE_DIR}"..."
export PYTHONPATH=${PROJECT_DIR}

# export TF_XLA_FLAGS="--tf_xla_auto_jit=2"
export XLA_FLAGS="--xla_gpu_force_compilation_parallelism=1" # Hacky and don't want

singularity exec \
--nv --bind /fsx:/fsx t5x-env.sif \
python ${T5X_DIR}/t5x/train.py \
--gin_search_paths=${PROJECT_DIR} \
--gin_file="config-base.gin" \
--gin.MODEL_DIR=\"${MODEL_DIR}\" \
--gin.USE_CACHED_TASKS=False \
--alsologtostderr \
--multiprocess_gpu \
--coordinator_address="${SLURM_LAUNCH_NODE_IPADDR}:29500" \
--process_count "${SLURM_NPROCS}" \
--process_index "${SLURM_PROCID}"
```

But this doesn't seem to work.
```
2022-10-12 15:03:54.406737: E tensorflow/stream_executor/cuda/cuda_blas.cc:2981] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered
2022-10-12 15:03:55.491066: W tensorflow/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libnvinfer.so.7'; dlerror: libnvinfer.so.7: cannot open shared object file: No such file or directory; LD_LIBRARY_PATH: /opt/nvidia/hpc_sdk/Linux_x86_64/22.7/comm_libs/hpcx/hpcx-2.11/hcoll/lib:/opt/nvidia/hpc_sdk/Linux_x86_64/22.7/comm_libs/hpcx/hpcx-2.11/ompi/lib:/opt/nvidia/hpc_sdk/Linux_x86_64/22.7/comm_libs/hpcx/hpcx-2.11/nccl_rdma_sharp_plugin/lib:/opt/nvidia/hpc_sdk/Linux_x86_64/22.7/comm_libs/hpcx/hpcx-2.11/sharp/lib:/opt/nvidia/hpc_sdk/Linux_x86_64/22.7/comm_libs/hpcx/hpcx-2.11/ucx/mt/lib:/opt/nvidia/hpc_sdk/Linux_x86_64/22.7/comm_libs/hpcx/hpcx-2.11/ucx/mt/lib/ucx:/opt/nvidia/hpc_sdk/Linux_x86_64/22.7/comm_libs/11.7/nccl/lib:/opt/nvidia/hpc_sdk/Linux_x86_64/22.7/compilers/lib:/opt/nvidia/hpc_sdk/Linux_x86_64/22.7/cuda/11.7/lib64:/opt/nvidia/hpc_sdk/Linux_x86_64/22.7/math_libs/11.7/lib64:::/.singularity.d/libs
2022-10-12 15:03:55.492223: W tensorflow/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libnvinfer_plugin.so.7'; dlerror: libnvinfer_plugin.so.7: cannot open shared object file: No such file or directory; LD_LIBRARY_PATH: /opt/nvidia/hpc_sdk/Linux_x86_64/22.7/comm_libs/hpcx/hpcx-2.11/hcoll/lib:/opt/nvidia/hpc_sdk/Linux_x86_64/22.7/comm_libs/hpcx/hpcx-2.11/ompi/lib:/opt/nvidia/hpc_sdk/Linux_x86_64/22.7/comm_libs/hpcx/hpcx-2.11/nccl_rdma_sharp_plugin/lib:/opt/nvidia/hpc_sdk/Linux_x86_64/22.7/comm_libs/hpcx/hpcx-2.11/sharp/lib:/opt/nvidia/hpc_sdk/Linux_x86_64/22.7/comm_libs/hpcx/hpcx-2.11/ucx/mt/lib:/opt/nvidia/hpc_sdk/Linux_x86_64/22.7/comm_libs/hpcx/hpcx-2.11/ucx/mt/lib/ucx:/opt/nvidia/hpc_sdk/Linux_x86_64/22.7/comm_libs/11.7/nccl/lib:/opt/nvidia/hpc_sdk/Linux_x86_64/22.7/compilers/lib:/opt/nvidia/hpc_sdk/Linux_x86_64/22.7/cuda/11.7/lib64:/opt/nvidia/hpc_sdk/Linux_x86_64/22.7/math_libs/11.7/lib64:::/.singularity.d/libs
2022-10-12 15:03:55.492252: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Cannot dlopen some TensorRT libraries. If you would like to use Nvidia GPU with TensorRT, please make sure the missing libraries mentioned above are installed properly.
2022-10-12 15:03:59.741770: W tensorflow/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libcudnn.so.8'; dlerror: libcudnn.so.8: cannot open shared object file: No such file or directory; LD_LIBRARY_PATH: /opt/nvidia/hpc_sdk/Linux_x86_64/22.7/comm_libs/hpcx/hpcx-2.11/hcoll/lib:/opt/nvidia/hpc_sdk/Linux_x86_64/22.7/comm_libs/hpcx/hpcx-2.11/ompi/lib:/opt/nvidia/hpc_sdk/Linux_x86_64/22.7/comm_libs/hpcx/hpcx-2.11/nccl_rdma_sharp_plugin/lib:/opt/nvidia/hpc_sdk/Linux_x86_64/22.7/comm_libs/hpcx/hpcx-2.11/sharp/lib:/opt/nvidia/hpc_sdk/Linux_x86_64/22.7/comm_libs/hpcx/hpcx-2.11/ucx/mt/lib:/opt/nvidia/hpc_sdk/Linux_x86_64/22.7/comm_libs/hpcx/hpcx-2.11/ucx/mt/lib/ucx:/opt/nvidia/hpc_sdk/Linux_x86_64/22.7/comm_libs/11.7/nccl/lib:/opt/nvidia/hpc_sdk/Linux_x86_64/22.7/compilers/lib:/opt/nvidia/hpc_sdk/Linux_x86_64/22.7/cuda/11.7/lib64:/opt/nvidia/hpc_sdk/Linux_x86_64/22.7/math_libs/11.7/lib64:::/.singularity.d/libs
2022-10-12 15:03:59.743346: W tensorflow/core/common_runtime/gpu/gpu_device.cc:1934] Cannot dlopen some GPU libraries. Please make sure the missing libraries mentioned above are installed properly if you would like to use GPU. Follow the guide at https://www.tensorflow.org/install/gpu for how to download and setup the required libraries for your platform.
Skipping registering GPU devices...
I1012 15:03:59.743987 140385279873024 train.py:725] Initializing distributed system for multi-host GPU:
coordinator_address: :29500
process_count: 2
process_index: 0
I1012 15:03:59.744206 140385279873024 distributed.py:58] JAX distributed initialized with visible devices: 0
I1012 15:03:59.744523 140385279873024 distributed.py:67] Starting JAX distributed service on :29500
E1012 15:03:59.747199977 8449 server_chttp2.cc:40] {"created":"@1665587039.747188280","description":"Name or service not known","errno":-2,"file":"external/com_github_grpc_grpc/src/core/lib/iomgr/resolve_address_posix.cc","file_line":108,"os_error":"Name or service not known","syscall":"getaddrinfo","target_address":":29500"}
Fatal Python error: Segmentation fault

Thread 0x00007fadfebd1000 (most recent call first):
File "/usr/local/lib/python3.8/site-packages/jax/_src/distributed.py", line 68 in initialize
File "/usr/local/lib/python3.8/site-packages/jax/_src/distributed.py", line 159 in initialize
File "/fsx/lintangsutawika/architecture-objective/t5x/train.py", line 730 in _main
File "/fsx/lintangsutawika/architecture-objective/t5x/train.py", line 710 in main
File "/fsx/home-lintangsutawika/.local/lib/python3.8/site-packages/absl/app.py", line 254 in _run_main
File "/fsx/home-lintangsutawika/.local/lib/python3.8/site-packages/absl/app.py", line 308 in run
File "/fsx/lintangsutawika/t5x/t5x/gin_utils.py", line 107 in run
File "/fsx/lintangsutawika/architecture-objective/t5x/train.py", line 750 in
/var/spool/slurmd/job08329/slurm_script: line 53: 8420 Segmentation fault singularity exec --nv --bind /fsx:/fsx t5x-env.sif python ${T5X_DIR}/t5x/train.py --gin_search_paths=${PROJECT_DIR} --gin_file="config-base.gin" --gin.MODEL_DIR=\"${MODEL_DIR}\" --gin.USE_CACHED_TASKS=False --alsologtostderr --multiprocess_gpu --coordinator_address="${SLURM_LAUNCH_NODE_IPADDR}:29500" --process_count "${SLURM_NPROCS}" --process_index "${SLURM_PROCID}"

```

Another method I tried is try to launch two process with hard-coded `process_index` (I did this in an interactive shell)
```
singularity exec --nv --bind /fsx:/fsx /fsx/lintangsutawika/t5x-env.sif \
python ${T5X_DIR}/t5x/train.py \
--gin_search_paths=${PROJECT_DIR} \
--gin_file="config-base.gin" \
--gin.MODEL_DIR=\"${MODEL_DIR}\" \
--gin.USE_CACHED_TASKS=False \
--alsologtostderr \
--multiprocess_gpu \
--coordinator_address="${SLURM_LAUNCH_NODE_IPADDR}:29500" \
--process_count "${SLURM_NPROCS}" \
--process_index 0
singularity exec --nv --bind /fsx:/fsx /fsx/lintangsutawika/t5x-env.sif \
python ${T5X_DIR}/t5x/train.py \
--gin_search_paths=${PROJECT_DIR} \
--gin_file="config-base.gin" \
--gin.MODEL_DIR=\"${MODEL_DIR}\" \
--gin.USE_CACHED_TASKS=False \
--alsologtostderr \
--multiprocess_gpu \
--coordinator_address="${SLURM_LAUNCH_NODE_IPADDR}:29500" \
--process_count "${SLURM_NPROCS}" \
--process_index 1
```

process_index 1 seems to working as intended
```
I1012 16:17:41.957645 140701990588416 train.py:725] Initializing distributed system for multi-host GPU:
coordinator_address: 172.31.37.44:29500
process_count: 2
process_index: 1
I1012 16:17:41.957846 140701990588416 distributed.py:58] JAX distributed initialized with visible devices: 0
I1012 16:17:41.986860 140701990588416 distributed.py:78] Connecting to JAX distributed service on 172.31.37.44:29500
```
but process_index 0 fails.
```
I1012 16:18:04.230459 140621940224000 train.py:725] Initializing distributed system for multi-host GPU:
coordinator_address: 172.31.37.44:29500
process_count: 2
process_index: 0
I1012 16:18:04.230640 140621940224000 distributed.py:58] JAX distributed initialized with visible devices: 0
I1012 16:18:04.230890 140621940224000 distributed.py:67] Starting JAX distributed service on 172.31.37.44:29500
E1012 16:18:04.231958535 6069 server_chttp2.cc:40] {"created":"@1665591484.231935845","description":"No address added out of total 1 resolved","file":"external/com_github_grpc_grpc/src/core/ext/transport/chttp2/server/chttp2_server.cc","file_line":395,"referenced_errors":[{"created":"@1665591484.231933760","description":"Unable to configure socket","fd":30,"file":"external/com_github_grpc_grpc/src/core/lib/iomgr/tcp_server_utils_posix_common.cc","file_line":215,"referenced_errors":[{"created":"@1665591484.231931381","description":"Cannot assign requested address","errno":99,"file":"external/com_github_grpc_grpc/src/core/lib/iomgr/tcp_server_utils_posix_common.cc","file_line":189,"os_error":"Cannot assign requested address","syscall":"bind"}]}]}
Fatal Python error: Segmentation fault

Thread 0x00007fe518cb7000 (most recent call first):
File "/usr/local/lib/python3.8/site-packages/jax/_src/distributed.py", line 68 in initialize
File "/usr/local/lib/python3.8/site-packages/jax/_src/distributed.py", line 159 in initialize
File "/fsx/lintangsutawika/architecture-objective/t5x/train.py", line 730 in _main
File "/fsx/lintangsutawika/architecture-objective/t5x/train.py", line 710 in main
File "/fsx/home-lintangsutawika/.local/lib/python3.8/site-packages/absl/app.py", line 254 in _run_main
File "/fsx/home-lintangsutawika/.local/lib/python3.8/site-packages/absl/app.py", line 308 in run
File "/fsx/lintangsutawika/t5x/t5x/gin_utils.py", line 107 in run
File "/fsx/lintangsutawika/architecture-objective/t5x/train.py", line 750 in
Segmentation fault
```

Start with the distributed initialization in t5x/train.py (_main/main) and the launch path in t5x/gin_utils.py, using the reported SLURM and Singularity commands as the reproduction. Compare coordinator_address, process_count, and process_index with the JAX logs for both processes. Done means a documented multi-node launch reproduces training without the address-resolution or segmentation-fault failures.

Written by the indexing model from the issue text.