Multi-host parallelization with GPUs

[IrishWhiskey]

Hello,

I'm trying the JAX multi-host parallelization on AWS GPU instances but I can't get it working. What am I doing wrong? See below the process I followed.

I created two EC2 instances (p3.2xlarge) and set up the network configurations so that they can communicate with each other. After installing the required dependencies, I ran the following script on both instances:

import jax

jax.distributed.initialize(
    coordinator_address="172.31.28.213:4444",
    num_processes=2,
    process_id=0
)

xs = jax.numpy.ones(jax.local_device_count())

print(xs)
print(jax.pmap(lambda x: jax.lax.psum(x, 'i'), axis_name='i')(xs))

The script is not 100% same as in the second instance I set the process_id to 1 (the coordinator address refers to the first instance).

From the terminal output I can see that up until the print(xs) everything works fine but the last line causes an error. Below are the outputs I got from both instances.

First instance:

[1.]
2023-11-23 11:08:58.023623: E external/tsl/tsl/distributed_runtime/coordination/coordination_service.cc:1169] Shutdown barrier in coordination service has failed:
DEADLINE_EXCEEDED: Barrier timed out. Barrier_id: Shutdown::18196655163771244757. Timed out task names:
/job:jax_worker/replica:0/task:0
 [type.googleapis.com/tensorflow.CoordinationServiceError='']
This suggests that the workers are out of sync. Either at least one worker is too fast in its execution / crashed early or too slow / hanging. Check the logs for an earlier error to identify the root cause.
2023-11-23 11:08:58.023691: E external/tsl/tsl/distributed_runtime/coordination/coordination_service.cc:762] INTERNAL: Shutdown barrier has been passed with status: 'DEADLINE_EXCEEDED: Barrier timed out. Barrier_id: Shutdown::18196655163771244757. Timed out task names:
/job:jax_worker/replica:0/task:0
 [type.googleapis.com/tensorflow.CoordinationServiceError='']', but this task is not at the barrier yet. [type.googleapis.com/tensorflow.CoordinationServiceError='']
2023-11-23 11:08:58.023775: E external/tsl/tsl/distributed_runtime/coordination/coordination_service.cc:473] Stopping coordination service as shutdown barrier timed out. Check the task logs for an earlier error.
2023-11-23 11:09:09.989581: E external/tsl/tsl/distributed_runtime/coordination/coordination_service_agent.cc:770] Coordination agent is set to ERROR: INVALID_ARGUMENT: Unexpected heartbeat request from task: /job:jax_worker/replica:0/task:0. This usually implies an earlier error that caused coordination service to shut down before the workers disconnect. Check the task leader's logs for an earlier error to debug the root cause.
Additional GRPC error information from remote target unknown_target_for_coordination_leader while calling /tensorflow.CoordinationService/Heartbeat:
:{"created":"@1700737749.989459704","description":"Error received from peer ipv4:172.31.28.213:4444","file":"external/com_github_grpc_grpc/src/core/lib/surface/call.cc","file_line":1056,"grpc_message":"Unexpected heartbeat request from task: /job:jax_worker/replica:0/task:0. This usually implies an earlier error that caused coordination service to shut down before the workers disconnect. Check the task leader's logs for an earlier error to debug the root cause.","grpc_status":3} [type.googleapis.com/tensorflow.CoordinationServiceError='']
2023-11-23 11:09:09.989646: E external/xla/xla/pjrt/distributed/client.cc:96] Coordination service agent in error status: INVALID_ARGUMENT: Unexpected heartbeat request from task: /job:jax_worker/replica:0/task:0. This usually implies an earlier error that caused coordination service to shut down before the workers disconnect. Check the task leader's logs for an earlier error to debug the root cause.
Additional GRPC error information from remote target unknown_target_for_coordination_leader while calling /tensorflow.CoordinationService/Heartbeat:
:{"created":"@1700737749.989459704","description":"Error received from peer ipv4:172.31.28.213:4444","file":"external/com_github_grpc_grpc/src/core/lib/surface/call.cc","file_line":1056,"grpc_message":"Unexpected heartbeat request from task: /job:jax_worker/replica:0/task:0. This usually implies an earlier error that caused coordination service to shut down before the workers disconnect. Check the task leader's logs for an earlier error to debug the root cause.","grpc_status":3} [type.googleapis.com/tensorflow.CoordinationServiceError='']
2023-11-23 11:09:09.989670: F external/xla/xla/pjrt/distributed/client.h:77] Terminating process because the coordinator detected missing heartbeats. This most likely indicates that another task died; see the other task logs for more details. Disable Python buffering, i.e. `python -u`, to be sure to see all the previous output. Status: INVALID_ARGUMENT: Unexpected heartbeat request from task: /job:jax_worker/replica:0/task:0. This usually implies an earlier error that caused coordination service to shut down before the workers disconnect. Check the task leader's logs for an earlier error to debug the root cause.
Additional GRPC error information from remote target unknown_target_for_coordination_leader while calling /tensorflow.CoordinationService/Heartbeat:
:{"created":"@1700737749.989459704","description":"Error received from peer ipv4:172.31.28.213:4444","file":"external/com_github_grpc_grpc/src/core/lib/surface/call.cc","file_line":1056,"grpc_message":"Unexpected heartbeat request from task: /job:jax_worker/replica:0/task:0. This usually implies an earlier error that caused coordination service to shut down before the workers disconnect. Check the task leader's logs for an earlier error to debug the root cause.","grpc_status":3} [type.googleapis.com/tensorflow.CoordinationServiceError='']
Aborted (core dumped)

