Regression on 2.1 RC RoCm: data parallel error on `torch._C._broadcast_coalesced`
Labels: high priority, oncall: binaries, in progress, module: rocm, oncall: releng, triaged

### 🐛 Describe the bug

Hi, I noticed today a bug in data parallel on RoCm system that is not present on NVIDIA ones.

Reproduction: run the following script `CUDA_VISIBLE_DEVICES=0,1 python run_dp.py`:

```python
from transformers import AutoModelForSequenceClassification, Trainer, TrainingArguments
import torch
from datasets import Dataset
import torch.nn as nn

model = AutoModelForSequenceClassification.from_pretrained("textattack/bert-base-uncased-SST-2")

training_arguments = TrainingArguments(output_dir="foo", skip_memory_metrics=True, per_device_train_batch_size=4)

dummy = {}

sequence_length = 128
n_samples = 256

dummy["input_ids"] = torch.randint(
        low=0,
        high=10,
        size=(n_samples, sequence_length))
dummy["attention_mask"] = torch.randint(
        low=0,
        high=2,
        size=(n_samples, sequence_length))

dummy["labels"] = torch.randint(
        low=0,
        high=2,
        size=(n_samples,))

task_dataset = Dataset.from_dict(dummy)
task_dataset.set_format(
    type="torch",  # for now we're using pytorch tensors
    columns=list(task_dataset.features.keys()),
)

train_dataset = task_dataset

print("train_dataset", train_dataset)
print("train_dataset", train_dataset[0])

# Same issue with Transformers Trainer.
"""
trainer = Trainer(
    model=model,
    args=training_arguments,
    train_dataset=train_dataset,
)
trainer.train()
"""

model = nn.DataParallel(model)
model.to("cuda")

inp = {
    "input_ids": train_dataset[:8]["input_ids"].to("cuda"),
    "attention_mask": train_dataset[:8]["attention_mask"].to("cuda"),
    "labels": train_dataset[:8]["labels"].to("cuda")
}
print("inp", inp)
output = model(**inp)
```

raises
```
Traceback (most recent call last):
  File "/home/user/transformers-regression/run_dp.py", line 58, in <module>
    output = model(**inp)
  File "/home/user/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/user/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/user/.local/lib/python3.10/site-packages/torch/nn/parallel/data_parallel.py", line 184, in forward
    replicas = self.replicate(self.module, self.device_ids[:len(inputs)])
  File "/home/user/.local/lib/python3.10/site-packages/torch/nn/parallel/data_parallel.py", line 189, in replicate
    return replicate(module, device_ids, not torch.is_grad_enabled())
  File "/home/user/.local/lib/python3.10/site-packages/torch/nn/parallel/replicate.py", line 110, in replicate
    param_copies = _broadcast_coalesced_reshape(params, devices, detach)
  File "/home/user/.local/lib/python3.10/site-packages/torch/nn/parallel/replicate.py", line 83, in _broadcast_coalesced_reshape
    tensor_copies = Broadcast.apply(devices, *tensors)
  File "/home/user/.local/lib/python3.10/site-packages/torch/autograd/function.py", line 539, in apply
    return super().apply(*args, **kwargs)  # type: ignore[misc]
  File "/home/user/.local/lib/python3.10/site-packages/torch/nn/parallel/_functions.py", line 23, in forward
    outputs = comm.broadcast_coalesced(inputs, ctx.target_gpus)
  File "/home/user/.local/lib/python3.10/site-packages/torch/nn/parallel/comm.py", line 57, in broadcast_coalesced
    return torch._C._broadcast_coalesced(tensors, devices, buffer_size)
RuntimeError: NCCL Error 3: internal error
```

on torch 2.1 RC (`2.1.0+rocm5.6`), while no error is raised on 2.1 RC (`2.1.0+cu118`) on an NVIDIA system.

On 2.0.1, there is no issue (specifically on `rocm/pytorch:rocm5.6_ubuntu20.04_py3.8_pytorch_2.0.1`).

