Distributed training is stuck

[syyxsxx]

I use two 4090 host for data parallel distributed training by jax.distributed, like this:
jax.distributed.initialize(coordinator_address="[ip]:[port]",
num_processes=2,
process_id=[index])
the train is stuck when doing all_reduce ops

How can I debug this problem？
Are there any examples for parallel distributed training