Repository: pytorch
Issue: Ineffective Buffer Size Adjustment in Multi-Node NCCL Communication
Labels: triaged,module: nccl

### 🐛 Describe the bug

I am experiencing an issue with NCCL when performing data transfers between two nodes. The data transfer works fine when sending smaller tensor shapes, but hangs when larger tensor shapes are used. Specifically, when I attempt to send tensors with the shape (8, 128, 10) from node 0 to node 1, node 0 indicates that data has been sent, but node 1 does not respond. However, when I reduce the shape to (8, 128, 2), the data is sent and received successfully. Adjusting the `NCCL_BUFFSIZE` environment variable does not resolve the issue for the larger tensor shape.

```
# from llama2_pp import pretrain


import torch
import argparse
import torch.distributed as dist
import torch

import os


def set_args():
    global _GLOBAL_ARGS
    
    parser = argparse.ArgumentParser()
    parser.add_argument('--distributed-backend', default='nccl',
                    choices=['nccl', 'gloo'])
    parser.add_argument('--local-rank', type=int, default=None,
                      help='local rank passed from distributed launcher.')
    args = parser.parse_args()
    # Distributed args.
    args.rank = int(os.getenv('RANK', '0'))
    args.local_rank = int(os.environ['LOCAL_RANK'])
    args.world_size = int(os.getenv("WORLD_SIZE", '1'))
    _GLOBAL_ARGS = args
    # return _GLOBAL_ARGS


def get_args():
    return _GLOBAL_ARGS

def _initialize_distributed():
    """Initialize torch.distributed and mpu."""
    args = get_args()
    print(args)

    # Call the init process
    init_method = 'tcp://'
    master_ip = os.getenv('MASTER_ADDR', 'localhost')
    master_port = os.getenv('MASTER_PORT', '6000')
    init_method += master_ip + ':' + master_port
    torch.distributed.init_process_group(
        backend=args.distributed_backend,
        world_size=args.world_size, rank=args.rank,
        init_method=init_method)



def main():

    set_args()
    args = get_args()
    
    # 初始化分布式环境
    _initialize_distributed()
    tensor_shape = (8,128,10)  

    torch.distributed.barrier()
    print('barrier done')
    
    rank = args.rank
    device = args.local_rank  
    torch.cuda.set_device(device)
    
    if rank == 0:
        input_tensor = torch.zeros(tensor_shape, device=device, requires_grad=True)
    else:
        input_tensor = torch.empty(tensor_shape, device=device)
        print(f'Rank {rank} waiting to receive from Rank {rank-1}')
        dist.recv(tensor=input_tensor, src=rank-1)  # 改为同步recv
        print(f'Rank {rank} received data from Rank {rank-1}')
        
    print(f'Rank {rank} input_tensor: {input_tensor}')
    output_tensor = input_tensor + 1

    if rank != args.world_size - 1:
        print(f'Rank {rank} sending to Rank {rank+1}')
        dist.send(tensor=output_tensor, dst=rank+1)
        print(f'Rank {rank} sent data to Rank {rank+1}')
        
    print(f'Process {rank} done')

if __name__ == "__main__":
    main()
    

```



node0:
```
#! /bin/bash

# Runs the "345M" parameter model
NGPUS_PER_NODE=8
STAGES_PER_NODE=2
GPUS_PER_STAGE=$(($NGPUS_PER_NODE/$STAGES_PER_NODE))
# Change for multinode config
MASTER_ADDR=172.16.1.83
MASTER_PORT=6002
NNODES=2
NODE_RANK=0
PIPELINE_SIZE=$(($NNODES * $STAGES_PER_NODE))
export NCCL_BUFFSIZE=419430400  # 将缓冲区大小设置为4MB
export NCCL_MIN_NRINGS=2      # 增加环的数量，可能会提高带宽利用率
DISTRIBUTED_ARGS="--nproc_per_node $STAGES_PER_NODE --nnodes $NNODES --node_rank $NODE_RANK --master_addr $MASTER_ADDR --master_port $MASTER_PORT"
python -m torch.distributed.launch $DISTRIBUTED_ARGS \
       train_llama2.py \
       --distributed-backend nccl \

```

node1:

