본문으로 건너뛰기

[vllm] vLLM Inkling 모델: 인플레이스 연산으로 GPU 메모리 최적화 및 OOM 방지

PR 링크: vllm-project/vllm#49487 상태: Merged | 변경: +2 / -2

들어가며

대규모 언어 모델(LLM)의 추론 성능은 GPU 메모리 효율성과 직결됩니다. 특히 vLLM과 같이 고성능 LLM 서빙을 목표로 하는 프레임워크에서는 미세한 메모리 할당 최적화가 전체 시스템의 안정성과 처리량에 지대한 영향을 미칠 수 있습니다. 오늘 분석할 GitHub PR은 vLLM의 Inkling 모델에서 발생하는 불필요한 임시 텐서 할당을 제거하여, GPU 메모리 부족(OOM) 문제를 해결하고 전반적인 성능을 개선한 사례입니다.

이 PR이 해결하고자 하는 핵심 문제는 Inkling 모델의 특정 연산들이 결과를 저장하기 위해 불필요하게 새로운 텐서를 할당한다는 점입니다. 예를 들어, H100 환경에서 global_scale 적용 시 86 MiB, MoE(Mixture-of-Experts) 출력 합산 시 94 MiB의 임시 텐서가 할당되었습니다. 이러한 작은 할당들이 누적되면, 특히 고부하 환경(예: 8xH100 TP8 8K→1K 고동시성 설정)에서 치명적인 OOM 오류를 유발할 수 있습니다. 이 PR은 이러한 연산들을 인플레이스(in-place) 방식으로 변경하여 메모리 할당을 최소화하고 OOM을 방지하는 것을 목표로 합니다.

코드 분석: Inkling 모델의 메모리 최적화

이번 PR의 핵심은 PyTorch 텐서 연산에서 발생하는 임시 메모리 할당을 제거하기 위해 인플레이스 연산을 도입하는 것입니다. 이는 torch.Tensor 객체의 _ (언더스코어) 접미사가 붙은 메서드들을 활용하여 구현됩니다. 이러한 메서드들은 새로운 텐서를 생성하는 대신, 기존 텐서의 데이터를 직접 수정합니다.

vllm/models/inkling/nvidia/mlp.py 변경

첫 번째 변경사항은 Inkling 모델의 MLP(Multi-Layer Perceptron) 레이어에서 global_scale을 적용하는 부분입니다. 이전 코드에서는 x * self.global_scale 연산이 새로운 텐서를 생성하여 x에 재할당했습니다.

Before:

        x = silu_and_mul_triton(gate_up)
        x, _ = self.down_proj(x)
        if self.global_scale is not None:
            x = x * self.global_scale
        # TP-partial output: the layer's reduce-scatter fallback consumes it.
        return x

After:

        x = silu_and_mul_triton(gate_up)
        x, _ = self.down_proj(x)
        if self.global_scale is not None:
            x.mul_(self.global_scale)
        # TP-partial output: the layer's reduce-scatter fallback consumes it.
        return x

변경 후에는 x.mul_(self.global_scale)을 사용하여 x 텐서 자체의 값을 self.global_scale로 곱하도록 했습니다. 이 mul_ 메서드는 x 텐서의 메모리 공간을 재활용하므로, 86 MiB에 달하는 임시 텐서 할당을 방지할 수 있습니다. down_proj가 새로운 텐서를 반환하고, 이 텐서가 즉시 스케일링되어 반환되기 때문에 이 연산은 안전하게 인플레이스로 수행될 수 있습니다.

vllm/models/inkling/nvidia/moe.py 변경

두 번째 변경사항은 MoE 레이어에서 라우팅된 출력(out)에 싱크(sink) 전문가의 출력(sink_out)을 더하는 부분입니다. 이전 코드에서는 out + sink_out 연산이 두 텐서의 합을 저장할 새로운 텐서를 생성하여 반환했습니다.

Before:

        )
        self._routed_sel = None

        return out + sink_out

    # -- weight loading ----------------------------------------------------

After:

        )
        self._routed_sel = None

        return out.add_(sink_out)

    # -- weight loading ----------------------------------------------------

변경 후에는 out.add_(sink_out)을 사용하여 out 텐서에 sink_out 텐서의 값을 직접 더하도록 했습니다. 이 add_ 메서드 또한 out 텐서의 메모리 공간을 재활용하므로, 94 MiB에 달하는 임시 텐서 할당을 제거합니다. 라우팅된 MoE 출력(out)은 이 연산 이후 더 이상 소비자가 없는 프라이빗 버퍼이며, 모든 라우팅/싱크 스트림 작업이 완료된 후에 덧셈이 이루어지므로 이 역시 안전한 인플레이스 연산입니다.

