PyTorch Internals (2019)

• PyTorch의 내부 구조에 대한 설명으로, PyTorch의 C++ 코드베이스에 기여하고자 하는 사람들을 위한 안내서
• 이 글의 목표는 PyTorch의 텐서 라이브러리 구조 및 자동 미분(autograd) 기법을 이해하고 코드베이스에서 길을 찾는 데 도움을 주는 것임

PyTorch 텐서의 기본 구조

• PyTorch에서 텐서는 가장 기본적인 데이터 구조임
• 텐서는 n차원 데이터 구조로, 부동 소수점(float), 정수(int) 등과 같은 스칼라 값 저장 가능
• 텐서는 다음과 같은 메타데이터 포함:
  • 크기(size): 텐서의 차원 정보
  • dtype: 저장된 데이터 타입 (예: float32, int64 등)
  • device: 데이터가 저장된 위치 (CPU, CUDA 등)
  • stride: 데이터의 물리적 메모리에서의 오프셋 정보

• Stride의 역할

  • stride는 논리적 인덱스를 물리적 메모리 위치로 변환하는 데 사용됨
  • stride는 각 차원에 대해 오프셋을 설정하고 인덱스에 stride 값을 곱해 물리적 메모리 위치 결정
  • stride를 통해 새로운 텐서를 생성하지 않고 view 형태로 같은 데이터를 다른 방식으로 볼 수 있음

텐서와 저장소(Storage) 개념

• PyTorch에서는 텐서가 실제 데이터를 직접 보관하지 않음 → 저장소(Storage)에서 데이터를 관리함
• Tensor = 크기 + dtype + device + stride + offset 정보
• 여러 텐서가 하나의 저장소를 공유 가능 → 뷰(View) 개념 지원
• 저장소와 텐서의 분리 덕분에 메모리 효율적으로 사용 가능

텐서 연산의 디스패치(Dispatch) 과정

• PyTorch에서 연산은 두 단계의 디스패치를 거침:
  1. 디바이스 타입 및 레이아웃 기반 디스패치
     • CPU 텐서 vs. CUDA 텐서에 따라 서로 다른 구현 코드 실행
  2. dtype 기반 디스패치
     • float vs. int 등 데이터 타입에 따라 서로 다른 커널 호출

PyTorch 텐서 확장 모델

• 텐서의 세 가지 주요 확장 요소:

  • Device: CPU, GPU, TPU 등에서 메모리 할당 방식 정의
  • Layout: 텐서가 메모리에 저장되는 방식 정의 (예: 연속적 저장, 희소(sparse) 저장 등)
  • dtype: 텐서의 각 요소에 저장될 데이터 유형 정의

• 확장 옵션:

  • PyTorch 코드를 직접 수정해서 텐서 확장 가능
  • 기존 텐서를 감싸는 래퍼 클래스 작성 가능
  • 자동 미분 중에 래퍼가 필요하면 직접 확장 필요

자동 미분 (Autograd) 동작 원리

• PyTorch는 **역전파(reverse-mode differentiation)**를 기반으로 자동 미분 수행
• 순전파(forward) 연산 시 그래프 생성 → 역전파 시 그래프 탐색하며 미분 수행
• Autograd는 다음과 같은 추가 정보를 관리:
  • AutogradMeta: 텐서에 연결된 메타데이터로 역전파에 사용됨
  • 연산 결과를 기록하고 역전파 시에 미분 수행

PyTorch 코드 구조 및 파일 위치

• PyTorch 코드베이스에서 주요 디렉토리:
  • torch/ → Python 모듈 (Python 코드)
  • torch/csrc/ → Python과 C++ 바인딩 코드, 자동 미분 엔진, JIT 컴파일러 등
  • aten/ → 텐서 연산 정의 (핵심 연산 대부분 포함)
  • c10/ → 텐서 및 저장소와 같은 코어 구조체 정의

PyTorch 연산 실행 과정

• 예제: torch.add() 호출 시 실행 과정:
  1. Python에서 C++ 코드로 인수 변환
  2. VariableType에서 디스패치 수행
  3. Device/레이아웃 기반 디스패치 수행
  4. 최종 커널 실행

커널 작성 과정 및 도구

• PyTorch에서 커널은 다음과 같은 단계로 작성됨:
  1. 연산 메타데이터 작성: 함수 시그니처, 지원 디바이스 및 데이터 타입 정의
  2. 입력 검증: 차원, 타입 등 입력 검증 수행
  3. 출력 텐서 할당
  4. dtype 디스패치: 데이터 타입에 따라 커널 실행
  5. 병렬 처리: CPU에서는 OpenMP, CUDA에서는 내장 병렬화 사용
  6. 데이터 접근 및 계산: TensorAccessor, TensorIterator 등 사용

주요 디스패치 매크로

• AT_DISPATCH_ALL_TYPES → dtype에 따른 디스패치 수행
• 다양한 데이터 타입에 대해 매크로 지원 → 성능 최적화 가능

성능 최적화 및 작업 효율 향상 팁

• 헤더 파일 수정 최소화 → 수정 시 전체 코드 리빌드 발생
• 로컬 개발 환경 설정 → CI 사용 시 시간 소모 최소화
• ccache 사용 → 재컴파일 시간 절약 가능
• 강력한 서버 사용 → C++ 컴파일 및 CUDA 빌드 시 시간 단축 가능

PyTorch 기여 가이드

• 시작하기 좋은 기여 대상:
  • triaged 라벨이 달린 이슈 → PyTorch 개발자가 확인 완료한 이슈
  • 문서 개선 및 버그 재현 도움
  • PyTorch의 RFC(기능 제안)에 대한 의견 제시
• PyTorch는 오픈 소스 기여자를 통해 성장해 왔으며, 커뮤니티 참여 환영