```

#! /bin/bash

# Runs the "345M" parameter model
NGPUS_PER_NODE=8
STAGES_PER_NODE=2
GPUS_PER_STAGE=$(($NGPUS_PER_NODE/$STAGES_PER_NODE))
# Change for multinode config
MASTER_ADDR=172.16.1.83
MASTER_PORT=6002
NNODES=2
NODE_RANK=1
PIPELINE_SIZE=$(($NNODES * $STAGES_PER_NODE))
# WORLD_SIZE=$(($STAGES_PER_NODE*$NNODES))
export NCCL_BUFFSIZE=419430400  # 将缓冲区大小设置为4MB
export NCCL_MIN_NRINGS=2      # 增加环的数量，可能会提高带宽利用率
DISTRIBUTED_ARGS="--nproc_per_node $STAGES_PER_NODE --nnodes $NNODES --node_rank $NODE_RANK --master_addr $MASTER_ADDR --master_port $MASTER_PORT"
python -m torch.distributed.launch $DISTRIBUTED_ARGS \
       train_llama2.py \
       --distributed-backend nccl \

```

the error:

node0:
```
/root/anaconda3/envs/gees_pytorch_3.10/lib/python3.10/site-packages/torch/distributed/launch.py:181: FutureWarning: The module torch.distributed.launch is deprecated
and will be removed in future. Use torchrun.
Note that --use-env is set by default in torchrun.
If your script expects `--local-rank` argument to be set, please
change it to read from `os.environ['LOCAL_RANK']` instead. See 
https://pytorch.org/docs/stable/distributed.html#launch-utility for 
further instructions

  warnings.warn(
[2024-07-15 22:49:55,968] torch.distributed.run: [WARNING] 
[2024-07-15 22:49:55,968] torch.distributed.run: [WARNING] *****************************************
[2024-07-15 22:49:55,968] torch.distributed.run: [WARNING] Setting OMP_NUM_THREADS environment variable for each process to be 1 in default, to avoid your system being overloaded, please further tune the variable for optimal performance in your application as needed. 
[2024-07-15 22:49:55,968] torch.distributed.run: [WARNING] *****************************************
Namespace(distributed_backend='nccl', local_rank=0, rank=0, world_size=4)
Namespace(distributed_backend='nccl', local_rank=1, rank=1, world_size=4)
barrier donebarrier done

Rank 1 waiting to receive from Rank 0
[W ProcessGroupNCCL.cpp:1856] Warning: 0NCCL_AVOID_RECORD_STREAMS=1 has no effect for point-to-point collectives. (function operator())
Rank 0 input_tensor: tensor([[[0., 0., 0.,  ..., 0., 0., 0.],
         [0., 0., 0.,  ..., 0., 0., 0.],
         [0., 0., 0.,  ..., 0., 0., 0.],
         ...,
         [0., 0., 0.,  ..., 0., 0., 0.],
         [0., 0., 0.,  ..., 0., 0., 0.],
         [0., 0., 0.,  ..., 0., 0., 0.]],

        [[0., 0., 0.,  ..., 0., 0., 0.],
         [0., 0., 0.,  ..., 0., 0., 0.],
         [0., 0., 0.,  ..., 0., 0., 0.],
         ...,
         [0., 0., 0.,  ..., 0., 0., 0.],
         [0., 0., 0.,  ..., 0., 0., 0.],
         [0., 0., 0.,  ..., 0., 0., 0.]],

        [[0., 0., 0.,  ..., 0., 0., 0.],
         [0., 0., 0.,  ..., 0., 0., 0.],
         [0., 0., 0.,  ..., 0., 0., 0.],
         ...,
         [0., 0., 0.,  ..., 0., 0., 0.],
         [0., 0., 0.,  ..., 0., 0., 0.],
         [0., 0., 0.,  ..., 0., 0., 0.]],

        ...,

        [[0., 0., 0.,  ..., 0., 0., 0.],
         [0., 0., 0.,  ..., 0., 0., 0.],
         [0., 0., 0.,  ..., 0., 0., 0.],
         ...,
         [0., 0., 0.,  ..., 0., 0., 0.],
         [0., 0., 0.,  ..., 0., 0., 0.],
         [0., 0., 0.,  ..., 0., 0., 0.]],

        [[0., 0., 0.,  ..., 0., 0., 0.],
         [0., 0., 0.,  ..., 0., 0., 0.],
         [0., 0., 0.,  ..., 0., 0., 0.],
         ...,
         [0., 0., 0.,  ..., 0., 0., 0.],
         [0., 0., 0.,  ..., 0., 0., 0.],
         [0., 0., 0.,  ..., 0., 0., 0.]],

        [[0., 0., 0.,  ..., 0., 0., 0.],
         [0., 0., 0.,  ..., 0., 0., 0.],
         [0., 0., 0.,  ..., 0., 0., 0.],
         ...,
         [0., 0., 0.,  ..., 0., 0., 0.],
         [0., 0., 0.,  ..., 0., 0., 0.],
         [0., 0., 0.,  ..., 0., 0., 0.]]], device='cuda:0', requires_grad=True)
Rank 0 sending to Rank 1
[W ProcessGroupNCCL.cpp:1856] Warning: 0NCCL_AVOID_RECORD_STREAMS=1 has no effect for point-to-point collectives. (function operator())
Rank 1 received data from Rank 0
Rank 0 sent data to Rank 1
Process 0 done
Rank 1 input_tensor: tensor([[[1., 1., 1.,  ..., 1., 1., 1.],
         [1., 1., 1.,  ..., 1., 1., 1.],
         [1., 1., 1.,  ..., 1., 1., 1.],
         ...,
         [1., 1., 1.,  ..., 1., 1., 1.],
         [1., 1., 1.,  ..., 1., 1., 1.],
         [1., 1., 1.,  ..., 1., 1., 1.]],

        [[1., 1., 1.,  ..., 1., 1., 1.],
         [1., 1., 1.,  ..., 1., 1., 1.],
         [1., 1., 1.,  ..., 1., 1., 1.],
         ...,
         [1., 1., 1.,  ..., 1., 1., 1.],
         [1., 1., 1.,  ..., 1., 1., 1.],
         [1., 1., 1.,  ..., 1., 1., 1.]],

        [[1., 1., 1.,  ..., 1., 1., 1.],
         [1., 1., 1.,  ..., 1., 1., 1.],
         [1., 1., 1.,  ..., 1., 1., 1.],
         ...,
         [1., 1., 1.,  ..., 1., 1., 1.],
         [1., 1., 1.,  ..., 1., 1., 1.],
         [1., 1., 1.,  ..., 1., 1., 1.]],

        ...,

        [[1., 1., 1.,  ..., 1., 1., 1.],
         [1., 1., 1.,  ..., 1., 1., 1.],
         [1., 1., 1.,  ..., 1., 1., 1.],
         ...,
         [1., 1., 1.,  ..., 1., 1., 1.],
         [1., 1., 1.,  ..., 1., 1., 1.],
         [1., 1., 1.,  ..., 1., 1., 1.]],

        [[1., 1., 1.,  ..., 1., 1., 1.],
         [1., 1., 1.,  ..., 1., 1., 1.],
         [1., 1., 1.,  ..., 1., 1., 1.],
         ...,
         [1., 1., 1.,  ..., 1., 1., 1.],
         [1., 1., 1.,  ..., 1., 1., 1.],
         [1., 1., 1.,  ..., 1., 1., 1.]],

        [[1., 1., 1.,  ..., 1., 1., 1.],
         [1., 1., 1.,  ..., 1., 1., 1.],
         [1., 1., 1.,  ..., 1., 1., 1.],
         ...,
         [1., 1., 1.,  ..., 1., 1., 1.],
         [1., 1., 1.,  ..., 1., 1., 1.],
         [1., 1., 1.,  ..., 1., 1., 1.]]], device='cuda:1')
Rank 1 sending to Rank 2
[W ProcessGroupNCCL.cpp:1856] Warning: 0NCCL_AVOID_RECORD_STREAMS=1 has no effect for point-to-point collectives. (function operator())
Rank 1 sent data to Rank 2
Process 1 done
[2024-07-15 22:55:17,999] torch.distributed.elastic.agent.server.api: [ERROR] Error waiting on exit barrier. Elapsed: 300.0857963562012 seconds
```


