저희는 최근 개방형 최첨단 범용 DBRX (LLM)를 도입했습니다. DBRX는 Mosaic AI 트레이닝을 사용하여 훈련, 미세 조정 및 평가되었으며, 3072개의 NVIDIA H100으로 훈련을 확장하고 그 과정에서 12조 이상의 토큰을 처리했습니다.
LLM, 특히 DBRX와 같은 MoE 모델을 교육하는 것은 어렵습니다. 이를 위해서는 많은 인프라, 성능, 과학적 과제를 극복해야 합니다. Mosaic AI 트레이닝은 이러한 문제를 해결하기 위해 의도적으로 구축되었으며, DBRX, MPT 시리즈 모델, 그리고 Ola의 Krutrim, AI2의 OLMo, Dynamo AI의 Dynamo 8B, << 스팬>리필의 LLM-, 등과 같은 다른 많은 LLM의 트레이닝을 통해 실전 테스트를 거쳤습니다.
Mosaic AI 특정 비즈니스 상황, 언어 및 도메인에 맞게 조정되고 주요 비즈니스 사용 사례를 효율적으로 지원할 수 있는 자체 엔터프라이즈 데이터에 대한 사용자 지정 모델을 구축할 수 있는 교육이 오늘부터 Databricks에서 제공됩니다.
이러한 사용자 지정 모델은 특정 도메인(예: 법률, 금융 등)에 적용하거나 리소스가 적은 언어를 처리하도록 교육할 때 특히 유용합니다. 고객들이 커스텀 모델을 사전 교육하거나 오픈 소스 모델을 광범위하게 계속 교육하여 고유한 비즈니스 사용 사례를 강화하는 것을 보았으며, 더 많은 고객이 자체 LLM을 교육하고 소유할 수 있도록 지원할 수 있기를 기대하고 있습니다.
이 블로그 게시물에서는 Mosaic AI Training의 핵심 기능과 이러한 기능이 DBRX의 성공적인 훈련에 어떻게 중요한 역할을 했는지 자세히 설명합니다. 기업 데이터를 활용하는 자체 DBRX급 맞춤형 LLM 구축을 시작하려면 지금 바로 Mosaic AI Training 제품 페이지를 확인하거나 유니티에 문의하세요.
"Mosaic AI 교육 인프라는 OLMo 및 기타 진정한 개방형 대규모 언어 모델을 교육하는 데 매우 중요했습니다." - Dirk Groenveld, Allen 인공 지능 연구소의 수석 소프트웨어 엔지니어.
Mosaic AI 교육 스택
LLM 및 기타 대규모 AI 모델을 학습하려면 하드웨어 장치 드라이버부터 작업 오케스트레이션, 신경망 학습 루프에 이르기까지 수많은 구성 요소를 통합해야 합니다. 이 프로세스는 복잡하고 다양한 전문 지식이 필요합니다. 네트워크 드라이버를 잘못 설정하는 등의 간단한 실수만으로도 훈련 속도가 5배나 느려질 수 있습니다!
이러한 복잡성을 단순화하고 "제대로 작동하는" 경험을 제공하기 위해 Mosaic AI Training은 대규모 분산 훈련의 모든 측면을 처리하는 최적화된 훈련 스택을 제공합니다. 이 스택은 여러 GPU 클라우드 제공업체(AWS, Azure, OCI, Coreweave 등)를 지원하며, NVIDIA CUDA 및 AMD ROCm을 포함한 최신 GPU 드라이버로 구성되고, 핵심 신경망 및 트레이닝 라이브러리 (PyTorch, 메가블록, 컴포저, 스트리밍). 마지막으로, LLM을 훈련, 미세 조정 및 평가하기 위한 실전 테스트를 거친 스크립트가 LLMFoundry에서 제공되므로 고객은 즉시 자체 LLM을 훈련할 수 있습니다.
"Refuel은 특정 사용 사례를 위해 특별히 제작된 맞춤형 LLM의 가치를 잘 알고 있습니다. 새로 출시된 모델인 RefuelLLM-2는 Mosaic AI 교육 인프라에서 훈련되었으며, 비정형 데이터부터 하드웨어 장치 드라이버, 교육 라이브러리 및 작업 오케스트레이션에 이르기까지 수많은 구성 요소를 사용하기 쉬운 하나의 리소스에서 조정할 수 있는 기능을 제공합니다." - Nihit Desai, 공동 창립자 겸 CTO, Refuel
분산된 훈련
수천 개의 GPU에서 모델을 훈련하기 위해 Mosaic AI 훈련은 텐서를 샤딩하고 복제하는 방법을 설명하기 위해 PyTorch 의 간단한 프레임워크인 DTensor를 활용합니다. 모델 트레이닝을 하나의 GPU에서 여러 GPU로 확장할 때는 모델과 최적화 프로그램을 여러 GPU에 복제하는 분산 데이터 병렬(DDP) 트레이닝을 활용하는 경우가 많습니다. DBRX의 경우처럼 단일 GPU의 메모리에 맞지 않는 모델을 훈련할 때는 대신 모델과 옵티마이저를 여러 GPU에 걸쳐 샤딩하고 행렬 곱셈과 같은 연산을 위해 수집하는 완전 샤딩 데이터 병렬(FSDP) 훈련을 사용합니다. 여러 GPU에 걸쳐 수집 및 샤딩하는 각 작업은 머신 수가 증가함에 따라 속도가 느려지므로 고정된 복제본 크기에 걸쳐 모델과 최적화기를 샤딩한 다음 이 복제본을 여러 번 복사하여 전체 클러스터를 채우는 하이브리드 샤딩 데이터 병렬(HSDP) 트레이닝을 사용합니다(그림 3). DTensor를 사용하면 텐서를 샤딩하고 복제하는 방법을 설명하는 장치 메시를 구성하고 이를 PyTorch의 FSDP 구현에 전달할 수 있습니다. DBRX와 같은 MoE 모델의 경우, 훈련 병렬 처리를 개선하여 전문가 병렬 처리를 포함하도록 하여 서로 다른 전문가를 서로 다른 GPU에 배치합니다. 성능을 향상시키기 위해 전문가 병렬 처리의 효율적인 구현인 MegaBlocks 을 활용합니다 . 또한, 긴 컨텍스트 훈련을 지원하기 위해 시퀀스 길이 병렬화를 적용하여 검색 증강 생성(RAG)과 같은 주요 사용 사례에 대한 지원을 크게 개선했습니다. 이러한 모든 기술을 내부에서 병렬 처리를 처리하는 오픈 소스 트레이닝 라이브러리인 Composer에 통합하여 사용자가 모델과 데이터에 집중할 수 있는 간단하고 깔끔한 인터페이스를 제공합니다 .
효율적인 병렬 처리 설정이 주어지면 오픈 소스 스트리밍 라이브러리를 활용하여 훈련 중에 즉시 데이터를 가져옵니다. 스트리밍은 시작 오버헤드와 훈련 오버헤드 없이 1분 미만의 시간으로 수천 대의 머신에서 결정론적 셔플링을 보장합니다. 하드웨어 장애가 발생하면 스트리밍은 중단된 지점부터 결정론적으로 즉시 재개할 수 있습니다. 이러한 결정론은 하드웨어 문제나 재시작에 관계없이 재현 가능한 결과를 보장하는 데 매우 중요합니다. 데이터 셔플링 외에도 최신 LLM 교육에는 커리큘럼 교육이라는 프로세스에서 여러 단계의 교육이 포함되며, 교육이 끝날 때 중요한 고품질 데이터가 업샘플링되는 경우가 많습니다. 몇 가지 parameter 오버라이드만으로 스트리밍은 기본적으로 이와 같은 고급 기능을 지원하여 나중에 데이터 세트 혼합을 트레이닝에 자동으로 업데이트합니다.
분산 체크포인트
하드웨어 장애(이 블로그 게시물의 뒷부분에서 하드웨어 장애에 대해 자세히 설명합니다!)로부터 신속한 복구를 보장하기 위해 Mosaic AI Training은 고도로 최적화된 체크포인트 저장 및 로딩을 구현합니다. 사전 학습 시에는 각 GPU가 모델 파라미터와 최적화 상태의 샤딩된 부분만 저장하거나 로드하는 샤딩된 체 크포인트를 활용합니다. 저장 및 로드 작업을 수백 대의 머신에 분산하여 저장 및 로드 시 차단 시간을 30초 미만으로 최소화함으로써 DBRX와 같이 수천억 개의 파라미터와 체크포인트가 수 테라바이트 크기로 구성된 모델을 학습할 때도 차단 시간을 최소화할 수 있습니다. HSDP를 사용하면 첫 번째 복제본 내에서만 체크포인트를 저장 및 로드하고 필요한 경우 다른 모든 복제본에 브로드캐스트함으로써 이 프로세스를 더욱 최적화할 수 있습니다 (그림 4).
교육 성능 최적화
특히 고가의 하이엔드 GPU 수천 대를 사용하는 LLM 훈련의 규모를 고려할 때, 훈련 기간과 비용을 최소화하려면 LLM 훈련 중에 가능한 최고의 성능을 달성하는 것이 중요합니다. 예를 들어, 3개월에 걸쳐 3072개의 H100 GPU를 사용한 DBRX 트레이닝이 있었습니다. Mosaic AI 트레이닝은 메모리 사용량과 계산을 적극적으로 최적화하여 이러한 SOTA 트레이닝 효율성을 달성합니다.
트레이닝 LLM 의 메모리 사용량에는 모델 파라미터, 옵티마이저 상태, 그라데이션, 활성화가 포함됩니다. 이로 인해 실행 가능한 훈련 배치 크기가 제한되며, 모델이 확장됨에 따라 심각한 성능 병목 현상이 발생할 수 있습니다. 배치 크기가 클수록 매트릭스 곱셈이 커져 GPU 활용도가 높아집니다. 모델 파라미터, 옵티마이저 상태 및 그라데이션을 처리하기 위해 앞서 설명한 대로 PyTorch FSDP 및 HSDP를 활용합니다. 활성화 메모리를 처리하기 위해 선택적 활성화 체크포인트와 압축을 혼합하여 사용합니다. 활성화 체크포인트는 이러한 중간 텐서 중 일부를 저장하는 대신 재계산하므로 추가 계산을 통해 메모리를 절약할 수 있습니다. 이 추가 메모리 덕분에 DBRX 트레이닝에서 더 큰 배치 크기를 사용할 수 있습니다. 활성화 체크포인트를 효율적으로 적용하기 위해 모델 세분성(전체 블록, 관심도 등)과 범위(예: 첫 번째-n 레이어, 범위(m, k) 레이어)를 쉽게 지정할 수 있는 구성 가능한 방식으로 활성화를 체크포인트할 수 있도록 지원하고 있습니다. 활성화 체크포인트와 함께 활성화 압축을 위한 사용자 지정 로우레벨 커널을 개발했습니다. 모든 활성화를 16비트 정밀도로 저장하는 대신, 이러한 활성화를 사용자 지정의 낮은 정밀도 형식으로 압축하여 모델 품질에 미치는 영향을 최소화하면서 활성화 메모리를 크게 줄입니다. 이러한 메모리 최적화를 통해 DBRX와 같은 대형 모델의 경우 메모리 사용 공간을 3배 가까이 줄일 수 있습니다.
Mosaic AI 훈련은 또한 계산을 최적화하여 GPU가 가능한 한 오랫동안 최대 FLOPS로 실행되도록 합니다. 먼저 PyTorch 프로파일러와 커스텀 벤치마크를사용해 모델의 상세한 이벤트 프로파일(CPU, CUDA 커널, NCCL 등)을 생성합니다 . 이러한 결과는 추가 개선에 큰 동기를 부여하여 항상 가장 중요한 병목 현상을 타겟팅할 수 있게 해줍니다. 예를 들어, 플래시어텐션 2, 메가블록 및 기타 커스텀 커널의 고성능 연산자를 활용하여 DBRX의 주요 레이어에서 CPU 및 메모리 병목 현상을 방지합니다. 빠른 드롭리스 MoE를 구현하기 위해 장치 동기화로 인해 CUDA 커널이 차단되지 않도록 GPU 및 CPU 이벤트 순서를 재조정하여 CPU 병목 현상을 방지합니다. 통신 병목 현상을 피하기 위해 유니티는 GPU당 플롭 수를 최대화하는 다른 최적화 방법을 사용합니다. 이러한 작업을 종합하면 다양한 규모에서 표준 트랜스포머와 거의 동일한 속도로 고급 MoE 아키텍처를 훈련할 수 있습니다.
GPU 내결함성
LLM을 트레이닝하려면 수백 개의 노드에서 수천 개의 GPU를 오케스트레이션해야 하므로 인프라 확장 및 안정성 문제가 발생합니다. 이 규모에서는 노드 장애가 매우 흔하게 발생합니다. 성공을 가능하게 한 Mosaic AI 교육의 몇 가지 주요 기능에 대해 설명합니다.
디바이스 수준의 문제를 빠르게 감지하고 해결하기 위해 Mosaic AI Training은 ECC 수정 불가 오류 및 리매핑 오류, GPU 연결 끊김 및 GPU ROM 오류와 같은 일반적인 GPU 오류 (XID 오류)를 감지하는 NVIDIA의 Prometheus 메트릭 익스포터(DCGM Exporter) 를 활용하는 강력한 GPU 경고 세트와 함께 제공됩니다. 또한, GPU에서 보고된 열 및 전력 위반과 같은 추가적인 GPU 성능 문제를 식별하기 위해 GPU 스로틀링을 모니터링하도록 DCGM 운영자 구성을 확장합니다.
Mosaic AI 또한 훈련 중 Mosaic AI 모델 FLOPs 활용도(MFU)가 미리 구성된 임계값 아래로 떨어지는 등 구성 가능한 특정 조건이 충족되면 더 적은 진단 워크로드로 자동 클러스터 스윕을 트리거합니다. 진단 테스트에 실패하면 트레이닝은 Composer의 자동 재개 기능과 분산된 체크포인팅을 사용하여 지연 노드를 묶고 최신 체크포인트부터 자동으로 실행을 재개합니다.
이러한 내결함성 기능의 조합은 DBRX 훈련에 매우 유용했습니다. GPU 용량이 제한되어 있기 때문에 사용 가능한 노드 수가 조금만 감소해도 훈련 작업을 실행하는 데 필요한 모든 리소스를 사용할 수 있는 시간 비율을 나타내는 Google에서 만든 용어인 스케줄링 성능에 큰 영향을 미칠 수 있었기 때문입니 다. Google의 알림 및 수정 시스템은 다양한 GPU 오류를 식별하고 실행이 예약되기 전에 불량 노드를 재활용하여 높은 스케줄링 처리량을 보장했습니다.
네트워크 패브릭 내 결함성
Mosaic AI 교육은 스위치 포트 장애 및 전체 스위치 장애와 같은 GPUDirect 원격 직접 메모리 액세스(GDRDMA) 패브릭 내의 특정 문제를 정확히 찾아낼 수 있는 기능을 제공합니다.
이러한 메트릭을 수집하기 위해 Mosaic AI 교육은 주기적으로 빠른 실행 차단 상태 검사 집합을 예약합니다. 이는 compute 및 통신 벤치마크를 모두 수집하기 위해 노드 집합을 스트레스 테스트하는 카나리아 배포입니다. compute 테스트에는 GPU가 완전히 활용된 워크로드에서 최적으로 실행되는지 확인하기 위한 멀티 GPU CUDA 스트레스 테스트가 포함됩니다. 통신 테스트에는 NVIDIA nccl-test 방법론을 사용한 인접 노드 간의 올-리듀스 NCCL 테스트가 포함됩니다 . Mosaic AI 그런 다음 트레이닝은 이러한 테스트의 메트릭을 다양한 NCCL 페이로드 크기에 대해 사용하고 이상 징후나 장애가 있는 클러스터 노드를 차단합니다.
이러한 네트워크 패브릭 점검을 통 해 3개월의 교육 기간 동안 발생한 여러 스위치 문제를 확인했습니다. 이를 통해 클라우드 제공업체가 스위치 중단을 서비스하는 동안 더 작은 노드 파티션에서 실행을 재개할 수 있었습니다.
그림 6은 네트워크 패브릭 점검이 어떻게 조율되었는지에 대한 대략적인 개요를 보여줍니다:
- Mosaic AI 상태 검사 API는 검사할 준비가 된 대상 노드를 추적하고 대기합니다.
- API가 분산 교육을 생성합니다. 프로세스를 실행하는 모든 프로세스는 테스트를 줄입니다.
- API 집단 통신 및 compute 메트릭을 사용하여 상태 확인을 위한 RPC 콜백을 진행합니다.
- 상태 확인 API는 비정상적인 메트릭이 있는 노드를 코딩합니다.
Experiment 추적
알림 및 GPU 상태 모니터링 외에도 Mosaic AI 트레이닝은 다음과 통합됩니다. MLflow 및 기타 Experiment 트래커와 통합하여 훈련 지표, 훈련 진행 상황, 시스템 지표 및 평가 결과를 실시간으로 추적할 수 있습니다. 이를 사용하여 초기 실험, 상태 확인, 최종 대규모 훈련 실행 등 DBRX 훈련의 모든 단계에서 진행 상황을 모니터링했습니다 (그림 5 참조).
DBRX 트레이닝 진행 상황을 모니터링하고 대략적인 최종 모델 품질을 추정하기 위해 모델이 트레이닝된 전체 12조 개의 토큰에 대해 모자이크 평가 건틀릿을 사용하여 트레이닝된 모든 1T 토큰에 대해 전체 평가 작업을 실행했습니다 . 중간 및 미세 조정된 체크포인트를 저장하여 MLflow Model Registry에 중간 및 조정된 체크포인트를 저장함으로써 평가용 모델의 여러 버전을 쉽게 관리하고 Databricks Model Serving을 통해 배포할 수 있었습니다.
"Databricks Mosaic AI 을 통해 기업 규정 준수와 고객을 위한 책임감 있는 AI에 초점을 맞춘 업계 최고의 기반 모델인 Dynamo8B를 훈련할 수 있었습니다. 구성된 컨테이너 이미지와 즉시 사용 가능한 트레이닝 스크립트 덕분에 몇 주에 달하는 개발 시간을 절약할 수 있었고, Mosaic AI Training의 기본 내장 속도 향상과 온디맨드 GPU 가용성 덕분에 80억 개에 달하는 parameter 다국어 LLM 를 사전 트레이닝하는 데 단 10일이 걸렸습니다." - Eric Lin, Dynamo AI 응용 AI 부문 책임자
나만의 맞춤형 LLM 교육 시작하기
엔터프라이즈 데이터를 활용하는 자체 DBRX 등급 사용자 지정 LLM 구축을 시작하려면 Mosaic AI 교육 제품 페이지를 확인하거나 자세한 내용을 문의하세요.