안전성 고려사항

인플레이스 연산은 메모리 효율적이지만, 원본 텐서의 데이터를 직접 수정하므로 부작용을 일으킬 수 있습니다. 이 PR에서는 다음 사항들을 고려하여 안전성을 확보했습니다:

  • 프라이빗 버퍼: 대상 버퍼(x와 out)는 해당 연산 이후 다른 소비자가 없는 프라이빗 텐서입니다.
  • 추론 모드: vLLM은 추론 모드에서 이 경로를 실행하므로, autograd에 의해 저장된 텐서가 변이될 염려가 없습니다.
  • 데이터 타입 및 형상 보존: 두 연산 모두 원래의 데이터 타입, 형상, 산술적 의미를 보존합니다.

왜 이 최적화가 좋은가?

이 PR은 단순해 보이는 코드 변경이지만, GPU 메모리 관리와 성능 최적화 측면에서 매우 중요한 의미를 가집니다.

OOM 방지 및 안정성 향상

가장 직접적인 이점은 OOM(Out Of Memory) 오류 방지입니다. PR 설명에 따르면, H100 환경에서 부모 커밋(parent commit)은 return out + sink_out 부분에서 94 MiB 할당 시 OOM이 발생했고, x = x * self.global_scale 부분에서도 86 MiB 할당 시 OOM이 발생했습니다. 이 PR은 이 두 가지 임시 할당(총 180 MiB)을 제거함으로써, 동일한 H100 고동시성 벤치마크(concurrency 256 및 512)를 오류 없이 성공적으로 완료했습니다. 이는 특히 제한된 GPU 메모리 환경에서 모델을 안정적으로 운영하는 데 필수적입니다.

성능 및 자원 효율성 개선

OOM 방지 외에도, 불필요한 텐서 할당 및 해제는 GPU 메모리 할당자(allocator)에 부하를 주고, CUDA 커널 실행 시 오버헤드를 증가시킬 수 있습니다. 인플레이스 연산은 이러한 할당자 트래픽을 줄여 전반적인 성능을 향상시키고, GPU 자원을 더욱 효율적으로 사용할 수 있게 합니다. PR의 벤치마크 결과는 OOM 없이 높은 동시성에서 안정적인 출력 처리량(350+ tok/s)을 보여주며, 이는 최적화의 효과를 입증합니다.

일반적 교훈

이 PR은 다음과 같은 중요한 개발 교훈을 제공합니다:

  1. 메모리 할당 최소화의 중요성: GPU 프로그래밍, 특히 딥러닝 모델 추론에서는 작은 임시 메모리 할당이라도 누적되면 OOM의 원인이 될 수 있습니다. 가능한 한 메모리 재사용을 고려해야 합니다.
  2. 인플레이스 연산의 활용: PyTorch와 같은 텐서 라이브러리에서 _ 접미사가 붙은 메서드(예: add_, mul_, copy_)는 인플레이스 연산을 수행하여 새로운 메모리 할당 없이 텐서의 내용을 직접 수정합니다. 이는 메모리 효율성을 극대화하는 강력한 도구입니다.
  3. 안전성 검토: 인플레이스 연산을 사용할 때는 해당 텐서가 다른 곳에서 참조되지 않는지, autograd 그래프에 영향을 주지 않는지 등 안전성을 철저히 검토해야 합니다.
  4. 통합 테스트의 중요성: 리뷰어 mgoin의 댓글에서 유닛 테스트의 필요성에 대한 논의가 있었지만, 이 PR의 경우 OOM 재현 및 해결을 위한 H100 통합 테스트가 핵심적인 검증 역할을 했습니다. 실제 운영 환경과 유사한 조건에서의 테스트는 미세한 성능 및 메모리 문제를 발견하는 데 필수적입니다.

결론

이번 vLLM Inkling 모델의 PR은 코드 몇 줄의 변경만으로 GPU 메모리 사용량을 크게 줄이고 OOM을 방지하여 시스템의 안정성과 성능을 향상시킨 모범적인 사례입니다. 이는 고성능 딥러닝 추론 시스템을 개발할 때, 코드의 미세한 부분까지 메모리 효율성을 고려하는 것이 얼마나 중요한지를 다시 한번 일깨워줍니다. 이러한 최적화는 단순히 버그를 수정하는 것을 넘어, 더 많은 사용자가 더 큰 모델을 더 효율적으로 사용할 수 있도록 하는 기반이 됩니다.

참고 자료

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

댓글

관련 포스트

PR Analysis 의 다른글