NCCL

tree all-reduce. N 개 GPU 를 binary tree 로 — log(N) 단계의 reduce 후 같은 단계만큼 broadcast.

• step 수: 2 log(N). ring 의 2(N-1) 보다 훨씬 적음. N=64 면 ring 은 126 step, tree 는 12 step.
• 그러나 — 각 link 에서 흐르는 데이터가 ring 보다 N× 많다. 큰 message 에서는 그 cost 가 ring 의 step 수 cost 를 초과.

NCCL 은 message size 에 따라 자동 선택. 작으면 tree, 크면 ring. cross-over 점은 보통 수 MB 근처.

double-binary tree — tree 의 한 link 가 idle 한 시간을 reverse direction 의 다른 tree 가 메우는 변형. 두 개의 binary tree 가 동시에 도는데 root 가 다르고 — link 사용률이 거의 100%.

NCCL 의 default. nsys 에서 NCCL kernel 이름이 ncclTreeAllReduce 또는 ncclRingAllReduce 로 뜸 — 어떤 algorithm 이 dispatch 되었는지 정보.

NCCL 의 처방.

• hierarchical ring — node 안에서 partial reduce, node 간에 한 번 inter-node, 다시 node 안에서 broadcast. 각 단계가 자기 layer 의 가장 적합한 algorithm.
• topology auto-detect — NCCL_DEBUG=INFO 로 NCCL 이 자기가 어떤 ring 을 그렸는지 출력. 잘못 그렸으면 환경변수로 hint.

device 측에서 일어나는 일.

1. NCCL kernel 이 launch 되면 — 일부 SM 에 떠서 ring 안 GPU 들과 send/recv 시작.
2. 그 동안 backward 의 GEMM 이 돌고 있다면 — 두 kernel 이 SM 을 분배해서 도는 중.
3. NCCL kernel 의 SM 수가 너무 많으면 GEMM throughput 이 떨어짐. 너무 적으면 NCCL 이 link 을 다 못 채움.

NCCL 이 사용하는 SM 수는 자동 결정 — 보통 큰 message 에서 더 많은 SM. 이 자동 결정이 가끔 backward GEMM 과 안 맞아서 손해가 나는 경우가 있다.

# ddp_simple.py — 강의 repo 그대로
import torch.distributed as dist
from torch.nn.parallel import \
    DistributedDataParallel as DDP

class ToyModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.w = nn.Parameter(torch.tensor(5.0))
    def forward(self, x):
        return self.w * 7.0 * x

def demo_basic():
    dist.init_process_group("nccl")
    rank = dist.get_rank()
    model = ToyModel().to(rank)
    ddp_model = DDP(model, device_ids=[rank])

    with profile() as prof:
        x = torch.tensor(rank, dtype=torch.float)
        y = ddp_model(x)
        y.backward()
        # 여기 ↑ 안에 ncclAllReduce 가 자동 발생
    if rank == 0:
        prof.export_chrome_trace("trace_ddp_simple.json")

실행: torchrun --nproc_per_node=2 ddp_simple.py.

trace 위에서 보이는 것.

• backward 안에 autograd::CppNode 가 gradient 를 만들면서 — Reducer::push_rebuilt_params_for_all_indices 같은 PyTorch 내부 hook 이 발동.
• 그 hook 이 ProcessGroupNCCL::allreduce 호출.
• 그게 device 에 ncclAllReduce kernel 을 launch.
• flow event 가 CPU 의 backward → CPU 의 nccl 호출 → GPU 의 nccl kernel 의 chain 을 그린다.

bucket gradient — DDP 가 모든 gradient 를 하나하나 all-reduce 하지 않고 N MB 단위 bucket 으로 묶어 한 번에. DDP(model, bucket_cap_mb=25) 의 25 MB 가 그 단위. 강의의 ddp_example.py 가 이 패턴.