DanLing NestedTensor: 가변 크기 텐서 구조화 기술
DanLing NestedTensor: Composable Multi-Ragged Tensors for Deep Learning
arXiv딥러닝에서 가변 크기 입력 처리 시 발생하는 패딩 및 메모리 낭비 문제를 해결하는 PyTorch 기반의 새로운 텐서 추상화입니다.
- 이 텐서는 데이터의 논리적 차원 순서와 분할 정보를 유지하며, 브로드캐스팅이나 변환 같은 연산도 패딩 없이 효율적으로 수행 가능합니다.
- 실제 테스트 결과, 기존 패딩 방식 대비 최대 3배 이상의 속도 향상과 메모리 절감을 입증하여 성능 최적화에 크게 기여합니다.
- 개발자는 복잡한 오프셋 관리 없이 가변 크기 모델을 구현할 수 있으나, 현재는 논문 출판 후 공개 예정입니다.
딥러닝에서 가변 크기 입력(variable-size inputs)은 흔하지만, 기존의 밀집 배치 방식은 공유된 영역을 할당하고 패딩에 계산 자원을 소모합니다. DanLing NestedTensor는 PyTorch 텐서 추상화로 다중 불규칙 구조(multi-ragged structure)를 텐서 자체의 속성으로 구현했습니다. 이 방식으로 인해, 기존에는 복잡했던 가변 크기 연산이 오프셋 관리 없이 효율적으로 구성될 수 있습니다.
새로운 텐서는 패킹된 값에 텐서 기반 파티션과 논리적 차원 순서를 포함합니다. 따라서 브로드캐스팅을 통해 불규칙한 축을 생성하거나, 특징 변환 및 축소 연산을 수행할 때 이 구조를 유지하며 처리할 수 있습니다. 이러한 표현 방식은 자동 미분(autograd) 및 eager/compiled 실행 환경 모두에서 지원됩니다.
실제 성능 테스트 결과에 따르면, A100 에서 해당 텐서는 동일 모드 패딩 대비 BERT 스케일 네 가지를 거쳐 eager 실행에서 기하 평균 속도 향상 2.74배, 컴파일된 실행에서 3.39배의 성능을 보였습니다. 또한, Pairformer 스타일 워크로드는 패딩 참조 방식보다 최대 4.32배 빠르게 작동했으며, 최고 할당 메모리는 고변동 배치에서 38.08 GiB에서 5.41 GiB까지 감소했습니다.
용어 풀이
- GPU
- 많은 계산을 한꺼번에 처리하는 칩. AI 학습과 실행에 가장 많이 써요.