Not able to run 3d_parallel.py
Author: quic-meetkuma
Pull request: no
Labels: name: bug

### System Info

- `transformers` version: 4.57.1
- Platform: Linux-5.4.0-218-generic-x86_64-with-glibc2.31
- Python version: 3.10.18
- Huggingface_hub version: 0.36.0
- Safetensors version: 0.6.2
- Accelerate version: 1.11.0
- Accelerate config:    not found
- DeepSpeed version: not installed
- PyTorch version (accelerator?): 2.9.0+cu126 (CUDA)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using distributed or parallel set-up in script?: <fill in>
- Using GPU in script?: [3D_parallel.py](https://github.com/huggingface/transformers/blob/main/examples/3D_parallel.py)
- GPU type: NVIDIA A100 80GB PCIe

### Who can help?

_No response_

### Information

- [x] The official example scripts
- [ ] My own modified scripts

### Tasks

- [x] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

Installed the dependencies.
1. Copy [3D_parallel.py](https://github.com/huggingface/transformers/blob/main/examples/3D_parallel.py) to some local file system.
2. Execute: "TP_SIZE=2 DP_SIZE=2 torchrun --nproc_per_node=4 --rdzv_endpoint
=localhost:29503 3D_parallel.py"

### Expected behavior

I am trying to learn the implementation of [3d_parallel.py](https://github.com/huggingface/transformers/blob/main/examples/3D_parallel.py). The code does not work out of the box for me.

Command: TP_SIZE=2 DP_SIZE=2 torchrun --nproc_per_node=4 --rdzv_endpoint=localhost:29503 examples/3D_parallel.py
Pytorch version: 2.9.0+cu126
CUDA version: 12.4 on host
Driver version: 550.163.01
OS: Ubuntu 20.04
Facing below error. I am observing this on rank1 only. I did not get any error on rank0.


[rank1]: return funcol.all_reduce(
[rank1]: File "/workspace/temp_env/lib/python3.10/site-packages/torch/distributed/_functional_collectives.py", line 176, in
all_reduce
[rank1]: tensor = torch.ops._c10d_functional.all_reduce(self, reduceOp.lower(), group_name)
[rank1]: File "/workspace/temp_env/lib/python3.10/site-packages/torch/_ops.py", line 1243, in call
[rank1]: return self._op(*args, **kwargs)
[rank1]: torch.distributed.DistBackendError: [1] is setting up NCCL communicator and retrieving ncclUniqueId from [0] via c10d key-value s
tore by key '0', but store->get('0') got error: wait timeout after 600000ms, keys: /default_pg/0//3//cuda//0
[rank1]: Exception raised from doWait at /pytorch/torch/csrc/distributed/c10d/TCPStore.cpp:597 (most recent call first):
[rank1]: frame #0: c10::Error::Error(c10::SourceLocation, std::__cxx11::basic_string<char, std::char_traits, std::allocator >)

0x80 (0x7f80e20b5eb0 in /workspace/temp_env/lib/python3.10/site-packages/torch/lib/libc10.so)
[rank1]: frame https://github.com/huggingface/transformers/pull/1: + 0x5cd0381 (0x7f8123a9e381 in /workspace/temp_env/lib/python3.10/site-packages/torch/lib/libtorch_cpu.so)
[rank1]: frame https://github.com/huggingface/transformers/pull/2: + 0x1361e6c (0x7f811f12fe6c in /workspace/temp_env/lib/python3.10/site-packages/torch/lib/libtorch_cpu.so)


Comments:
This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.

Reopening the thread

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.