오늘날의 거대 언어 모델(LLM)은 2017년에 도입된 트랜스포머 모델 아키텍처를 기반으로 합니다. 2017년 이후로 AI 컴퓨팅 성능이 빠르게 발전하여 더 큰 트랜스포머 기반 LLM을 제작할 수 있게 되면서 기능이 대폭 향상되었는데요. 고급 트랜스포머 기반 LLM은 지능형 챗봇, 컴퓨터 코드 생성, 칩 설계와 같은 많은 흥미로운 애플리케이션을 지원하고 있습니다.
최첨단 LLM을 트레이닝하려면 효율적이고 다용도의 소프트웨어 스택이 필요합니다. NVIDIA NeMo는 LLM을 구축하고, 맞춤화하고, 배포할 수 있는 엔드 투 엔드 플랫폼을 제공합니다. NeMo 프레임워크에 깊숙이 통합된 Megatron-Core는 PyTorch 기반 라이브러리로, 거대 LLM을 트레이닝하는 데 필요한 필수 구성 요소와 최적화를 제공합니다. 모델 개발자가 새로운 모델 아키텍처를 탐색하면서 NVIDIA 플랫폼은 계속 확장되어 혁신을 지원합니다.
오늘 NVIDIA는 NeMo와 Megatron-Core가 이제 각각 상태 공간 모델(SSM)의 사전 트레이닝 및 파인튜닝을 지원한다고 발표했습니다. 또한 NeMo는 이제 Google DeepMind가 설명한 Griffin 아키텍처 기반의 트레이닝 모델을 지원합니다.
대체 모델 아키텍처를 살펴봐야 하는 이유는 무엇인가요?
트랜스포머 모델은 이제 유명한 어텐션 메커니즘을 통해 장거리 종속성을 캡처하는 데 탁월한 성능을 제공하므로 글로벌 컨텍스트 이해가 필요한 작업에 이상적인 선택입니다.
그러나 어텐션의 연산 복잡성은 시퀀스 길이에 따라 2차적으로 확장되므로 시퀀스 길이가 증가함에 따라 트레이닝 시간과 트레이닝 비용이 크게 증가합니다. 또한 추론 중에 어텐션은 시퀀스 길이에 따라 선형적으로 증가하는 키-값 쌍의 캐시(KV 캐시라고 함)를 저장해야 합니다. 이로 인해 시퀀스 길이가 증가함에 따라 메모리 공간이 늘어납니다.
최근 SSM은 어텐션의 몇 가지 한계를 극복하여 시퀀스 모델링 작업을 위한 강력한 모델 아키텍처로 떠오르고 있습니다.
보다 효율적인 긴 시퀀스 길이의 트레이닝을 지원하는 SSM
SSM은 시퀀스 모델링 작업을 위한 어텐션 기반 트랜스포머 모델에 대한 효율적인 대안으로 딥 러닝 커뮤니티에서 인기를 얻은 모델 클래스입니다.
SSM의 강력한 특징은 다음과 같습니다.
- 선형 복잡성: SSM은 연산 및 메모리 복잡성이 모두 선형인 반면, 어텐션은 두 가지 모두에서 2차적입니다. 즉, 어텐션보다 훨씬 더 효율적으로 시퀀스의 장거리 종속성을 모델링할 수 있습니다.
- 높은 품질 및 정확도: SSM은 어텐션과 마찬가지로 입력 시퀀스의 토큰을 살펴보고 모델이 가장 관련성이 높은 부분에 집중할 수 있도록 합니다. 그 결과 트랜스포머 기반 모델에 필적하는 품질과 정확도를 얻을 수 있습니다.
- 효율적인 추론: SSM은 KV 캐시가 아닌 일정한 크기의 벡터만 저장하면 되므로, 추론은 특히 시퀀스 길이가 더 긴 경우 메모리 효율성이 높아집니다.
SSM이 더 긴 시퀀스 길이에 제공하는 이점을 설명하기 위해 다음 차트는 시퀀스 길이가 증가함에 따라 트랜스포머 레이어를 트레이닝하는 것과 비교하여 Mamba-2(이 게시물의 뒷부분에서 설명하는 상태 공간 모델 변형) 레이어 트레이닝의 상대적인 속도 향상을 보여줍니다. 시퀀스 길이가 256K로 증가함에 따라 Mamba-2 레이어는 트랜스포머 레이어보다 18배 더 빠릅니다.
트랜스포머: 모델 차원 4,096, 32개 헤드. Mamba-2: 모델 차원 4,096, 상태 차원 128, 8개 그룹.
AI 커뮤니티에서는 Hyena, Mamba-1, 최근에는 Mamba-2 등 여러 SSM 변형이 인기를 끌고 있습니다.
구조화된 상태 공간 이중성 및 Mamba-2
Mamba-2는 여러 벤치마크에서 매우 높은 정확도를 달성한 최신 릴리스입니다. Mamba-2의 핵심에는 새로운 구조화된 상태 공간 이중성(SSD) 레이어가 있으며, 이는 실제로 Mamba-1 모델에 사용된 SSM 수학을 재구성한 것입니다. 이 재구성은 SSM 연산을 행렬 곱셈으로 재구성하여 NVIDIA Tensor 코어의 상당한 행렬 곱셈 성능을 활용할 수 있도록 합니다.
따라서 Mamba-2는 Mamba-1에 비해 훨씬 더 빠르게 트레이닝할 수 있습니다. 또한 Mamba-2는 언어 모델링 작업에서 트랜스포머와 경쟁하는 품질과 정확도를 제공하며 하이브리드 모델에서 몇 개의 어텐션 레이어가 SSD 레이어와 결합될 때 훨씬 더 나은 결과를 산출할 수 있습니다.
하지만 순수 SSM에도 제한이 없는 것은 아닙니다. 예를 들어, 이는 매우 긴 시퀀스에서 정보를 정밀하게 호출해야 하는 ‘사막에서 바늘 찾기’ 유형의 시나리오에서 어려움이 있는 것으로 나타났습니다.
결과를 개선하고 성능을 향상할 수 있는 하이브리드 모델
SSM, SSD, RNN, 트랜스포머를 결합하는 하이브리드 모델은 각 모델 아키텍처의 장점을 활용하면서 개별 약점을 완화할 수 있습니다.
최근 논문에서 NVIDIA ADLR(Applied Deep Learning Research) 팀원을 포함한 연구원들은 하이브리드 Mamba-트랜스포머 모델을 설명했습니다. 이러한 하이브리드 모델에서 표준 트랜스포머 레이어와 새로운 SSM 레이어는 임의의 구성으로 인터리빙할 수 있습니다. 예를 들어, 이 논문에 설명된 8B 하이브리드 모델에는 56개의 레이어가 있습니다. 레이어 중 4개는 셀프 어텐션 레이어, 24개는 Mamba-2 레이어, 28개는 MLP(Multilayer Perceptron) 레이어입니다. 레이어는 Mamba-2 레이어가 먼저 오고 어텐션 레이어가 뒤따라오면서 MLP 레이어가 모델 전체에 고르게 분산되는 방식으로 할당됩니다.
논문에 따르면 하이브리드 8B Mamba-2-하이브리드 모델은 팀이 평가한 “12가지 표준 작업 모두에서 8B 트랜스포머를 능가”합니다. 또한 8B Mamba-2-하이브리드는 “추론 시간에 토큰을 생성할 때 최대 8배 더 빠를 것으로 예측”됩니다.
Mamba-2-하이브리드 모델은 작업을 수행하는 향상된 기능과 추론 중 상당한 성능이라는 이점 외에도 더 높은 컴퓨팅 효율성을 보여줍니다. 아래 차트는 시퀀스 길이가 증가함에 따라 8B 트랜스포머 모델을 트레이닝하는 데 필요한 컴퓨팅과 비교하여 8B Mamba-2-하이브리드 모델을 트레이닝하는 데 필요한 컴퓨팅을 보여줍니다.
2,048개 토큰의 시퀀스 길이에서 두 가지에 필요한 컴퓨팅은 대략 유사하며 하이브리드 모델은 약간의 이점이 있습니다. 그러나 시퀀스 길이가 최대 32,768개의 토큰으로 확장되는 경우 8B 트랜스포머 모델에 필요한 컴퓨팅은 두 배가 되는 반면, 하이브리드 모델에서는 13%만 증가합니다. 최신 언어 모델은 1백만 개 이상의 토큰 시퀀스 길이를 지원하므로 SSM-트랜스포머-하이브리드 모델의 이점은 더욱 커질 것입니다.
새로운 클래스의 모델 아키텍처를 지원하는 첫 단계
모델 아키텍처 혁신은 새로운 수준의 인텔리전스를 제공하는 데 매우 중요합니다. NeMo 및 Megatron-Core는 트랜스포머 기반 모델 구축을 위한 세계적 수준의 지원 외에도 이제 커뮤니티에 SSM 및 SSD는 물론 트랜스포머 모델의 이점과 강점을 결합한 하이브리드 모델을 트레이닝할 수 있는 기능을 제공합니다.
이번 NeMo 릴리스에서는 커뮤니티가 빠르게 실험을 시작할 수 있도록 다음과 같은 초기 기능이 제공됩니다.
- Mamba-2를 포함한 SSD 모델 지원.
- RG-LRU(Grifin 아키텍처) 지원.
- 트랜스포머/SSM 하이브리드 모델 조합 지원.
- Recurrent Gemma(Grifin), 순수 Mamba-2 모델, 8B Mamba-2-하이브리드 모델에 대한 파인튜닝 지원.
- 샤딩 및 모델 병렬 처리 지원.
예정된 릴리스에서는 추가적인 하위 2차 모델 아키텍처, 추가 성능 최적화, FP8 트레이닝에 대한 지원이 계획되어 있습니다.
관련 리소스
- GTC 세션: LLM 기반 에이전트의 세계 탐색(발표: Quantiphi, Inc.)
- GTC 세션: NeMo, TensorRT-LLM, Triton 추론 서버의 가속화된 LLM 모델 정렬 및 배포
- GTC 세션: NVIDIA NeMo 및 AWS를 통한 LLM 트레이닝 최적화
- SDK: NeMo
- SDK: NeMo LLM 서비스
- SDK: NeMo Megatron