node1:
```
/root/anaconda3/envs/gees_pytorch_3.10/lib/python3.10/site-packages/torch/distributed/launch.py:181: FutureWarning: The module torch.distributed.launch is deprecated
and will be removed in future. Use torchrun.
Note that --use-env is set by default in torchrun.
If your script expects `--local-rank` argument to be set, please
change it to read from `os.environ['LOCAL_RANK']` instead. See 
https://pytorch.org/docs/stable/distributed.html#launch-utility for 
further instructions

  warnings.warn(
[2024-07-15 22:46:39,966] torch.distributed.run: [WARNING] 
[2024-07-15 22:46:39,966] torch.distributed.run: [WARNING] *****************************************
[2024-07-15 22:46:39,966] torch.distributed.run: [WARNING] Setting OMP_NUM_THREADS environment variable for each process to be 1 in default, to avoid your system being overloaded, please further tune the variable for optimal performance in your application as needed. 
[2024-07-15 22:46:39,966] torch.distributed.run: [WARNING] *****************************************
Namespace(distributed_backend='nccl', local_rank=1, rank=3, world_size=4)
Namespace(distributed_backend='nccl', local_rank=0, rank=2, world_size=4)
barrier done
barrier done
Rank 2 waiting to receive from Rank 1
[W ProcessGroupNCCL.cpp:1856] Warning: 0NCCL_AVOID_RECORD_STREAMS=1 has no effect for point-to-point collectives. (function operator())
Rank 3 waiting to receive from Rank 2
[W ProcessGroupNCCL.cpp:1856] Warning: 0NCCL_AVOID_RECORD_STREAMS=1 has no effect for point-to-point collectives. (function operator())
Rank 2 received data from Rank 1
Traceback (most recent call last):
  File "/root/chj/gees/GeeSibling/examples/pytorch/pipeline/test_p2p/train_llama2.py", line 85, in <module>
    main()
  File "/root/chj/gees/GeeSibling/examples/pytorch/pipeline/test_p2p/train_llama2.py", line 70, in main
    dist.recv(tensor=input_tensor, src=rank-1)  # 改为同步recv
  File "/root/anaconda3/envs/gees_pytorch_3.10/lib/python3.10/site-packages/torch/distributed/c10d_logger.py", line 47, in wrapper
    return func(*args, **kwargs)
  File "/root/anaconda3/envs/gees_pytorch_3.10/lib/python3.10/site-packages/torch/distributed/distributed_c10d.py", line 1640, in recv
    pg.recv([tensor], src, tag).wait()
RuntimeError: [3] is setting up NCCL communicator and retrieving ncclUniqueId from [0] via c10d key-value store by key '2:3', but store->get('2:3') got error: Connection reset by peer. This may indicate a possible application crash on rank 0 or a network set up issue.
[2024-07-15 22:52:05,279] torch.distributed.elastic.multiprocessing.api: [WARNING] Sending process 955437 closing signal SIGTERM
[2024-07-15 22:52:08,047] torch.distributed.elastic.multiprocessing.api: [ERROR] failed (exitcode: 1) local_rank: 1 (pid: 955438) of binary: /root/anaconda3/envs/gees_pytorch_3.10/bin/python
Traceback (most recent call last):
  File "/root/anaconda3/envs/gees_pytorch_3.10/lib/python3.10/runpy.py", line 196, in _run_module_as_main
    return _run_code(code, main_globals, None,
  File "/root/anaconda3/envs/gees_pytorch_3.10/lib/python3.10/runpy.py", line 86, in _run_code
    exec(code, run_globals)
  File "/root/anaconda3/envs/gees_pytorch_3.10/lib/python3.10/site-packages/torch/distributed/launch.py", line 196, in <module>
    main()
  File "/root/anaconda3/envs/gees_pytorch_3.10/lib/python3.10/site-packages/torch/distributed/launch.py", line 192, in main
    launch(args)
  File "/root/anaconda3/envs/gees_pytorch_3.10/lib/python3.10/site-packages/torch/distributed/launch.py", line 177, in launch
    run(args)
  File "/root/anaconda3/envs/gees_pytorch_3.10/lib/python3.10/site-packages/torch/distributed/run.py", line 797, in run
    elastic_launch(
  File "/root/anaconda3/envs/gees_pytorch_3.10/lib/python3.10/site-packages/torch/distributed/launcher/api.py", line 134, in __call__
    return launch_agent(self._config, self._entrypoint, list(args))
  File "/root/anaconda3/envs/gees_pytorch_3.10/lib/python3.10/site-packages/torch/distributed/launcher/api.py", line 264, in launch_agent
    raise ChildFailedError(
torch.distributed.elastic.multiprocessing.errors.ChildFailedError:
```

