그래프코어 IPU를 통한 희소 대규모 언어 모델 학습

새로운 그래프코어 연구 논문에서는 IPU를 통한 대규모 언어 모델 사전 학습을 예시로 들어 희소 학습을 효율적으로 구현하는 방법을 보여줍니다.

논문 읽기

희소성의 잠재력 탐구

희소 신경망 연구는 혁신의 선순환을 주도합니다. 알고리즘 개선을 통해 소프트웨어 라이브러리가 희소 패턴을 효과적으로 표현할 수 있습니다. 그런 다음 프로그래머는 이러한 라이브러리를 통해 사용 가능한 하드웨어를 최대한 활용하고, 그 결과로 컴퓨팅 요구 사항을 줄일 수 있습니다.

동시에 하드웨어 혁신을 통해 보다 효율적이고 유연한 소프트웨어 라이브러리를 도입하여 연구자가 알고리즘 혁신을 탐색하도록 지원할 수 있습니다.

최근 이 혁신 주기는 모든 수준에서 흥미로운 발전을 가져왔습니다. 예를 들어 복권 가설(Frankle 및 Carbin, 2018, Frankle et al., 2019, Ramanujan et al., 2019, Pensia et al., 2020)에서는 학습된 모델 가중치의 작은 하위 집합이 학습된 기능을 높은 정확도로 나타내기에 충분하다는 알고리즘적 인사이트를 강조한 바 있습니다(Han et al. 2015). 이 아이디어에 힘입어 새로운 학습 방법이 제안되었습니다. 바로 밀집 학습의 컴퓨팅 비용을 발생시키지 않으면서 정확도가 높은 희소 하위 네트워크를 찾는 것입니다(Evci et al. 2020, Jayakumar et al. 2020). 

이 작업과 더불어, 희소 컴퓨팅에 대한 주요 기계 학습 프레임워크의 소프트웨어 지원을 확장하기 위한 중요한 노력이 있었습니다. OpenAI의 조치로 Block Sparse Toolbox를 PyTorch에 통합한 것이 향상된 지원의 한 예입니다.

하드웨어 수준에서도 주요한 혁신이 진행 중입니다. 그래프코어의 IPU(지능형 처리 장치)와 같이 희소 컴퓨팅을 지원하도록 특별히 설계된 새로운 가속 솔루션은 업계와 학계의 혁신가들에게 널리 보급되고 있습니다.

이러한 모든 개발은 최근 AI 모델 크기의 폭발적인 성장을 고려할 때 매우 시의적절합니다. 오늘날 널리 사용되는 GPT3와 같은 대규모 Transformer 기반 언어 모델은 최신 아키텍처에서 계속 증가하는 모델 크기의 적절한 예입니다. 이처럼 더 큰 모델은 일관된 정확도 개선이라는 이점을 안겨주지만, 학습 및 추론 면에서는 훨씬 더 많은 계산이 요구됩니다. 이러한 맥락에서 개선된 희소 방법은 비용 및 에너지 소비를 줄일 수 있는 큰 잠재력을 가지고 있어 높은 전력 소비로 인한 환경 문제를 완화하면서 모델을 더 널리 액세스할 수 있게 돕습니다.

사례 연구에서는 SNN Workshop 2021에서 발표된 바와 같이 IPU 프로세서에서 확장 가능한 언어 모델 사전 학습의 구체적인 예를 사용하여 희소 컴퓨팅의 가능성을 탐구합니다. 특히, BERT 언어 모델의 강력한 희소 사전 학습 기준점을 구현하며 부동 소수점 연산(FLOP) 기반과 비교할 때 밀집 모델보다 탁월함을 입증합니다. 이 사례 연구와 관련된 하드웨어, 소프트웨어 및 알고리즘의 다양한 측면에 대한 탐구를 넘어 우리의 발견이 이 분야에서의 더 활발한 연구를 위한 실질적인 출발점이 되기를 바랍니다. 

이제 워크샵에서 발표한 주요 결과 중 일부에 초점을 맞추겠습니다. IPU에 대한 개요와 희소 컴퓨팅에 IPU를 적합하게 만드는 기능부터 시작하겠습니다. 다음으로 희소 알고리즘을 구현하는 데 사용하는 오픈 소스 소프트웨어 라이브러리를 소개하겠습니다. 마지막으로 효율적인 실행에 도움이 될 수 있는 몇 가지 알고리즘 속성을 조명하는 실험 결과에 대해 논의합니다.

하드웨어: IPU 아키텍처 및 희소 컴퓨팅  

