[DeepSpeed: 모두를 위한 대규모 모델 트레이닝 🧑‍💻] GPT-3의 파라미터가 1750억 개라고 하죠. Google, MS도 크고 아름다운(...) 모델 만드는 작업에 박차를 가하고 있습니다. 이제 파라미터 개수는 조 단위를 넘을 겁니다. (학부 시절 다중 회귀 분석 배울 때 파라미터 3개도 충공깽이었는데... 😔) 이런 굇수 모델을 학습하는데 물리적인 문제는 없을까요? 일단 파라미터, 그래디언트 용량만 해도 GPU 1개 메모리 크기의 몇백 배는 될 겁니다. 계산 처리량 또한 단일 장치가 소화할 수준이 아니고요. 병렬화해야겠죠. 병렬화 전략으로 데이터 병렬화, 모델 병렬화, 다소 생소하지만 파이프라인 병렬화 세 가지가 있습니다. (1) 데이터 병렬화는 구현하기 쉽지만 모델을 통째로 복사하기 때문에 메모리 효율성이 떨어져요. 그래디언트를 집계하기 위한 통신 비용도 제법 들고요. (2) 모델 병렬화는 모델을 조각조각 쪼개기 때문에 메모리 효율성이 좋지만 모델 아키텍처마다 별도 구현해줘야 하고 통신 비용 또한 매우 큽니다. (3) 파이프라인 병렬화는 모델의 층 별로 분할하는 걸 의미합니다. 활성화 값만 주고받으면 되므로 통신 비용이 낮습니다. 그러나 노예 GPU가 본인 텃밭 감자 다 캤다고 놀면 안 되니... 😔 데이터를 마이크로 배치로 분할하여 계속 넣어줘야 합니다. MS는 데이터 병렬화의 메모리 효율성 문제를 ZeRO라는 기술로 상당히 해결했습니다. ZeRO의 핵심 아이디어는 파라미터, 그래디언트, 옵티마이저 상태 정보를 분할하여 각 GPU마다 따로 저장하되 GPU 본인(?)이 들고 있는 정보를 필요한 다른 녀석들(?)에게 지속 전파시키는 것입니다. 이 ZeRO를 포함해서 여러 딥러닝 최적화 기술을 구현한 게 deepspeed 라이브러리입니다. PyTorch 기반 API를 제공하고요. (MS: 뭐? TensorFlow 지원? 안돼. 안 받아줘. 빨리 돌아가. 😔) MS는 지난 9월 10일 리서치 블로그를 통해 deepspeed의 놀라운 기능 4가지를 자랑(...)했습니다. 1. 3D 병렬화 위에서 말한 데이터, 모델, 파이프라인 병렬화를 조합하여 함께 사용하는 기술을 3D 병렬화라고 합니다. deepspeed는 3D 병렬화를 지원하며 GPU 클러스터 구조를 인식하여 병렬화를 최적화합니다. 즉 통신 비용이 큰 모델 병렬화는 최대한 같은 노드 내에서 GPU 할당을, 통신 비용이 작은 파이프라인 병렬화는 서로 다른 노드 간에 할당이 이뤄지게 합니다. 데이터 병렬화의 경우 위에서 말한 ZeRO 기술 적용으로 메모리 효율성을 달성함과 동시에 유효 대역폭 또한 증가시킵니다. 이 기술을 통해 파라미터 수조 개 짜리 모델도 트레이닝 가능하답니다. 2. ZeRO 오프로드 GPU 클러스터가 아닌 GPU 서버 한대를 생각해보죠. 역시나 모델이 큰 경우 메모리 부족이 가장 큰 문제입니다. ZeRO 변형 기술, 즉 ZeRO 오프로드를 이용하면 CPU, GPU 메모리를 모두 사용함으로써 메모리 효율성을 달성할 수 있습니다. GPU는 그래디언트 계산이 끝나는 족족 CPU로 오프로드합니다. CPU는 그래디언트를 취합하여 파라미터를 업데이트하면서 GPU로 스왑 합니다. 이 기술을 통해 단일 GPU에서 10배 큰 모델 트레이닝이 가능하답니다. 3. 희소 어텐션 트랜스포머 계열이 사용하는 셀프 어텐션은 시계열 상 자기 시점 이전의 모든 셀에 대해 스코어 계산을 하기 때문에 n 제곱의 계산 복잡도를 갖습니다. 그래서 아주 긴 시퀀스 입력을 받아들이기 어려워요. 희소 어텐션의 아이디어는 모든 셀 말고 특정 패턴에 해당하는 셀들(예컨대 인접한 몇 개)에 대해서만 스코어 계산을 하는 것입니다. 이렇게 하면 기존 어텐션 대비 정확도가 비슷하거나 상승한다고 하니, 이것은 실화입니까...😔 해당 기능을 사용하면 10배 더 긴 시퀀스 지원, 6배 더 빠른 수행이 가능하답니다. 4. 1비트 Adam 그래디언트와 옵티마이저 상태 값은 분산 장치들이 자주 전달받거나 전달해줘야 하는, 즉 통신에서 병목이 될 수 있는 데이터입니다. 그래서 SGD나 모멘텀의 경우 이를 1비트로 압축하여 통신하고 오차를 보정해주는 기술을 사용하죠. 그러나 Adam은 분산 항으로 인한 비선형성이 있어 이 기술을 사용하기 어렵습니다. 1비트 Adam 기술은 이 문제를 한 가지 트릭으로 풀어냈습니다. 처음 몇 번은 비압축으로 진행하고 어느 단계에 다다르면 분산 항을 상수로 고정한 후 모멘텀 항만 압축하여 전송하는 것입니다. 이 기술은 비압축 Adam과 동일한 수렴률을 보여준다고 하네요. 그 결과 통신량을 최대 5배 감소시킬 수 있다고 합니다. GPT-3 이후 AI는 'Size does matter'의 시대로 접어든 것 같습니다. 테크 기업들은 핵 군비 경쟁처럼 각자 초대형 모델을 만들고 있어요. 학습된 최종 모델이 핵탄두라면 deepspeed 같은 딥러닝 최적화 라이브러리는 그걸 작동하게 만드는 대륙간 탄도 미사일(ICBM) 같은 거겠죠. 잘 드러나지 않지만 되려 더 대단하고 무섭습니다. 😔