Repository: pytorch
Issue: "Attribute arch of node cpu not found" error when running FSDP+TP on 8 H100s
Labels: oncall: distributed,release notes: distributed (fsdp2)

### 🐛 Describe the bug

When trying to run the [FSDP + TP official example](https://github.com/pytorch/examples/blob/main/distributed/tensor_parallelism/fsdp_tp_example.py) on a machine with 8 H100s or a [modified version of the code that uses the new FSDP2 api](https://gist.github.com/vturrisi/9ddfdfd6cbdd3ace177dad31135f3c86), we are getting the following error (that comes only from the last four GPUs):

```
[rank4]: Traceback (most recent call last):
[rank4]:   File "/mnt/task_runtime/main.py", line 143, in <module>
[rank4]:     model = parallelize_module(
[rank4]:   File "/miniconda/lib/python3.10/site-packages/torch/distributed/tensor/parallel/api.py", line 112, in parallelize_module
[rank4]:     parallelize_module(submodule, device_mesh, parallelize_style)
[rank4]:   File "/miniconda/lib/python3.10/site-packages/torch/distributed/tensor/parallel/api.py", line 85, in parallelize_module
[rank4]:     return parallelize_plan._apply(module, device_mesh)
[rank4]:   File "/miniconda/lib/python3.10/site-packages/torch/distributed/tensor/parallel/style.py", line 269, in _apply
[rank4]:     return distribute_module(
[rank4]:   File "/miniconda/lib/python3.10/site-packages/torch/distributed/tensor/_api.py", line 852, in distribute_module
[rank4]:     partition_fn(name, submod, device_mesh)
[rank4]:   File "/miniconda/lib/python3.10/site-packages/torch/distributed/tensor/parallel/style.py", line 242, in _partition_embedding_fn
[rank4]:     dist_param = nn.Parameter(distribute_tensor(param, device_mesh, [Shard(0)]))
[rank4]:   File "/miniconda/lib/python3.10/site-packages/torch/distributed/tensor/_api.py", line 735, in distribute_tensor
[rank4]:     local_tensor = placement._shard_tensor(local_tensor, device_mesh, idx)
[rank4]:   File "/miniconda/lib/python3.10/site-packages/torch/distributed/tensor/placement_types.py", line 176, in _shard_tensor
[rank4]:     mesh_scatter(output, scatter_list, mesh, mesh_dim=mesh_dim)
[rank4]:   File "/miniconda/lib/python3.10/site-packages/torch/distributed/tensor/_collective_utils.py", line 113, in mesh_scatter
[rank4]:     fut = scatter(
[rank4]:   File "/miniconda/lib/python3.10/site-packages/torch/distributed/c10d_logger.py", line 83, in wrapper
[rank4]:     return func(*args, **kwargs)
[rank4]:   File "/miniconda/lib/python3.10/site-packages/torch/distributed/distributed_c10d.py", line 3718, in scatter
[rank4]:     work = group.scatter(output_tensors, input_tensors, opts)
[rank4]: torch.distributed.DistBackendError: NCCL error in: /tmp/tmp.ASHKnc31xJ/pytorch/torch/csrc/distributed/c10d/NCCLUtils.hpp:317, internal error - please report this issue to the NCCL developers, NCCL version 2.19.1
[rank4]: ncclInternalError: Internal check failed.
[rank4]: Last error:
[rank4]: Attribute arch of node cpu not found
```

The same issue doesn't occur if using only the first four GPUs, but it always occurs when using all the GPUs, or including a subset of the last four.
This command will work:
```
CUDA_VISIBLE_DEVICES=0,1,2,3 torchrun --nnodes=1 --nproc_per_node=4 --standalone fsdp_tp_example.py
```

While these won't work:
```
CUDA_VISIBLE_DEVICES=4,5,6,7 torchrun --nnodes=1 --nproc_per_node=4 --standalone fsdp_tp_example.py
torchrun --nnodes=1 --nproc_per_node=6 --standalone fsdp_tp_example.py
torchrun --nnodes=1 --nproc_per_node=8 --standalone fsdp_tp_example.py
```

### Versions

```
PyTorch version: 2.5.1
Is debug build: False
CUDA used to build PyTorch: 12.3
ROCM used to build PyTorch: N/A

OS: Ubuntu 20.04.6 LTS (x86_64)
GCC version: (Ubuntu 9.4.0-1ubuntu1~20.04.2) 9.4.0
Clang version: Could not collect
CMake version: version 3.30.5
Libc version: glibc-2.31

Python version: 3.10.13 | packaged by conda-forge | (main, Dec 23 2023, 15:36:39) [GCC 12.3.0] (64-bit runtime)
Python platform: Linux-5.10.192-183.736.amzn2.x86_64-x86_64-with-glibc2.31
Is CUDA available: True
CUDA runtime version: 12.3.107
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration:
GPU 0: NVIDIA H100 80GB HBM3
GPU 1: NVIDIA H100 80GB HBM3
GPU 2: NVIDIA H100 80GB HBM3
GPU 3: NVIDIA H100 80GB HBM3
GPU 4: NVIDIA H100 80GB HBM3
GPU 5: NVIDIA H100 80GB HBM3
GPU 6: NVIDIA H100 80GB HBM3
GPU 7: NVIDIA H100 80GB HBM3

Nvidia driver version: 535.129.03
cuDNN version: Probably one of the following:
/usr/lib/x86_64-linux-gnu/libcudnn.so.9.5.1
/usr/lib/x86_64-linux-gnu/libcudnn_adv.so.9.5.1
/usr/lib/x86_64-linux-gnu/libcudnn_cnn.so.9.5.1
/usr/lib/x86_64-linux-gnu/libcudnn_engines_precompiled.so.9.5.1
/usr/lib/x86_64-linux-gnu/libcudnn_engines_runtime_compiled.so.9.5.1
/usr/lib/x86_64-linux-gnu/libcudnn_graph.so.9.5.1
/usr/lib/x86_64-linux-gnu/libcudnn_heuristic.so.9.5.1
/usr/lib/x86_64-linux-gnu/libcudnn_ops.so.9.5.1
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Architecture:                       x86_64
CPU op-mode(s):                     32-bit, 64-bit
Byte Order:                         Little Endian
Address sizes:                      48 bits physical, 48 bits virtual
CPU(s):                             192
On-line CPU(s) list:                0-191
Thread(s) per core:                 2
Core(s) per socket:                 48
Socket(s):                          2
NUMA node(s):                       2
Vendor ID:                          AuthenticAMD
CPU family:                         25
Model:                              1
Model name:                         AMD EPYC 7R13 Processor
Stepping:                           1
CPU MHz:                            2031.335
BogoMIPS:                           5300.00
Hypervisor vendor:                  KVM
Virtualization type:                full
L1d cache:                          3 MiB
L1i cache:                          3 MiB
L2 cache:                           48 MiB
L3 cache:                           384 MiB
NUMA node0 CPU(s):                  0-47,96-143
NUMA node1 CPU(s):                  48-95,144-191
Vulnerability Gather data sampling: Not affected
Vulnerability Itlb multihit:        Not affected
Vulnerability L1tf:                 Not affected
Vulnerability Mds:                  Not affected
Vulnerability Meltdown:             Not affected
Vulnerability Mmio stale data:      Not affected
Vulnerability Retbleed:             Not affected
Vulnerability Spec rstack overflow: Mitigation; safe RET
Vulnerability Spec store bypass:    Mitigation; Speculative Store Bypass disabled via prctl and seccomp
Vulnerability Spectre v1:           Mitigation; usercopy/swapgs barriers and __user pointer sanitization
Vulnerability Spectre v2:           Mitigation; Retpolines, IBPB conditional, IBRS_FW, STIBP always-on, RSB filling
Vulnerability Srbds:                Not affected
Vulnerability Tsx async abort:      Not affected
Flags:                              fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush mmx fxsr sse sse2 ht syscall nx mmxext fxsr_opt pdpe1gb rdtscp lm constant_tsc rep_good nopl nonstop_tsc cpuid extd_apicid aperfmperf tsc_known_freq pni pclmulqdq monitor ssse3 fma cx16 pcid sse4_1 sse4_2 x2apic movbe popcnt aes xsave avx f16c rdrand hypervisor lahf_lm cmp_legacy cr8_legacy abm sse4a misalignsse 3dnowprefetch topoext perfctr_core invpcid_single ssbd ibrs ibpb stibp vmmcall fsgsbase bmi1 avx2 smep bmi2 invpcid rdseed adx smap clflushopt clwb sha_ni xsaveopt xsavec xgetbv1 clzero xsaveerptr rdpru wbnoinvd arat npt nrip_save vaes vpclmulqdq rdpid

Versions of relevant libraries:
[pip3] dctorch==0.1.2
[pip3] flake8==7.0.0
[pip3] mypy==1.11.2
[pip3] mypy-extensions==1.0.0
[pip3] numpy==1.26.4
[pip3] numpydoc==1.7.0
[pip3] optree==0.13.0
[pip3] parametrized_torchvision==0.0.1
[pip3] pytorch-lightning==2.4.0
[pip3] pytorch-msssim==1.0.0
[pip3] pytorchvideo==0.1.5
[pip3] torch==2.5.1
[pip3] torchaudio==2.5.0
[pip3] torchdiffeq==0.2.4
[pip3] torchlayers==0.1.1
[pip3] torchmetrics==1.5.1
[pip3] torchsde==0.2.6
[pip3] torchsummary==1.5.1
[pip3] torchvision==0.20.0
[pip3] triton==3.1.0
[conda] _anaconda_depends         2024.10             py310_mkl_0
[conda] blas                      1.0                         mkl
[conda] dctorch                   0.1.2                    pypi_0    pypi
[conda] magma-cuda121             2.6.1                         1    pytorch
[conda] mkl                       2023.1.0         h213fc3f_46344
[conda] mkl-include               2025.0.0                 pypi_0    pypi
[conda] mkl-service               2.4.0           py310h5eee18b_1
[conda] mkl-static                2025.0.0                 pypi_0    pypi
[conda] mkl_fft                   1.3.10          py310h5eee18b_0
[conda] mkl_random                1.2.7           py310h1128e8f_0
[conda] numpy                     1.26.4          py310h5f9d8c6_0
[conda] numpy-base                1.26.4          py310hb5e798b_0
[conda] numpydoc                  1.7.0           py310h06a4308_0
[conda] optree                    0.13.0                   pypi_0    pypi
[conda] parametrized-torchvision  0.0.1                    pypi_0    pypi
[conda] pytorch-lightning         2.4.0                    pypi_0    pypi
[conda] pytorch-msssim            1.0.0                    pypi_0    pypi
[conda] pytorchvideo              0.1.5                    pypi_0    pypi
[conda] torch                     2.5.1                    pypi_0    pypi
[conda] torchaudio                2.5.0                    pypi_0    pypi
[conda] torchdiffeq               0.2.4                    pypi_0    pypi
[conda] torchlayers               0.1.1                    pypi_0    pypi
[conda] torchmetrics              1.5.1                    pypi_0    pypi
[conda] torchsde                  0.2.6                    pypi_0    pypi
[conda] torchsummary              1.5.1                    pypi_0    pypi
[conda] torchvision               0.20.0                   pypi_0    pypi
[conda] triton                    3.1.0                    pypi_0    pypi
```

cc @H-Huang @awgu @kwen2501 @wanchaol @fegin @fduwjj @wz337 @wconstab @d4l3k @c-p-i-o