그래프코어 IPU는 머신러닝 워크로드용으로 특별히 설계된 새로운 대규모 병렬 프로세서입니다. 2세대 칩에는 각각 6개의 스레드가 있는 1,472개의 완전히 독립적인 프로세서 코어(타일이라고 함)가 있어 한 번에 8,832개의 독립적인 동시 명령어(다중 명령어 다중 데이터(MIMD))를 실행할 수 있습니다. 특히, 각 타일에는 카드당 180TB/s의 속도로 액세스할 수 있는 자체 프로세서 내 메모리 청크가 있으며, 이는 오프칩 DRAM의 메모리 액세스 대역폭의 100배 이상입니다. IPU에 컴퓨팅과 메모리를 함께 배치하면 에너지 효율성이 높아질 뿐만 아니라 세분화되고 유연한 접근 방식을 사용하여 병렬화해야 하는 매우 불규칙한 계산에 아주 적합합니다. 

IPU는 많은 양의 병렬 처리를 위해 특별히 설계되었기 때문에 희소 워크로드에 특히 적합합니다. 네트워크 토폴로지의 알고리즘적 희소화는 종종 미세한 입도의 매우 불규칙한 계산 그래프들을 초래하므로, 가장 작은 계산 구조조차 효율적으로 병렬화할 수 있는 IPU의 능력이 장점이 될 수 있습니다. 

IPU 아키텍처는 희소 구조의 입도가 높은 경우에도 희소 학습 및 추론을 위한 계산 효율성을 유지할 수 있습니다. 크기 16의 곱셈누적 유닛을 사용하면 IPU가 1x1, 4x4, 8'8 또는 16'16의 매우 작은 블록 크기에 대해서도 우수한 스레드 활용도와 성능을 얻을 수 있습니다. IPU는 또한 처리량을 높이는 데 도움이 되는 일반적인 희소 원시 컴퓨팅을 가속화하도록 설계된 여러 명령들을 지원합니다. 

소프트웨어: Poplar와 Popsparse 소프트웨어 라이브러리 

희소 컴퓨팅을 위한 유연한 저수준 C++ 라이브러리인 Popsparse를 통해 사용할 수 있는 최적화된 커널로 모든 하드웨어 및 아키텍처 기능을 활용할 수 있습니다. Popsparse는 IPU에서 실행되는 모든 것을 뒷받침하는 그래프코어의 그래프 구성 및 컴파일 API인 Poplar를 기반으로 하는 그래프코어의 오픈 소스 소프트웨어 라이브러리(PopLibs로 알려짐) 중 하나입니다.

희소성 커널은 PopART API를 통해 직접 타게팅되거나 이 작업에 사용된 TensorFlow 및 Keras 레이어를 통해 기존 TensorFlow 모델에 쉽게 통합될 수 있습니다. FP16 및 FP32 모두에서 블록 크기가 1x1, 4x4, 8x8, 16x16인 희소성이 지원됩니다. 그래프코어는 GitHub의 공개 예제를 통해 오픈 소스 구현 및 튜토리얼을 제공합니다.

Poplar 소프트웨어 스택은 알고리즘의 희소성 패턴이 고정되어 있는지 또는 시간에 걸쳐 발전하는지에 따라 정적 및 동적 희소성을 모두 지원합니다. 정적 희소성을 사용하면 컴파일 시 최적화가 향상되어 실행 효율성이 개선되므로 이미 학습된 모델과 관련된 추론 워크로드에 적합합니다. 그러나 많은 학습 알고리즘은 학습 과정에서 알고리즘의 희소성 패턴을 천천히 발전시킵니다. 이러한 경우 동적 희소성이 더 큰 유연성을 제공할 수 있습니다.

동적 희소성 기술의 구현을 단순화하기 위해, IPU로 스트리밍하기 전에 호스트에서 희소성 패턴을 쉽게 조작할 수 있는 다양한 기능이 제공됩니다. 사용자는 구현에 가장 적합한 희소성 형식을 선택한 다음 이를 IPU의 내부 CSR 기반 희소 형식으로 변환할 수 있습니다. 이를 통해 희소성을 기존 코드에 쉽게 통합하고 Python으로 패턴 조작을 표현할 수 있습니다.  

RNN 아키텍처에서 GPT2에 이르는 희소 모델의 전체 모델 구현과 RigL 및 희소 어텐션과 같은 정적 및 동적 희소 최적화 기술도 오픈 소스로 제공됩니다.   

알고리즘: 부분 신경망 학습을 위한 동적 희소성 사용 

IPU의 Poplar 소프트웨어 라이브러리는 희소 학습을 위한 최신 알고리즘을 구현하는 데 필요한 모든 것을 제공합니다. 그러나 BERT와 같은 기존의 밀집 모델의 경우 최상의 알고리즘 접근 방식이 무엇일지 즉시 알 수 없습니다.

