OpenAI의 Techniques for Training Large Neural Networks을 번역했습니다. 오역이 있다면 알려주세요.
뉴럴 네트워크 대형 모델은 최근 AI 발전의 핵심 요소입니다. 그러나 이것을 훈련하는 일은 동기화된 단일 계산을 수행하기 위해 GPU 클러스터를 조정해야 하는, 리서치와 엔지니어링 상의 난제입니다. 클러스터와 모델 크기가 점점 커짐에 따라 기계 학습 실무자는 모델 훈련 작업을 다수의 GPU에서 병렬화하기 위해 다양한 기술을 개발해왔습니다. 언뜻 보기에 이러한 병렬 처리 기술은 이해하기 까다로울 수 있지만 계산 구조에 대해 몇 가지 가정을 두면 훨씬 더 명료하게 느껴지실 겁니다. 그런 다음 여러분은 패킷 주변에서 셔틀을 전환하는 네트워크처럼 A에서 B로 불투명한 비트를 계속 왕복하며 이해하면 됩니다.
3개의 레이어를 갖는 모델에 대한 다양한 병렬화 전략 예시. 각 색상은 하나의 레이어를 나타내고 점선은 서로 다른 GPU를 구분합니다.
병렬 처리 미적용
뉴럴 네트워크 훈련은 반복적인 과정입니다. 매 반복마다 데이터 배치, 각 훈련 샘플에 대한 출력 값을 계산하기 위해 모델 레이어로 순방향 전달을 진행합니다. 그런 다음 레이어를 통해 역방향 전달을 진행하는데 각 매개변수에 대한 그래디언트를 계산하면서 각 매개변수가 최종 출력 값에 미치는 영향을 전파합니다. 배치, 매개변수 그리고 일부 매개변수 별 최적화 상태에 관한 평균 그래디언트를 Adam 같은 옵티마이저에 전달합니다. 그러면 옵티마이저는 다음 반복을 위해 (데이터 상의 성능을 약간이나마 개선할) 매개변수와 새로운 매개변수 별 최적화 상태를 계산해냅니다. 훈련 과정이 데이터 배치 상에서 반복됨에 따라 모델은 점점 더 정확한 출력을 생성하게끔 진화합니다.
이러한 훈련 프로세스를 각종 병렬 처리 기술이 아래 내용을 포함하여 다양한 차원으로 분할합니다.
- 데이터 병렬 처리 - 서로 다른 GPU에서 배치의 서로 다른 하위 집합을 실행합니다.
- 파이프라인 병렬 처리 - 서로 다른 GPU에서 모델의 서로 다른 레이어를 실행합니다.
- 텐서 병렬 처리 - 여러 개의 GPU에 분할되는 행렬 곱셈처럼 단일 작업에 대한 수리적 연산을 나눕니다.
- 전문가 혼합 - 각 레이어의 일부만 사용하여 각 샘플을 처리합니다.
(이 게시물은 GPU를 사용하여 신경망을 훈련한다고 가정하지만 다른 종류의 신경망 가속기를 사용하는 경우에도 동일한 아이디어를 적용할 수 있습니다.)
데이터 병렬 처리
데이터 병렬 훈련은 동일한 매개변수를 여러 GPU(종종 "작업자"라고 함)에 복사하고 각 GPU마다 동시에 처리해야 할, 서로 다른 샘플들을 할당하는 과정을 의미합니다. 데이터 병렬 처리를 하더라도 모델 크기는 여전히 단일 GPU 메모리 내에서 수용 가능해야 합니다. 다만, 매개변수의 여러 중복 복사본을 저장하는 비용과 맞바꿔 GPU 다수의 계산 능력을 사용 가능하게 됩니다. 또한 처리하는 도중 매개변수를 일시적으로 CPU 메모리에 오프로드하는 식으로 GPU의 가용한 유효 RAM을 늘리는 전략을 쓸 수 있습니다.
여러 데이터 병렬 작업자가 매개변수 복사본을 업데이트할 때 각 작업자가 거의 동일한 매개변수 값을 지속적으로 가질 수 있게 조정해줘야 합니다. 가장 간단한 접근 방법은 작업자 사이에 블로킹 통신을 도입하는 것입니다. (1) 각 작업자마다 그래디언트를 독립적으로 계산합니다. (2) 작업자들의 그래디언트 평균을 구합니다. (3) 각 작업자마다 동일한 매개변수의 새로운 값을 독립적으로 계산합니다. (2) 번 단계는 훈련 처리량에 악영향을 줄 수 있는, 상당량의 데이터(작업자 수와 매개변수 크기에 비례할)를 전송해야만 하는 블로킹 방식의 평균 연산입니다. 이 오버헤드를 제거하기 위해 다양한 비동기식 동기화 방법이 있지만 학습 효율성이 떨어집니다. 실제로 동기식 접근 방법이 일반적으로 선택됩니다.
파이프라인 병렬 처리
파이프라인 병렬 훈련에서는 모델을 순차적인 청크로 나눠 여러 GPU에 할당합니다. 각 GPU는 매개변수의 일부만 보유하므로 같은 모델이라면 GPU는 개당 비례적으로 더 적은 메모리를 소비합니다.
대형 모델을 연속적인 레이어의 청크로 나누는 건 간단합니다. 다만 레이어의 입력과 출력 사이에 순차적 종속성이 있기 때문에 이를 단순하게 구현하면 작업자가 이전 시스템 출력 값이 입력으로 들어오길 기다리면서 유휴 시간이 길게 발생하게 됩니다. 이러한 대기 시간 청크를 "버블"이라고 하며 유휴 상태 장비에서 수행할 수 있는 연산 능력을 낭비하게 만듭니다.
모델을 레이어에 따라 4개 파티션으로 수직 분할한, 단순하게 구현한 파이프라인 병렬 처리 설정 그림입니다. 작업자 1은 네트워크의 첫 번째 레이어(입력과 가장 가까운)의 모델 매개변수를 호스팅 하는 반면 작업자 4는 네 번째 레이어(출력과 가장 가까운)를 호스팅 합니다. "F", "B"와 "U"는 각각 순방향, 역방향과 업데이트 작업을 나타냅니다. 아래 첨자는 작업을 실행하는 작업자를 나타냅니다. 순차적인 종속성으로 인해 데이터를 한 번에 한 작업자가 처리하므로 유휴 시간에 거대한 "버블"이 발생합니다.
데이터 병렬 처리 아이디어를 더 잘 구현해봅시다. 각 작업자가 한 번에 데이터 요소 하위 집합의 한 단위씩 처리하도록 함으로써 버블 비용을 줄일 수 있고 대기 시간에 신규 연산을 영리하게 배정할 수 있습니다. 핵심 아이디어는 한 개의 배치를 여러 개의 마이크로 배치로 나누는 것입니다. 각 마이크로 배치는 크기에 반비례하여 처리 속도가 빠를 것이고 각 작업자는 다음 마이크로 배치가 사용 가능해지면 바로 작업을 시작할 것이기 때문에 파이프라인 실행 속도가 빨라질 것입니다. 마이크로 배치가 충분하다면 작업자는 단계 시작과 끝에서 발생하는 최소한의 버블만 제외하고 대부분의 시간을 활용할 수 있습니다. 그래디언트를 마이크로 배치 전체에서 평균화하며 모든 마이크로 배치가 완료된 다음 매개변수에 대한 업데이트가 발생합니다.
모델을 분할하는 작업자 수를 보통 파이프라인 깊이라고 부릅니다.
순방향 전달 과정에서 작업자는 레이어 청크의 출력 값(활성화라고 함)을 다음 작업자에게 보내기만 하면 됩니다. 역방향 전달 과정에서는 해당 활성화의 그래디언트만 이전 작업자에게 보냅니다. 이러한 전달을 예약하는 방법과 마이크로 배치에 걸쳐 그래디언트를 집계하는 방법에 있어서 상이한 설계 방향이 존재할 수 있습니다. GPipe는 각 작업자가 순방향과 역방향 전달을 순서를 지켜 처리하게 하고 마지막에 여러 마이크로 배치의 그래디언트를 동기적으로 집계합니다. 반면에 PipeDream은 각 작업자가 순방향과 역방향 전달을 번갈아 처리하도록 예약합니다.
배치당 4개의 마이크로 배치를 사용하는 GPipe와 PipeDream 파이프라인 작업 계획 비교. 마이크로 배치 1-8은 두 개의 연속적인 데이터 배치에 해당합니다. 이미지에서 "(숫자)"는 작업이 수행되는 마이크로 배치를 나타내고 아래 첨자는 작업자 ID를 나타냅니다. PipeDream은 과거 상태의 매개변수 값으로 일부 계산을 수행하여 더 높은 효율성을 달성합니다.
텐서 병렬 처리
파이프라인 병렬 처리는 모델을 레이어 별로 "수직으로" 분할합니다. 또한 일반적으로 텐서 병렬 훈련이라고 해서 레이어 안에서 특정 작업을 "수평으로" 분할할 수도 있습니다. 여러 최신 모델(예: 트랜스포머)의 계산 병목 현상은 활성화 배치 행렬에 큰 가중치 행렬을 곱하는 작업에서 발생합니다. 행렬 곱셈은 행과 열의 쌍 별 내적으로 생각할 수 있습니다. 서로 다른 GPU에서 독립적으로 내적을 계산하거나 GPU 각자가 내적의 일부를 계산해서 결과를 합산하는 일 또한 가능합니다. 어느 전략이든 가중치 행렬을 고른 크기의 "샤드"로 분할하고, 각 샤드를 서로 다른 GPU에 호스팅 하고, 전체 행렬 곱 중 해당 샤드와 관련 있는 부분을 각자 계산하고 결과를 결합하기 위해 통신을 수행합니다.
트랜스포머의 셀프 어텐션과 MLP 레이어 내에서 행렬 곱셈을 병렬화하는 Megatron-LM를 하나의 예로 들 수 있습니다. PTD-P는 텐서, 데이터와 파이프라인 병렬 처리를 사용합니다. 파이프라인 스케줄은 각 장치마다 여러 개의 비연속적 레이어를 할당하는데 이때 네트워크 통신 비용은 늘어나지만 대신 버블 오버헤드는 줄어듭니다.
때때로 네트워크 입력은 높은 수준의 병렬 연산을 사용하는 차원 축을 따라 교차 통신을 고려하여 병렬화할 수 있습니다. 시퀀스 병렬 처리는 입력 시퀀스를 시간에 따라 여러 개의 하위 샘플로 분할해서 더 세분화한 크기의 샘플로 계산을 수행하도록 하여 피크 메모리 소비를 비례적으로 줄이는, 한 아이디어입니다.
전문가 혼합(MoE)
전문가 혼합(MoE) 접근 방법의 경우 네트워크 일부만 하나의 입력에 대한 출력 값을 계산하는 데 사용됩니다. 한 가지 예로 들 수 있는 접근 방법은 다수의 가중치 집합을 갖고 있으면서 추론 때 게이팅 메커니즘을 통해 네트워크가 사용할 집합을 선택하는 것입니다. 이것은 계산 비용의 증가 없이 더 많은 매개변수를 사용 가능하게 합니다. 이러한 개별 가중치 집합을 "전문가"라고 하며 전문적인 계산과 기술을 각각의 전문가에게 할당하는 방법을 네트워크가 학습하길 희망합니다. 서로 다른 GPU가 서로 다른 전문가를 호스팅 할 수 있으므로 이것은 모델이 사용하는 GPU 수를 확장할 수 있는 명확한 방법입니다.
전문가 혼합(MoE) 레이어의 그림. 해당 게이팅 네트워크에 의해 n 개의 전문가 중 2개만 선택됩니다. (이미지 출처: Shazeer et al., 2017)
GShard는 MoE 레이어만 여러 개의 TPU 장치에 분할하고 다른 레이어는 완전히 복제하는 방식으로 MoE 트랜스포머를 최대 6000억 개의 매개변수로 확장합니다. Switch Transformer는 하나의 입력을 단일 전문가에게 라우팅 하여 훨씬 더 희소한, 수조 개의 매개변수로 모델 크기를 확장합니다.
메모리 절약을 위한 기타 설계
점점 커지는 뉴럴 네트워크 훈련을 좀 더 다루기 용이하게끔 다양한 계산 전략이 개발되었습니다. 예를 들면,
-
그래디언트를 계산하려면 다량의 장치 RAM을 소모하는 활성화 원천 값을 저장해놔야 합니다. 체크포인팅(활성화 재계산이라고도 함)은 활성화의 하위 집합을 저장하고 역방향 전달을 수행할 때 중간 부분을 적시에 재계산합니다. 이것은 기껏해야 한 번 더 추가되는 전체 순방향 전달의 계산 비용으로 많은 메모리를 절약할 수 있습니다. 또한 선택적 활성화 재계산을 통해 컴퓨팅 비용과 메모리 비용을 지속적으로 절충해 나갈 수 있습니다.
-
혼합 정밀도 훈련은 정밀도가 낮은 숫자(가장 일반적으로는 FP16)를 사용하여 모델을 훈련하는 것입니다. 최신 가속기는 정밀도가 낮은 숫자로 훨씬 더 많은 FLOP 수에 도달할 수 있으며 장치의 RAM도 절약할 수 있습니다. 적절한 주의를 기울인다면 모델 결과 정확도에 거의 타격 없습니다.
-
오프로딩은 사용하지 않는 데이터를 CPU 또는 다른 장치에 일시적으로 오프로드하고 나중에 필요할 때 다시 읽는 방법입니다. 단순한 구현은 훈련 속도를 상당히 지연시키지만 정교한 구현은 데이터를 미리 가져다 놓으므로 장치가 데이터 입수를 기다릴 필요가 없습니다. 이 아이디어의 구현물 중 하나는 가용한 모든 하드웨어에 매개변수, 그래디언트와 옵티마이저 상태를 할당하고 필요에 따라 구체화하는 ZeRO입니다.
-
Adafactor 같은 메모리 효율적 옵티마이저는 옵티마이저가 유지, 관리하는 실행 상태의 메모리 공간을 줄이기 위해 제안되었습니다.
-
압축은 네트워크의 중간 결과를 저장하는 데에도 사용할 수 있습니다. 예를 들어, Gist는 역방향 전달을 위해 저장한 활성화를 압축합니다. DALL·E는 그래디언트를 동기화하기 전에 압축합니다.
OpenAI는 제반 인프라부터 현실에 배포하는 문제에 이르기까지 대형 모델을 훈련하고 또 개선하고 있습니다. 이 게시물의 아이디어를 현실에 만들어보고 싶다면 특히 확장과 응용 연구 팀들과 연관된, 채용 공고에 지원하세요!
