Cohere가 Google Cloud TPU로 언어 모델 학습을 가속화하는 방법

지난 수년 동안 대규모 언어 모델(LLM) 학습이 발전하면서 자연어 처리(NLP)가 소수 기업만 누릴 수 있는 최첨단 기술에서 다양한 일반 애플리케이션에서 활용되는 강력한 구성 요소로 진화했습니다. 챗봇부터 콘텐츠 검토, 카테고리 지정에 이르기까지 NLP의 일반적인 규칙은 모델의 규모가 커질수록 언어를 이해하고 생성할 때 정확도를 더 높일 수 있다는 점입니다.

하지만 더 크고 강력한 언어 모델을 만들기 위한 과정에서 규모는 중요한 도전과제가 되었습니다. 모델 규모가 너무 커져서 단일 장치에 맞지 않을 경우 분산형 학습 전략이 필요하고, 이로 인해 방대한 메모리 용량과 신속한 상호 연결을 갖춘 광범위한 컴퓨팅 리소스가 요구됩니다. 또한 하드웨어와 시간 리소스를 최적화할 특수한 알고리즘도 필요합니다.

Cohere 엔지니어들은 이러한 규모 문제를 해결할 솔루션을 개발하고 있으며 이미 성과를 내고 있습니다. Cohere는 일반적으로 이러한 프로젝트에 필요한 인프라나 심층적인 ML 전문성이 없어도 강력한 LLM으로 작동할 수 있는 플랫폼을 개발자에게 제공합니다. JAX pjit와 TPUv4를 사용한 언어 모델의 확장 가능한 학습이라는 최신 기술 문서에서 Cohere 엔지니어들은 Google Cloud가 최근 발표한 Cloud TPU v4 포드에 배포된 새로운 FAX 프레임워크가 LLM을 수천억 개의 매개변수로 확장하는 문제를 어떻게 해결하는지 설명합니다. 특히 이 보고서는 Cohere가 텐서 및 데이터 동시 로드를 통해 달성할 수 있었던 학습 효율성의 혁신적인 면모를 잘 보여줍니다.

이 프레임워크는 확장성과 신속한 프로토타입 제작이라는 두 가지 중요한 개선사항을 통해 대규모 언어 모델의 연구, 개발, 생산을 가속화하는 것을 목표로 합니다. Cohere는 규모가 더 큰 모델을 보다 빠르게 학습시켜 모델을 개선함으로써 고객에게 향상된 모델을 신속하게 제공합니다. 또한 이 프레임워크는 새로운 아이디어를 실험하고 테스트하여 고객 서비스 챗봇을 작동시키는 생성 모델을 만드는 등의 구체적인 목표를 해결하는 모델의 신속한 프로토타입 제작도 지원합니다. Cohere는 모델 유형을 전환하고 다양한 목표에 맞게 최적화할 수 있는 기능을 통해 궁극적으로 특정 사용 사례에 최적화된 모델을 제공할 수 있습니다.

FAX 프레임워크는 기기와 워크로드의 관계를 추상화하는 JAX의 파티션을 나눈 적시 컴파일(pjit) 기능에 크게 의존합니다. 이를 통해 Cohere 엔지니어들은 현재 작업에 이상적인 구성으로 기기와 프로세스를 조정하여 효율성과 성능을 최적화할 수 있습니다. Pjit는 다른 호스트에 상주하는 기기를 비롯해 여러 기기에서 실행되는 단일 프로그램(XLA 연산)으로 임의의 함수를 컴파일하여 작동합니다.

또한 Cohere의 새로운 솔루션은 Google Cloud의 새로운 TPU v4 포드를 활용하여 이전 파이프라인 동시 로드보다 효율적인 텐서 동시 로드도 수행합니다. 이름에서 알 수 있듯, 파이프라인 동시 로드 방식은 하나의 긴 조립 라인처럼 선형 방식으로 가속기를 사용하여 워크로드를 확장합니다. 가속기는 데이터의 마이크로 배치를 각각 다음 배치로 전달하기 전에 처리한 후 역순으로 역방향 전달을 실행해야 합니다. 

텐서 동시 로드는 '파이프라인 버블'이라고도 하는 파이프라인 동시 로드의 가속기 유휴 시간을 제거합니다. 텐서 동시 로드에는 여러 기기에서 동시에 계산을 수행하기 위해 여러 가속기에 대형 텐서(단락 내 단어와 같이 여러 객체 간 관계를 정의하는 수학적 배열)를 분할하는 작업이 포함됩니다. 파이프라인 동시 로드가 조립 라인을 계속 늘린다면, 텐서 동시 로드는 엔진과 차체 등을 각각 만드는 일련의 병렬 조립 라인입니다. 여러 조립 라인이 동시에 작동하여 신속하게 완성차를 제작합니다.

그런 다음 이러한 계산 결과가 정렬됩니다. 이 프로세스는 Google Cloud TPU v4 VM이 이전 버전인 v3보다 연산 능력이 두 배 이상 향상되면서 실용화된 것입니다. Cohere는 v4 칩의 우수한 성능 덕분에 아이디어를 반복하고 검증하여 이전보다 1.7배 더 빠르게 계산을 수행할 수 있었습니다.

Google Cloud가 LLM 학습에 필요한 이유

Cohere는 다년간 기술 파트너십의 일환으로 Google Cloud의 고급 AI 및 ML 인프라를 활용하여 자체 플랫폼을 구동했습니다. Cohere는 대규모 ML에 최적화된 Google Cloud의 커스텀 설계 머신러닝 칩인 Cloud TPU에 제품을 개발하고 배포합니다. 최근 Cohere는 Google Cloud TPU에서 FAX를 사용한 LLM 학습으로 새롭게 개선된 모델과 확장성을 발표했습니다. 이 모델은 TPU v3에서 TPU v4 전환으로 지금까지 총 1.7배 더 빠른 속도를 달성한 것으로 나타났습니다. TPU는 성능을 크게 향상시킬 뿐만 아니라 새로운 TPU VM 아키텍처로 우수한 사용자 경험을 제공합니다. 무엇보다 Google Cloud는 90% 무탄소 에너지로 운영되는세계 최대 규모의 일반 공개 ML 허브에서 최고 수준의 지속 가능성을 토대로 Cohere의 첨단 ML 학습이 수행되도록 합니다.

Cloud TPU 채택으로 Cohere에서는 LLM 학습을 더 빠르고 경제적이고 민첩하게 수행하고 있습니다. 이를 통해 규모가 더 크고 정확한 LLM을 개발자에게 제공하고, 규모에 상관없이 모든 기업과 개발자가 NLP 기술을 사용할 수 있도록 지원합니다.

LLM 학습의 발전에 대한 자세한 내용은 JAX pjit와 TPUv4를 사용한 언어 모델의 확장 가능한 학습전체 문서를 참조하세요. Cohere의 권장사항과 AI 원칙에 대해 자세히 알아보려면 Open AI 및 AI 21 Labs와 공동 작성한 이 자료를 확인하세요.