### Versions

```
PyTorch version: 2.1.2+cu121
Is debug build: False
CUDA used to build PyTorch: 12.1
ROCM used to build PyTorch: N/A

OS: Ubuntu 22.04 LTS (x86_64)
GCC version: (Ubuntu 11.4.0-1ubuntu1~22.04) 11.4.0
Clang version: Could not collect
CMake version: version 3.22.1
Libc version: glibc-2.35

Python version: 3.10.6 | packaged by conda-forge | (main, Aug 22 2022, 20:36:39) [GCC 10.4.0] (64-bit runtime)
Python platform: Linux-5.15.0-88-generic-x86_64-with-glibc2.35
Is CUDA available: True
CUDA runtime version: 12.4.131
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: 
GPU 0: NVIDIA L20
GPU 1: NVIDIA L20
GPU 2: NVIDIA L20
GPU 3: NVIDIA L20
GPU 4: NVIDIA L20
GPU 5: NVIDIA L20
GPU 6: NVIDIA L20
GPU 7: NVIDIA L20

Nvidia driver version: 535.171.04
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
Address sizes:                      52 bits physical, 57 bits virtual
Byte Order:                         Little Endian
CPU(s):                             192
On-line CPU(s) list:                0-191
Vendor ID:                          GenuineIntel
Model name:                         Intel(R) Xeon(R) Platinum 8463B
CPU family:                         6
Model:                              143
Thread(s) per core:                 2
Core(s) per socket:                 48
Socket(s):                          2
Stepping:                           8
Frequency boost:                    enabled
CPU max MHz:                        2601.0000
CPU min MHz:                        800.0000
BogoMIPS:                           5200.00
Flags:                              fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush dts acpi mmx fxsr sse sse2 ss ht tm pbe syscall nx pdpe1gb rdtscp lm constant_tsc art arch_perfmon pebs bts rep_good nopl xtopology nonstop_tsc cpuid aperfmperf tsc_known_freq pni pclmulqdq dtes64 ds_cpl vmx smx est tm2 ssse3 sdbg fma cx16 xtpr pdcm pcid dca sse4_1 sse4_2 x2apic movbe popcnt tsc_deadline_timer aes xsave avx f16c rdrand lahf_lm abm 3dnowprefetch cpuid_fault epb cat_l3 cat_l2 cdp_l3 invpcid_single intel_ppin cdp_l2 ssbd mba ibrs ibpb stibp ibrs_enhanced tpr_shadow vnmi flexpriority ept vpid ept_ad fsgsbase tsc_adjust bmi1 avx2 smep bmi2 erms invpcid cqm rdt_a avx512f avx512dq rdseed adx smap avx512ifma clflushopt clwb intel_pt avx512cd sha_ni avx512bw avx512vl xsaveopt xsavec xgetbv1 xsaves cqm_llc cqm_occup_llc cqm_mbm_total cqm_mbm_local split_lock_detect avx_vnni avx512_bf16 wbnoinvd dtherm ida arat pln pts avx512vbmi umip pku ospke waitpkg avx512_vbmi2 gfni vaes vpclmulqdq avx512_vnni avx512_bitalg tme avx512_vpopcntdq la57 rdpid bus_lock_detect cldemote movdiri movdir64b enqcmd fsrm md_clear serialize tsxldtrk pconfig arch_lbr amx_bf16 avx512_fp16 amx_tile amx_int8 flush_l1d arch_capabilities
Virtualization:                     VT-x
L1d cache:                          4.5 MiB (96 instances)
L1i cache:                          3 MiB (96 instances)
L2 cache:                           192 MiB (96 instances)
L3 cache:                           210 MiB (2 instances)
NUMA node(s):                       2
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: 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; Enhanced IBRS, IBPB conditional, RSB filling, PBRSB-eIBRS SW sequence
Vulnerability Srbds:                Not affected
Vulnerability Tsx async abort:      Not affected

Versions of relevant libraries:
[pip3] numpy==1.26.4
[pip3] torch==2.1.2
[pip3] torchaudio==2.1.2
[pip3] torchdata==0.7.1
[pip3] torchvision==0.16.2
[pip3] triton==2.1.0
[conda] numpy                     1.26.4                    <pip>
[conda] torch                     2.1.2                     <pip>
[conda] torchaudio                2.1.2                     <pip>
[conda] torchdata                 0.7.1                     <pip>
[conda] torchvision               0.16.2                    <pip>
[conda] triton                    2.1.0                     <pip>
```