Repository: pytorch
Issue: torch.distributed.all_to_all_single & alltoall_base, size limit INT_MAX
Labels: oncall: distributed

### 🐛 Describe the bug

I am using torch.distributed.all_to_all_single function do alltoall communication. My input and output tensor are both 2d tensor. From my error code, it seems this function only support tensor which size <= INT_MAX. Besides, the `size` is not `length` of dimension 0,  it is actually the size of tensor.  However, as I check in the source code, I think it should be the length of first dimension of tensor. I do find there was a check in older version(branch `1x1` ,function `computeLengthsAndOffsets`), but not in PyTorch 2.0.1. In PyTorch 2.0.1, I also find cpp code use `int64_t` as datatype, so there should not be a INT_MAX limitation.

I want to know is there anything wrong with my software version or code. I have some self-define function and whole script is long, so it is not easy to upload the code for reproduction, but I print some information about input/output tensor and splits.

Thank you for you attention and help.

```python
# all_to_all_single
print_all("rank, ", my_rank, "shape[input|output], ", input.shape, output.shape)
print_all("rank, ", my_rank, "type[input|output], ", input.dtype, output.dtype)
print_all("rank, ", my_rank, "input split, ", batch_split_lengths)
print_all("rank, ", my_rank, "output split, ", table_split_lengths)
req = dist.all_to_all_single(
    output,
    input,
    output_split_sizes = table_split_lengths,
    input_split_sizes = batch_split_lengths,
    async_op=True)

```

Here is my error report.
```
rank,  1 shape[input|output],  torch.Size([26214400, 128]) torch.Size([29491200, 128])
rank,  1 type[input|output],  torch.uint8 torch.uint8
rank,  1 input split,  [6553600, 6553600, 6553600, 6553600]
rank,  1 output split,  [9830400, 6553600, 6553600, 6553600]
rank,  3 shape[input|output],  torch.Size([26214400, 128]) torch.Size([29491200, 128])
rank,  3 type[input|output],  torch.uint8 torch.uint8
rank,  3 input split,  [6553600, 6553600, 6553600, 6553600]
rank,  3 output split,  [9830400, 6553600, 6553600, 6553600]
rank,  2 shape[input|output],  torch.Size([26214400, 128]) torch.Size([29491200, 128])
rank,  2 type[input|output],  torch.uint8 torch.uint8
rank,  2 input split,  [6553600, 6553600, 6553600, 6553600]
rank,  2 output split,  [9830400, 6553600, 6553600, 6553600]
16:36:53
Using All2All_v_Req
embedding len is,  128
rank,  0 shape[input|output],  torch.Size([39321600, 128]) torch.Size([29491200, 128])
rank,  0 type[input|output],  torch.uint8 torch.uint8
rank,  0 input split,  [9830400, 9830400, 9830400, 9830400]
rank,  0 output split,  [9830400, 6553600, 6553600, 6553600]
Traceback (most recent call last):
  File "/geode2/home/u030/haofeng/BigRed200/dlrm/ext_dist_quan.py", line 63, in <module>
Traceback (most recent call last):
  File "/geode2/home/u030/haofeng/BigRed200/dlrm/ext_dist_quan.py", line 63, in <module>
    a2q_v_req = ext_dist.alltoall_v(send_tensor, send_cnt = send_cnt, receive_cnt = receive_cnt)
    a2q_v_req = ext_dist.alltoall_v(send_tensor, send_cnt = send_cnt, receive_cnt = receive_cnt)
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/geode2/home/u030/haofeng/BigRed200/dlrm/extend_distributed.py", line 693, in alltoall_v
    output = All2All_v_Req.apply(a2a_info, inputs) # I remove * , it used to be *inputs
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/N/slate/haofeng/anaconda3/envs/new_dlrm/lib/python3.11/site-packages/torch/autograd/function.py", line 506, in apply
    return super().apply(*args, **kwargs)  # type: ignore[misc]
File "/geode2/home/u030/haofeng/BigRed200/dlrm/extend_distributed.py", line 572, in forward
    req = dist.all_to_all_single(
          ^^^^^^^^^^^^^^^^^^^^^^^
File "/N/slate/haofeng/anaconda3/envs/new_dlrm/lib/python3.11/site-packages/torch/distributed/distributed_c10d.py", line 3155, in all_to_all_single
    work = default_pg.alltoall_base(
           ^^^^^^^^^^^^^^^^^^^^^^^^^
RuntimeError: length <= std::numeric_limits<int>::max() && offset <= std::numeric_limits<int>::max() INTERNAL ASSERT FAILED at "/opt/conda/conda-bld/pytorch_1682343970094/work/torch/csrc/distributed/c10d/Utils.hpp":463, please report a bug to PyTorch. Length or offset larger than INT_MAX not supported
```

### Versions

Collecting environment information...
PyTorch version: 2.0.1
Is debug build: False
CUDA used to build PyTorch: 11.7
ROCM used to build PyTorch: N/A

OS: SUSE Linux Enterprise Server 15 SP3 (x86_64)
GCC version: (SUSE Linux) 7.5.0
Clang version: Could not collect
CMake version: Could not collect
Libc version: glibc-2.31

Python version: 3.11.3 (main, May 15 2023, 15:45:52) [GCC 11.2.0] (64-bit runtime)
Python platform: Linux-5.3.18-150300.59.90-default-x86_64-with-glibc2.31
Is CUDA available: False
CUDA runtime version: 11.3.58
CUDA_MODULE_LOADING set to: N/A
GPU models and configuration: No devices found.
Nvidia driver version: Could not collect
cuDNN version: Could not collect
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):                          128
On-line CPU(s) list:             0-127
Thread(s) per core:              2
Core(s) per socket:              64
Socket(s):                       1
NUMA node(s):                    4
Vendor ID:                       AuthenticAMD
CPU family:                      25
Model:                           1
Model name:                      AMD EPYC 7713 64-Core Processor
Stepping:                        1
Frequency boost:                 enabled
CPU MHz:                         2224.573
CPU max MHz:                     2000.0000
CPU min MHz:                     1500.0000
BogoMIPS:                        3992.40
Virtualization:                  AMD-V
L1d cache:                       2 MiB
L1i cache:                       2 MiB
L2 cache:                        32 MiB
L3 cache:                        256 MiB
NUMA node0 CPU(s):               0-15,64-79
NUMA node1 CPU(s):               16-31,80-95
NUMA node2 CPU(s):               32-47,96-111
NUMA node3 CPU(s):               48-63,112-127
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 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 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 erms 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 wbnoinvd amd_ppin arat npt lbrv svm_lock nrip_save tsc_scale vmcb_clean flushbyasid decodeassists pausefilter pfthreshold v_vmsave_vmload vgif umip pku ospke vaes vpclmulqdq rdpid overflow_recov succor smca fsrm

Versions of relevant libraries:
[pip3] numpy==1.25.0
[pip3] torch==2.0.1
[pip3] triton==2.0.0
[conda] blas                      1.0                         mkl  
[conda] mkl                       2023.1.0         h6d00ec8_46342  
[conda] numpy                     1.25.0                   pypi_0    pypi
[conda] pytorch                   2.0.1           py3.11_cuda11.7_cudnn8.5.0_0    pytorch
[conda] pytorch-cuda              11.7                 h778d358_5    pytorch
[conda] pytorch-mutex             1.0                        cuda    pytorch
[conda] torchtriton               2.0.0                     py311    pytorch

cc @mrshenli @pritamdamania87 @zhaojuanmao @satgera @rohan-varma @gqchen @aazzolini @osalpekar @jiayisuse @H-Huang @kwen2501 @awgu