### Versions

```
Collecting environment information...
PyTorch version: 2.1.0+rocm5.6
Is debug build: False
CUDA used to build PyTorch: N/A
ROCM used to build PyTorch: 5.6.31061-8c743ae5d

OS: Ubuntu 20.04.6 LTS (x86_64)
GCC version: (Ubuntu 9.4.0-1ubuntu1~20.04.1) 9.4.0
Clang version: 16.0.0 (https://github.com/RadeonOpenCompute/llvm-project roc-5.6.0 23243 be997b2f3651a41597d7a41441fff8ade4ac59ac)
CMake version: version 3.26.4
Libc version: glibc-2.31

Python version: 3.8.16 (default, Jun 12 2023, 18:09:05)  [GCC 11.2.0] (64-bit runtime)
Python platform: Linux-5.15.0-84-generic-x86_64-with-glibc2.17
Is CUDA available: True
CUDA runtime version: Could not collect
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: AMD Instinct MI210
Nvidia driver version: Could not collect
cuDNN version: Could not collect
HIP runtime version: 5.6.31061
MIOpen runtime version: 2.20.0
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):                             96
On-line CPU(s) list:                0-95
Thread(s) per core:                 1
Core(s) per socket:                 48
Socket(s):                          2
NUMA node(s):                       2
Vendor ID:                          AuthenticAMD
CPU family:                         25
Model:                              1
Model name:                         AMD EPYC 7643 48-Core Processor
Stepping:                           1
Frequency boost:                    enabled
CPU MHz:                            1500.000
CPU max MHz:                        3640.9170
CPU min MHz:                        1500.0000
BogoMIPS:                           4600.40
Virtualization:                     AMD-V
L1d cache:                          3 MiB
L1i cache:                          3 MiB
L2 cache:                           48 MiB
L3 cache:                           512 MiB
NUMA node0 CPU(s):                  0-47
NUMA node1 CPU(s):                  48-95
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 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 disabled, RSB filling, PBRSB-eIBRS Not affected
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 rapl pni pclmulqdq monitor ssse3 fma cx16 pcid sse4_1 sse4_2 movbe popcnt aes xsave avx f16c rdrand lahf_lm cmp_legacy svm extapic cr8_legacy abm sse4a misalignsse 3dnowprefetch osvw ibs skinit wdt tce topoext perfctr_core perfctr_nb bpext perfctr_llc mwaitx cpb cat_l3 cdp_l3 invpcid_single hw_pstate ssbd mba ibrs ibpb stibp vmmcall fsgsbase bmi1 avx2 smep bmi2 invpcid cqm rdt_a rdseed adx smap clflushopt clwb sha_ni xsaveopt xsavec xgetbv1 xsaves cqm_llc cqm_occup_llc cqm_mbm_total cqm_mbm_local clzero irperf xsaveerptr rdpru wbnoinvd amd_ppin arat npt lbrv svm_lock nrip_save tsc_scale vmcb_clean flushbyasid decodeassists pausefilter pfthreshold v_vmsave_vmload vgif v_spec_ctrl umip pku ospke vaes vpclmulqdq rdpid overflow_recov succor smca

Versions of relevant libraries:
[pip3] mypy==0.960
[pip3] mypy-extensions==1.0.0
[pip3] numpy==1.24.4
[pip3] pytorch-triton-rocm==2.1.0+34f8189eae
[pip3] torch==2.1.0+rocm5.6
[pip3] torchvision==0.15.2a0+fa99a53
[pip3] triton==2.0.0
[conda] No relevant packages
root@f3d18347865e:/var/lib/jenkins# [conda] No relevant packages
bash: [conda]: command not found

```

cc @ezyang @gchanan @zou3519 @kadeng @seemethere @malfet @jeffdaily @sunway513 @jithunnair-amd @pruthvistony @ROCmSupport @dllehr-amd @jataylo @hongxiayang