[vllm] vLLM에서 Qwen MoE 모델의 All-Reduce를 Reduce-Scatter로 최적화하기

PR 링크: vllm-project/vllm#47006 상태: Merged | 변경: +113 / -11

들어가며

vLLM에서 대규모 언어 모델(LLM), 특히 Mixture-of-Experts(MoE) 구조를 가진 모델을 서빙할 때 통신 비용은 성능의 병목이 되곤 합니다. 기존 구현에서는 Attention 레이어의 o_proj 이후 all_reduce를 수행하고, 바로 이어지는 MoE 레이어에서 다시 시퀀스 차원으로 데이터를 청크(chunk)하는 과정이 있었습니다. 이 PR은 이 비효율적인 통신 패턴을 reduce_scatter로 통합하여 불필요한 all_gather와 데이터 복사 과정을 제거함으로써 성능을 최적화합니다.

코드 분석

1. Qwen3NextAttention 및 MLP 레이어의 통신 최적화 (vllm/model_executor/models/qwen3_next.py)

핵심 변경 사항은 Attention 출력과 MoE 입력 사이의 통신을 최적화하는 것입니다. 기존에는 all_reduce를 수행한 후 다시 시퀀스 병렬화를 위해 데이터를 나누었으나, 이를 reduce_scatter로 한 번에 처리합니다.

# Before: 별도의 all_reduce 및 시퀀스 청킹
# After: reduce_scatter를 통한 통합 최적화
if self.use_attn_reduce_scatter_for_moe:
    tp_world_size = get_tensor_model_parallel_world_size()
    sp_pad = (-hidden_states.shape[0]) % tp_world_size
    hidden_states = torch.nn.functional.pad(hidden_states, (0, 0, 0, sp_pad))
    hidden_states = tensor_model_parallel_reduce_scatter(hidden_states, 0)
    if not input_is_sequence_parallel:
        residual = sequence_parallel_chunk(residual)

또한, MoE 레이어의 forward 함수에서 already_sequence_parallel 플래그를 도입하여, 이미 분산된 상태의 데이터를 불필요하게 다시 모으지 않도록 로직을 개선했습니다.

2. 레이어 간 통신 제어 (vllm/model_executor/models/qwen3_5.py)

모델 설정 단계에서 use_attn_reduce_scatter_for_moe 플래그를 통해 MoE 레이어 여부를 판단하고, reduce_results 파라미터를 동적으로 제어하여 통신 방식을 최적화합니다.

# Qwen3.5 모델 초기화 시 MoE 레이어 여부 판단
is_moe_layer = config.model_type == "qwen3_5_moe_text"
self.use_attn_reduce_scatter_for_moe = (
    parallel_config.use_sequence_parallel_moe
    and parallel_config.pipeline_parallel_size == 1
    and is_moe_layer
)

왜 이게 좋은가

이 최적화는 두 가지 측면에서 성능 향상을 가져옵니다:

1. 통신 오버헤드 감소: all_reduce는 reduce_scatter와 all_gather의 결합입니다. 기존에는 all_reduce 후 MoE에서 다시 데이터를 쪼개는 과정이 있었는데, 이를 reduce_scatter 하나로 통합하여 통신 횟수를 줄였습니다.
2. 지연 시간 개선: MoE 연산이 끝난 직후 all_gather를 수행하지 않고, 다음 레이어의 Attention 연산 직전에만 필요한 만큼 데이터를 모으도록 하여, 각 TP(Tensor Parallel) 랭크가 독립적으로 RMSNorm을 수행할 수 있게 되었습니다.

벤치마크 결과, 전체 토큰 처리량(Total token throughput)이 약 2.3% 향상되었으며, P99 TTFT(Time to First Token) 또한 개선되는 효과를 보였습니다.

리뷰어 피드백 반영

리뷰 과정에서 torch.compile과 torch inductor가 패딩된 텐서를 처리할 때 발생할 수 있는 이슈가 논의되었습니다. PR 작성자는 reduce_scatter 연산 시 텐서 크기가 월드 사이즈로 나누어떨어지지 않을 경우를 대비해 수동 패딩을 적용하였고, 이를 통해 컴파일 최적화와 안정성을 동시에 확보했습니다.

결론

이번 최적화는 분산 환경에서 MoE 모델을 서빙할 때 발생하는 통신 병목을 효과적으로 해결한 사례입니다. 특히 reduce_scatter를 활용한 통신 통합은 대규모 모델 서빙의 효율성을 극대화하는 중요한 기법입니다.

참고 자료

• https://pytorch.org/docs/stable/distributed.html
• https://docs.vllm.ai/en/latest/models/parallelism.html

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