텐서 병렬 처리는 단일 신경망 계층의 계산을 여러 하드웨어 장치에 분산하는 기법이다. 이 접근 방식에서 계층의 가중치 행렬은 샤드로 분할되며, 각 장치는 매개변수의 일부를 보유하고 행렬 곱셈의 해당 부분을 수행한다. 그런 다음 부분 결과는 all-reduce와 같은 집합 통신 연산을 통해 결합되어 계층의 출력을 생성한다. 이 방법은 모델 병렬 처리의 한 형태로, 각 장치가 모델의 전체 사본을 보유하고 서로 다른 데이터 샘플을 처리하는 데이터 병렬 처리와 구별된다. 텐서 병렬 처리는 단일 가속기의 메모리 용량을 초과하는 대규모 언어 모델과 같은 매우 큰 모델을 훈련하고 추론을 실행하는 데 필수적이다.
텐서 병렬 처리의 주요 동기는 현대 딥러닝 모델의 메모리 및 계산 요구 사항이다. 예를 들어, 수천억 개의 매개변수를 가진 모델은 일반적으로 16~80GB 범위인 단일 GPU 또는 TPU의 메모리에 맞출 수 없다. 가중치 행렬을 샤딩함으로써 텐서 병렬 처리는 모델을 여러 장치에 분산시켜 더 큰 모델 사용과 더 빠른 훈련 시간을 가능하게 한다. 또한 장치별 메모리 사용량을 줄이고 더 큰 배치 크기를 허용한다. 그러나 텐서 병렬 처리는 장치가 부분 결과를 동기화해야 하므로 통신 오버헤드를 도입하며, 이는 특히 느린 상호 연결에서 병목 현상이 될 수 있다.
텐서 병렬 처리는 일반적으로 파이프라인 병렬 처리 및 데이터 병렬 처리와 같은 다른 병렬 처리 전략과 함께 사용되어 수천 개의 장치에 걸쳐 훈련을 확장한다. 이는 OpenAI, Anthropic, Google DeepMind가 개발한 모델을 포함한 많은 최첨단 모델의 훈련 인프라의 핵심 구성 요소이다. 이 기법은 추론에도 관련이 있으며, 여러 가속기에 계산을 분산하여 대규모 모델을 낮은 지연 시간으로 제공할 수 있게 한다.
역사적 배경
행렬 연산을 여러 프로세서로 분할하는 개념은 초기 병렬 컴퓨팅 연구로 거슬러 올라간다. 1980년대와 1990년대에 연구자들은 신경망의 병렬 구현을 탐구했지만 하드웨어와 소프트웨어 도구는 제한적이었다. 현대적 형태의 텐서 병렬 처리는 2010년대 GPU 기반 딥러닝의 부상과 함께 등장했다. 가장 초기의 주목할 만한 사용 중 하나는 토론토 대학교 및 다른 기관의 연구자들이 모델 병렬 처리를 실험한 대규모 합성곱 신경망 훈련이었다. 그러나 2017년에 도입된 Transformer 아키텍처와 같은 트랜스포머 기반 모델의 출현으로 텐서 병렬 처리가 주류 기법이 되었다. 트랜스포머의 어텐션 계층과 피드포워드 네트워크는 큰 행렬 곱셈으로 구성되어 있어 샤딩의 자연스러운 대상이다.
2019년에 NVIDIA 및 다른 조직의 연구자들은 대규모 언어 모델 훈련을 위한 텐서 병렬 처리에 대한 상세한 설명을 발표했다. NVIDIA가 개발한 Megatron-LM 프레임워크는 트랜스포머 모델에 대한 효율적인 텐서 병렬 처리를 입증하여 여러 GPU에서 거의 선형 확장을 달성했다. 이 작업은 Google의 GShard 및 Microsoft의 DeepSpeed와 같은 후속 대규모 훈련 시스템의 기초를 마련했으며, 이들은 텐서 병렬 처리를 핵심 기능으로 통합한다.
텐서 병렬 처리 작동 방식
일반적인 신경망 계층에서 순방향 전달은 행렬 곱셈 Y = XW를 계산하며, 여기서 X는 입력 활성화, W는 가중치 행렬, Y는 출력이다. 텐서 병렬 처리에서 가중치 행렬 W는 하나 이상의 차원을 따라 분할된다. 완전 연결 계층의 경우 일반적인 접근 방식은 W를 열 방향(즉, 출력 차원을 따라)으로 분할하는 것이다. 각 장치는 W의 열 하위 집합을 보유하고 부분 출력을 계산한다. 그런 다음 부분 출력은 all-reduce 연산을 사용하여 결합되어 최종 Y를 생성한다. 이를 열 병렬 분할이라고 한다. 대안으로 W를 행 방향(입력 차원을 따라)으로 분할할 수 있으며, 입력에 all-gather 연산과 출력에 all-reduce가 필요하다.
트랜스포머 모델의 경우 텐서 병렬 처리는 멀티헤드 어텐션과 피드포워드 네트워크 모두에 적용된다. 멀티헤드 어텐션에서 쿼리, 키, 값 가중치 행렬은 장치 간에 샤딩되고 어텐션 헤드는 분산된다. 그런 다음 출력 투영이 결합된다. 피드포워드 네트워크에서 두 개의 선형 계층은 일반적으로 분할된다. 첫 번째 계층은 열 방향으로 분할되고 두 번째 계층은 행 방향으로 분할되어 중간 활성화가 샤딩되고 출력이 축소된다. 이 설계는 특정 지점에서만 all-reduce 연산을 요구하여 통신을 최소화한다.
통신 패턴은 중요하다. 각 all-reduce 연산은 텐서 병렬 그룹의 모든 장치 간에 데이터를 교환해야 한다. 통신량은 가중치가 아닌 활성화 크기에 비례하며, 이는 종종 더 작다. 그러나 매우 큰 모델의 경우 통신은 여전히 상당할 수 있다. 이를 완화하기 위해 텐서 병렬 처리는 일반적으로 NVIDIA의 NVLink 또는 AMD의 Infinity Fabric과 같은 고속 상호 연결이 있는 단일 노드 내에서 사용되는 반면, 데이터 병렬 처리는 노드 간에 사용된다.
다른 병렬 처리 기법과의 비교
텐서 병렬 처리는 여러 모델 병렬 처리 전략 중 하나이다. 예를 들어 파이프라인 병렬 처리는 모델을 계층별로 분할하여 각 장치가 연속적인 계층 집합을 담당한다. 이는 계층 경계에서 활성화만 교환되므로 텐서 병렬 처리에 비해 통신을 줄인다. 그러나 파이프라인 병렬 처리는 파이프라인 버블로 인한 유휴 시간이 발생할 수 있다. 데이터 병렬 처리는 반면에 각 장치에 모델을 복제하고 다른 데이터 샘플을 처리하며 각 단계 후에 그래디언트 동기화가 필요하다. 텐서 병렬 처리는 둘 다와 상호 보완적이다. 파이프라인 병렬 처리와 결합하여 장치당 메모리를 줄이고 데이터 병렬 처리와 결합하여 처리량을 높일 수 있다.
또 다른 관련 기법은 트랜스포머 모델에서 시퀀스 차원을 샤딩하는 시퀀스 병렬 처리이다. 이는 메모리 사용량을 더 줄이기 위해 텐서 병렬 처리와 함께 자주 사용된다. 혼합 전문가 모델에서 사용되는 전문가 병렬 처리는 전문가 모듈을 장치 간에 샤딩하며, 이는 모델 병렬 처리의 한 형태이지만 엄밀히 텐서 병렬 처리는 아니다.
응용 및 사용 사례
텐서 병렬 처리는 대규모 언어 모델의 훈련과 추론에 널리 사용된다. 예를 들어 OpenAI가 개발한 GPT-3 모델은 수천 개의 NVIDIA GPU에 걸쳐 텐서 병렬 처리를 포함한 모델 병렬 처리의 조합으로 훈련되었다. 마찬가지로 Anthropic의 Claude 모델과 Google DeepMind의 Gemini 모델은 수천억 개의 매개변수로 확장하기 위해 텐서 병렬 처리에 의존한다. 추론에서 텐서 병렬 처리는 NVIDIA의 Triton Inference Server 및 Hugging Face의 Text Generation Inference에서 수행되는 것처럼 여러 GPU에 계산을 분산하여 GPT-4와 같은 모델을 낮은 지연 시간으로 제공할 수 있게 한다.
언어 모델 외에도 텐서 병렬 처리는 대규모 모델이나 대규모 배치 크기가 분산 계산을 요구하는 컴퓨터 비전 및 과학 컴퓨팅과 같은 다른 영역에서도 사용된다. 예를 들어 고해상도 이미지에서 대규모 비전 트랜스포머(ViT)를 훈련하는 것은 텐서 병렬 처리의 이점을 얻을 수 있다.
과제 및 한계
이점에도 불구하고 텐서 병렬 처리에는 여러 가지 과제가 있다. 주요 문제는 통신 오버헤드이다. All-reduce 연산은 고대역폭, 저지연 상호 연결이 필요하다. 이더넷과 같은 느린 네트워크가 있는 시스템에서는 텐서 병렬 처리가 비효율적일 수 있다. 따라서 일반적으로 단일 서버 또는 긴밀하게 결합된 클러스터 내에서 사용된다. 또 다른 과제는 부하 균형이다. 고르지 않은 샤딩은 일부 장치가 충분히 활용되지 않을 수 있다. 또한 텐서 병렬 처리는 각 장치가 가중치 샤드와 중간 활성화를 저장해야 하므로 신중한 메모리 관리가 필요하다. 마지막으로 텐서 병렬 처리를 수동으로 구현하는 것은 복잡하며 모델 코드 수정과 통신 기본 요소 처리가 필요하다. 다행히 PyTorch 및 TensorFlow와 같은 프레임워크는 텐서 병렬 처리에 대한 내장 지원을 제공하여 많은 복잡성을 추상화한다.
소프트웨어 및 프레임워크 지원
여러 딥러닝 프레임워크와 라이브러리가 텐서 병렬 처리를 지원한다. PyTorch는 텐서 샤딩을 위한 torch.distributed 모듈과 tensor_parallel API를 제공한다. TensorFlow는 모델 병렬 처리를 지원하는 tf.distribute를 제공한다. Megatron-LM, DeepSpeed 및 Hugging Face의 Transformers와 같은 특수 라이브러리는 트랜스포머 모델에 대한 텐서 병렬 처리를 구현했다. 이러한 도구를 통해 연구자와 엔지니어는 장치 수와 샤딩 전략을 지정하는 것만으로 최소한의 코드 변경으로 텐서 병렬 처리를 적용할 수 있다.
향후 방향
모델이 계속 성장함에 따라 텐서 병렬 처리는 중요한 기법으로 남을 것이다. 향후 개발에는 계층적 all-reduce와 같은 더 효율적인 통신 알고리즘과 AMD GPU, Intel 가속기, AWS Trainium 및 Groq의 LPU와 같은 맞춤형 칩을 포함한 이기종 하드웨어와의 더 나은 통합이 포함될 수 있다. 또한 시스템이 최적의 샤딩 전략을 결정하는 자동 병렬 처리 발견은 활발한 연구 영역이다. 대규모 AI에 대한 수요가 증가함에 따라 텐서 병렬 처리는 규모와 효율성의 과제를 해결하기 위해 계속 진화할 것이다.