Second instance:

[1.]
2023-11-23 11:03:57.149896: W external/xla/xla/service/gpu/runtime/support.cc:58] Intercepted XLA runtime error:
INTERNAL: external/xla/xla/service/gpu/nccl_utils.cc:297: NCCL operation ncclCommInitRank(&comm, nranks, id, rank) failed: unhandled system error (run with NCCL_DEBUG=INFO for details)
2023-11-23 11:03:57.149997: E external/xla/xla/pjrt/pjrt_stream_executor_client.cc:2716] Execution of replica 1 failed: INTERNAL: Failed to execute XLA Runtime executable: run time error: custom call 'xla.gpu.all_reduce' failed: external/xla/xla/service/gpu/nccl_utils.cc:297: NCCL operation ncclCommInitRank(&comm, nranks, id, rank) failed: unhandled system error (run with NCCL_DEBUG=INFO for details); current tracing scope: all-reduce-start; current profiling annotation: XlaModule:#hlo_module=pmap__lambda_,program_id=2#.
jax.errors.SimplifiedTraceback: For simplicity, JAX has removed its internal frames from the traceback of the following exception. Set JAX_TRACEBACK_FILTERING=off to include these.

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/home/ubuntu/jax_distributed/test.py", line 12, in <module>
    print(jax.pmap(lambda x: jax.lax.psum(x, 'i'), axis_name='i')(xs))
jaxlib.xla_extension.XlaRuntimeError: INTERNAL: Failed to execute XLA Runtime executable: run time error: custom call 'xla.gpu.all_reduce' failed: external/xla/xla/service/gpu/nccl_utils.cc:297: NCCL operation ncclCommInitRank(&comm, nranks, id, rank) failed: unhandled system error (run with NCCL_DEBUG=INFO for details); current tracing scope: all-reduce-start; current profiling annotation: XlaModule:#hlo_module=pmap__lambda_,program_id=2#.
2023-11-23 11:08:57.152821: E external/tsl/tsl/distributed_runtime/coordination/coordination_service_agent.cc:494] Failed to disconnect from coordination service with status: DEADLINE_EXCEEDED: Deadline Exceeded
Additional GRPC error information from remote target unknown_target_for_coordination_leader while calling /tensorflow.CoordinationService/ShutdownTask:
:{"created":"@1700737737.152669627","description":"Error received from peer ipv4:172.31.28.213:4444","file":"external/com_github_grpc_grpc/src/core/lib/surface/call.cc","file_line":1056,"grpc_message":"Deadline Exceeded","grpc_status":4}
Proceeding with agent shutdown anyway. This is usually caused by an earlier error during execution. Check the logs (this task or the leader) for an earlier error to debug further.
Exception ignored in atexit callback: <function shutdown at 0x7f94fdbf6320>
Traceback (most recent call last):
  File "/opt/tensorflow/lib/python3.10/site-packages/jax/_src/distributed.py", line 189, in shutdown
    global_state.shutdown()
  File "/opt/tensorflow/lib/python3.10/site-packages/jax/_src/distributed.py", line 101, in shutdown
    self.client.shutdown()
jaxlib.xla_extension.XlaRuntimeError: DEADLINE_EXCEEDED: Deadline Exceeded
Additional GRPC error information from remote target unknown_target_for_coordination_leader while calling /tensorflow.CoordinationService/ShutdownTask:
:{"created":"@1700737737.152669627","description":"Error received from peer ipv4:172.31.28.213:4444","file":"external/com_github_grpc_grpc/src/core/lib/surface/call.cc","file_line":1056,"grpc_message":"Deadline Exceeded","grpc_status":4}

I'm using python 3.10.12, jax 0.4.20 and jaxlib 0.4.20+cuda11.cudnn86.

[mjsML]

You need the EFA drivers and certain magic bits for the stack to work properly in AWS.
May I suggest you take a look at https://github.com/NVIDIA/JAX-Toolbox ?

the jax container from jax-toolbox should work for AWS cc: @yhtang

[magicofai]

https://docs.aws.amazon.com/sagemaker/latest/dg/your-algorithms-training-efa.html

so you want https://github.com/aws/aws-ofi-nccl.git also, see the bash script there.

[sshuaizhang]

I have the same question.