본문으로 건너뛰기

[ray] Ray RDT NIXL 메모리 풀 최적화: 불필요한 복사 제거와 전송 효율 극대화

PR 링크: ray-project/ray#65205 상태: Merged | 변경: +959 / -577

들어가며

Ray의 RDT(Remote Data Transfer) 프레임워크 내 NIXL 메모리 풀은 분산 환경에서 텐서 데이터를 효율적으로 전송하기 위한 핵심 컴포넌트입니다. 하지만 기존 구현체에는 두 가지 치명적인 성능 병목이 있었습니다.

  1. 과도한 메모리 복사: 텐서의 특정 부분(View)만 전송하려 해도 해당 텐서가 참조하는 전체 storage를 복사했습니다.
  2. 전송 오버헤드: 여러 텐서가 메모리상에서 연속되어 있더라도 각각 별도의 전송 디스크립터(Transfer Descriptor)를 생성하여, 텐서 개수가 늘어날수록 전송 속도가 선형적으로 느려졌습니다.

이번 PR은 이러한 문제를 해결하기 위해 메모리 풀의 할당 로직을 재설계하고, 텐서 패킹(Packing) 및 디스크립터 병합 알고리즘을 도입했습니다. 실제 코드 변경 사항을 통해 무엇이 개선되었는지 살펴보겠습니다.

코드 분석: 핵심 변경 사항

1. 전체 스토리지 복사에서 뷰(View) 기반 복사로 전환

기존에는 텐서의 실제 데이터 크기와 상관없이 전체 스토리지를 복사하는 비효율이 있었습니다. 변경 후에는 텐서의 numel과 element_size를 곱해 실제 필요한 바이트만 계산하여 복사합니다.

Before (기존 로직 추정):

# 기존에는 스토리지 전체 크기를 할당 사이즈로 사용
alloc_sizes.append(tensor.untyped_storage().nbytes())

After (개선된 로직):

# python/ray/experimental/rdt/nixl_memory_pool.py
tensor_layouts = [
    TensorLayout(t.numel() * t.element_size(), t.element_size())
    for t in tensors
]
sizes = [layout.nbytes for layout in tensor_layouts]

이 변경을 통해 메모리 사용량을 최적화하고, 불필요한 데이터 이동을 방지했습니다.

2. 데이터 정렬(Alignment) 및 패킹 로직 도입

여러 텐서를 하나의 연속된 메모리 블록에 넣기 위해서는 각 데이터 타입(dtype)에 맞는 정렬이 필요합니다. 예를 들어 complex128은 16바이트 정렬이 필요합니다. 이를 위해 _align_up 함수와 packed_offsets 함수가 도입되었습니다.

New Code:

def _align_up(value: int, alignment: int) -> int:
    return (value + alignment - 1) // alignment * alignment

def packed_offsets(tensor_layouts: Sequence[TensorLayout]) -> Tuple[List[int], int]:
    offsets: List[int] = []
    byte_index = 0
    for nbytes, alignment in tensor_layouts:
        byte_index = _align_up(byte_index, alignment)
        offsets.append(byte_index)
        byte_index += nbytes
    return offsets, byte_index

이 로직은 송신측과 수신측이 동일한 레이아웃을 공유할 수 있게 하여, 수신측에서 하나의 커다란 버퍼를 받은 뒤 즉시 여러 개의 텐서 뷰로 재구성할 수 있게 합니다.

3. 전송 디스크립터 병합 및 시간 복잡도 개선

가장 극적인 성능 향상은 전송 디스크립터를 관리하는 방식에서 왔습니다. 기존에는 텐서마다 디스크립터를 생성했지만, 이제는 연속된 메모리 블록을 사용하는 텐서들을 하나의 디스크립터로 묶습니다.

또한, 리뷰 과정에서 저자(Sparks0219)는 이전 구현의 시간 복잡도가 $O(N^2)$이었음을 언급했습니다. 오프셋을 매번 재계산하던 로직을 선형 시간($O(N)$)으로 최적화하여, 10,000개의 뷰를 처리할 때 200ms에서 2ms로 성능을 100배 개선했습니다.

Reviewer Feedback 반영: Stephanie Wang의 제안에 따라 desc_lens라는 모호한 변수명 대신 packed_tensor_sizes와 같은 명확한 명칭을 사용하고, NamedTuple을 도입하여 가독성을 높였습니다.

class TensorLayout(NamedTuple):
    nbytes: int
    alignment: int

왜 이게 좋은가?

  1. 성능의 비약적 향상: 텐서 개수가 많은 워크로드(예: 대규모 모델의 파라미터 전송)에서 메타데이터 처리 비용을 $O(N^2)$에서 $O(N)$으로 줄인 것은 시스템 전체의 처리량(Throughput)에 큰 영향을 미칩니다.
  2. 메모리 효율성: 텐서 뷰만 복사함으로써 GPU/시스템 메모리 파편화를 줄이고 더 많은 데이터를 풀에 수용할 수 있습니다.
  3. 프로토콜의 견고함: _align_up을 통한 엄격한 바이트 정렬은 PyTorch의 Tensor.view(dtype) 호출 시 발생할 수 있는 정렬 오류를 원천 차단합니다.

결론

이번 최적화는 단순히 코드를 깔끔하게 만드는 것을 넘어, 분산 시스템에서 데이터 전송 시 발생하는 전형적인 병목 지점들을 정밀하게 타격했습니다. 특히 대규모 텐서 뷰를 다룰 때의 성능 향상은 Ray RDT를 사용하는 사용자들에게 즉각적인 이득을 제공할 것입니다.

시니어 엔지니어로서 배울 수 있는 교훈은 "추상화된 데이터 구조(Tensor View) 뒤에 숨겨진 실제 메모리 레이아웃(Storage)을 이해하고, 메타데이터 처리 로직의 시간 복잡도를 항상 경계해야 한다"는 점입니다.

참고 자료

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

댓글

관련 포스트

PR Analysis 의 다른글