본문으로 건너뛰기

[vllm] [vLLM 성능 최적화] Nemotron-3 디코딩 속도를 13% 향상시킨 Latent-MoE All-Reduce 최적화 분석

PR 링크: vllm-project/vllm#52301 상태: Merged | 변경: +21 / -0

들어가며

최근 vLLM 프로젝트에서 Nemotron-3 모델의 디코딩 성능이 v0.26.0 대비 v0.27.x 버전에서 약 14% 하락하는 회귀(Regression) 현상이 발견되었습니다. 특히 이 문제는 Tensor Parallelism(TP)이 1보다 큰 분산 환경에서 두드러졌습니다.

원인은 v0.27.0에서 Latent-MoE(Mixture of Experts)를 지원하기 위해 추가된 로직 때문이었습니다. Latent-MoE는 전문가(Expert)의 출력을 결합하기 전에 특정 변환(Transform)을 거치는데, 이 과정에서 데이터의 일관성을 위해 All-Reduce 통신이 발생합니다. 하지만 Nemotron-3와 같은 특정 모델에서는 이 통신이 수학적으로 중복될 수 있다는 점이 발견되었습니다.

본 글에서는 불필요한 All-Reduce를 제거하여 성능을 원복시킨 PR의 핵심 변경 사항과 그 이면에 숨겨진 수학적 원리를 분석합니다.

코드 분석: 무엇이 바뀌었는가?

1. moe_runner.py: 조건부 All-Reduce 스킵 로직 도입

기존 코드에서는 routed_output_transform이 존재하면 무조건 변환 전에 All-Reduce를 수행했습니다. 하지만 이번 PR에서는 reduce_commutative라는 플래그를 도입하여, 변환 연산이 합산(Sum) 연산과 교환 법칙이 성립할 경우 이 과정을 건너뛸 수 있게 했습니다.

[Before & After]

# vllm/model_executor/layers/fused_moe/runner/moe_runner.py

# Before: transform이 있으면 무조건 All-Reduce 시도 가능성 존재
if (
    self.routed_output_transform is not None
    and not self.moe_config.is_sequence_parallel
    and (self.moe_config.tp_size > 1 or self.moe_config.ep_size > 1)
    and not fused_output_is_reduced
):
    # ... All-Reduce 수행 ...

# After: 'reduce_commutative' 속성을 체크하여 최적화 여부 결정
if (
    self.routed_output_transform is not None
    and not getattr(self.routed_output_transform, "reduce_commutative", False) # 추가된 조건
    and not self.moe_config.is_sequence_parallel
    and (self.moe_config.tp_size > 1 or self.moe_config.ep_size > 1)
    and not fused_output_is_reduced
):
    # ... All-Reduce 수행 ...

이 변경을 통해 reduce_commutative = True로 설정된 변환 레이어는 첫 번째 All-Reduce를 건너뛰고, 나중에 발생하는 최종 All-Reduce 한 번으로 통신을 통합할 수 있게 되었습니다.

2. nemotron_h.py: Nemotron 모델의 최적화 활성화

Nemotron 모델 정의 파일에서는 fc2_latent_proj 레이어가 선형(Linear) 연산이면서 바이어스(Bias)가 없고 양자화되지 않은 경우, 이 교환 법칙이 성립함을 명시합니다.

[Before & After]

# vllm/model_executor/models/nemotron_h.py

# Before: 단순히 레이어만 생성
self.fc2_latent_proj = RowParallelLinear(
    # ... parameters ...
)

# After: 교환 법칙 성립 여부를 판단하여 플래그 설정
self.fc2_latent_proj = RowParallelLinear(
    # ... parameters ...
)

# bias가 없고 양자화되지 않은 Linear 연산은 sum_r(W * x_r) == W * sum_r(x_r) 이 성립함
self.fc2_latent_proj.reduce_commutative = (
    not config.mlp_bias
    and isinstance(
        self.fc2_latent_proj.quant_method, UnquantizedLinearMethod
    )
)

여기서 흥미로운 점은 모델 전체의 quant_config를 보지 않고, 개별 레이어의 quant_method를 직접 확인한다는 것입니다. 이는 ModelOpt와 같은 도구가 Latent Projection 레이어만 양자화에서 제외하는 경우가 있기 때문인데, 이를 통해 양자화된 체크포인트에서도 최대한의 성능 이득을 얻을 수 있도록 설계되었습니다.

왜 이게 좋은 최적화인가?

1. 수학적 정당성: 선형성의 활용

이 최적화의 핵심은 선형 대수의 기본 원리인 선형성(Linearity)에 있습니다.

  • 기존 방식 (2회 통신): AllReduce(x_1, x_2, ...) -> Transform(Sum(x_i)) -> AllReduce(Final_Output)
  • 최적화 방식 (1회 통신): Transform(x_i) -> AllReduce(Sum(Transform(x_i)))

만약 Transform(T)가 선형 연산($T(∑x) = ∑T(x)$)이라면 두 결과는 동일합니다. Nemotron의 fc2_latent_proj는 Bias가 없는 Linear 레이어이므로 이 조건이 완벽히 충족됩니다. 결과적으로 동일한 결과값을 보장하면서 비싼 네트워크 통신 비용을 절반으로 줄인 것입니다.

2. 압도적인 성능 향상

PR에 첨부된 벤치마크 결과에 따르면, TP4 환경에서 Inter-Token Latency(ITL)가 4.64ms에서 4.08ms로 약 13.8% 개선되었습니다. 이는 v0.26.0 수준의 성능을 완벽하게 복구한 수치입니다.

특히 배치가 작은(Concurrency=1) 상황에서는 통신 횟수 자체가 지연 시간(Latency)의 병목이 되는데, 이 PR은 통신 횟수를 2회에서 1회로 줄임으로써 이 문제를 정면으로 해결했습니다.

3. 안전한 설계 (Opt-in 방식)

모든 모델에 이 로직을 강제 적용하는 대신, reduce_commutative라는 속성을 명시적으로 설정한 모델만 최적화가 적용되도록 했습니다. 예를 들어 RMSNorm과 같은 비선형 연산을 포함하는 Kimi K3 모델은 이 플래그를 설정하지 않으므로 기존의 안전한(하지만 느린) 경로를 그대로 유지합니다. 이는 대규모 오픈소스 프로젝트에서 기존 모델의 정확도를 깨뜨리지 않으면서 특정 모델의 성능을 올리는 아주 모범적인 방식입니다.

결론

이번 최적화는 분산 학습 및 추론에서 통신 오버헤드가 얼마나 치명적인지, 그리고 모델의 수학적 구조를 깊이 이해했을 때 어떤 성능 이득을 얻을 수 있는지를 잘 보여줍니다.

시니어 엔지니어로서 우리는 단순히 코드를 작성하는 것을 넘어, 우리가 다루는 연산의 수학적 성질(교환 법칙, 결합 법칙 등)을 활용해 시스템의 물리적 한계(네트워크 대역폭 및 지연 시간)를 극복할 수 있는 방법을 항상 고민해야 합니다.

참고 자료

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

댓글

관련 포스트

PR Analysis 의 다른글