전통적으로 가지치기 기술은 학습된 모델의 추론 속도를 높이는 간단하고 효과적인 방법입니다. 학습 중에 유사한 컴퓨팅 이점을 달성하는 것은 일반적으로 더 어렵습니다. 한 가지 일반적인 기술은 희소 토폴로지로 학습을 시작하고 시간이 흐를수록 연결 패턴을 발전시키는 것입니다(Evci et al. 2020, Jayakumar et al. 2020). 이 접근 방식을 사용하면 현재의 하위 네트워크를 가지치기하고 재성장시키면서 희소 연결이 주기적으로 수정됩니다. 이렇게 하면 계산 비용 절감의 이점을 유지하면서 원래 과모수화된 모델의 자유도를 효과적으로 탐색할 수 있습니다. 

물론 이러한 방법의 효율성은 특정 가지치기 및 재배치 기준에 크게 의존합니다. 옵션은 단순한 규모 기반 가지치기 및 임의 재할당에서 더 정교한 그래디언트 기반 접근법에 이르기까지 다양합니다.

여기서 다양한 블록 크기 1x1, 4x4, 8x8 및 16x16에 대해 규모 기반 가지치기 및 무작위 재성장을 사용하여 간단하지만 잘 조정된 동적 희소성 기준점을 구현하도록 선택합니다. 특히 BERT 언어 모델링 작업을 기반으로(Devlin et al. 2018), 우리는 완전히 연결된 인코더 가중치(임베딩 가중치 제외)를 0.9의 희소성 비율로 희소화하여 이 레이어에 있는 가중치의 10%만 남도록 했습니다. 규모 기반 가지치기와 가중치의 특정 비율(가지치기 비율 0.5, 시간 경과에 따른 코사인 감쇠, 그림 1 참조)의 무작위 재성장을 사용하여 희소성 패턴(사전 학습의 1단계에서 160회)을 정기적으로 업데이트합니다. 

그런 다음 학습률을 조정하여 최상의 작업 성능을 얻습니다. 이 접근 방식을 사용하여 다양한 모델 크기를 BERT Base까지 희소화합니다. 그림 2는 1단계 이후에 달성된 마스킹된 언어 모델 손실을 보여줍니다. 이는 마스킹된 토큰의 무작위 하위 집합을 예측하는 유효성 검사 성능에 해당합니다. FLOP 예산으로 비교할 때 동적으로 희소 학습된 모델은 모든 모델 크기에 대해 밀집 모델보다 성능이 뛰어납니다. 

이는 동적 희소 학습 접근법이 대규모 사전 학습의 계산 부담을 줄이는 효과적인 기술이 될 수 있음을 보여줍니다. 흥미롭게도 동적 희소 학습 접근법은 최대 16x16까지의 더 큰 블록 크기에서도 여전히 FLOP 효율적입니다(표 1). 더 큰 블록 크기를 사용하면 실행을 더욱 효율적으로 만들기 위해 활용할 수 있는 추가 최적화가 가능합니다.  

향후 작업을 통해 구조화된 희소성의 성능을 개선하여 대규모 언어 모델의 컴퓨팅 효율적인 학습을 위한 최상의 절충안을 식별할 수 있습니다.

희소 학습 효율성을 위한 새로운 방향

이 연구에서 우리는 새로운 IPU 하드웨어를 사용하여 대규모 언어 모델의 동적인 상시 희소 학습을 위한 간단한 기준점을 설정했습니다. 이러한 모델들에 대해 주어진 FLOP 예산에서 비교할 때 희소 학습이 밀집 학습보다 더 효율적일 수 있음을 보여주었습니다.

또한 더 큰 블록을 사용하면 밀집 학습에 비해 FLOP 이점을 부분적으로 유지할 수 있습니다. 이러한 맥락에서, 블록 희소성과의 호환성을 보장하기 위해 구조화되지 않은 학습 방법을 최소한으로만 수정했다는 점에는 주목할 만한 가치가 있습니다.  이는 추가 최적화와 알고리즘 혁신이 여기에서 달성한 블록 희소 정확도를 개선하고 희소 학습을 훨씬 더 효율적으로 만드는 데 도움이 될 수 있음을 의미합니다.

우리는 이 연구가 새로운 하드웨어와 소프트웨어가 희소 네트워크 접근 방식의 다음 혁신 주기를 주도하고 활성화할 수 있는 방법을 보여준다고 믿습니다.

논문 읽기

희소성 프로젝트에서 그래프코어 IPU를 사용하는 데 관심이 있는 경우 그래프코어 Academic Programme에 신청서를 제출해 주세요.