[flashinfer] FlashInfer 분산 오토튜닝 동기화: NCCL 데드락 해결을 위한 전략적 접근

PR 링크: flashinfer-ai/flashinfer#3187 상태: Merged | 변경: +435 / -6

들어가며

대규모 언어 모델(LLM)을 분산 환경(TP, Tensor Parallelism)에서 실행할 때, FlashInfer와 같은 커널 라이브러리는 최적의 성능을 내기 위해 autotuner를 사용하여 커널 실행 전략(tactic)을 결정합니다. 하지만 여러 GPU가 동시에 오토튜닝을 수행할 때, 각 GPU의 미세한 타이밍 차이로 인해 서로 다른 tactic을 선택하게 되는 문제가 발생합니다. 특히 NCCL_WIN_COLL_SYMMETRIC과 같은 대칭 메모리 할당을 사용하는 경우, 각 랭크(rank)가 서로 다른 크기의 스크래치 버퍼를 할당하려 시도하면서 NCCL 데드락이 발생합니다. 본 PR은 이 문제를 해결하기 위해 오토튜닝 과정에서 랭크 간 tactic 선택을 동기화하는 메커니즘을 도입했습니다.

코드 분석

1. flashinfer/autotuner/autotuner.py: 동기화 로직 추가

핵심은 _profile_single_kernel 내에서 각 랭크가 측정한 커널 실행 시간을 all_reduce를 통해 평균화하는 것입니다. 이를 통해 모든 랭크가 동일한 평균 시간을 참조하여 동일한 tactic을 선택하게 합니다.

Before:

# 기존에는 각 랭크가 독립적으로 시간을 측정함
avg_time = pure_profile(stream, self.repeat)

After:

# 동기화 그룹이 설정된 경우, 측정된 시간을 모든 랭크가 공유
if _tune_process_group is not None:
    import torch.distributed as dist
    time_tensor = torch.tensor([avg_time], dtype=torch.float64, device=device)
    dist.all_reduce(time_tensor, op=dist.ReduceOp.SUM, group=_tune_process_group)
    avg_time = time_tensor.item() / dist.get_world_size(_tune_process_group)

또한, 특정 랭크에서 OOM(Out of Memory)이 발생하더라도 전체 랭크가 동기화 루프를 이탈하지 않도록 예외 처리를 강화했습니다. 예외 발생 시 avg_time을 inf로 설정하여 해당 tactic을 배제하고, 모든 랭크가 all_reduce에 참여하도록 강제합니다.

2. flashinfer/autotuner/__init__.py: API 노출

사용자가 명시적으로 set_autotune_process_group을 호출하여 동기화 그룹을 설정할 수 있도록 API를 추가했습니다. 이는 옵트인(opt-in) 방식으로, 기존 단일 GPU 사용자에게는 아무런 영향을 주지 않습니다.

왜 이게 좋은가

이 최적화는 단순히 성능을 높이는 것이 아니라, 분산 환경에서의 결정론적 동작(determinism)을 보장하여 시스템의 안정성을 확보합니다.

1. 데드락 방지: 모든 랭크가 동일한 tactic을 선택함으로써, NCCL이 요구하는 대칭적 메모리 할당 조건이 충족됩니다. 실제 Qwen3-235B 모델 테스트에서 기존에 발생하던 CUDA-graph 캡처 시의 데드락이 완벽히 해결되었습니다.
2. 유연한 설계: gloo와 nccl 백엔드를 모두 지원하며, SUM 연산을 사용하여 백엔드 호환성을 높였습니다. 또한, 예외 상황에서도 all_reduce 호출 횟수를 보존하여 랭크 간 동기화가 깨지지 않도록 설계되었습니다.
3. 오버헤드 최소화: CPU 기반의 gloo 그룹을 사용할 경우, 오토튜닝 전체 시간에 미치는 영향은 무시할 수 있을 정도로 작습니다.

교훈

분산 시스템에서 각 노드가 독립적으로 의사결정을 내리는 것은 성능 최적화 측면에서는 유리할 수 있으나, collective 통신이 결합된 환경에서는 심각한 데드락을 유발할 수 있습니다. 이러한 경우, '의사결정의 동기화'가 시스템 전체의 안정성을 위한 필수적인 선결 조건임을 알 수 있습니다.

참고 자료

• https://pytorch.org/docs/stable/distributed.html
• https://docs.nvidia.com/nccl/user-guide/docs/usage/collectives.html

⚠️ 알림: 이 분석은 AI가 실제 코드 diff를 기반으로 작성했습니다.