Microsoft DeepSpeed 팀의 공식 중국어 블로그(공식 영어 블로그에서 번역됨: DeepSpeed Ulysses: Extreme Long Sequence Transformer Models 훈련 활성화를 위한 시스템 최적화). Microsoft DeepSpeed 팀의 공식 Zhihu 계정(zhihu.com/people/deepspeed)에서 재인쇄되었습니다 .
생성 AI부터 과학 연구 모델까지, 긴 시퀀스 훈련이 매우 중요해지고 있습니다. 생성 AI 분야에서 대화형 AI, 긴 문서 요약, 비디오 생성 등의 작업은 모두 공간적, 시간적 수준에서 긴 맥락에 대한 추론이 필요합니다. 예를 들어 음성, 이미지 및 파형을 동시에 처리하는 모델과 같은 다중 모드 기본 모델에는 매우 긴 시퀀스가 있는 고차원 입력에 대한 긴 컨텍스트 추론이 필요합니다. 마찬가지로 장과 책 수준의 요약(수만 또는 수십만 단어)도 대화형 AI 및 요약 작업에서 매우 중요합니다.
긴 시퀀스는 구조 생물학, 의료, 기후 및 날씨 예측, 거대분자 시뮬레이션을 더 잘 이해할 수 있는 문을 열어주는 과학적 AI에도 중요합니다. 예를 들어, 유전자 서열에 대한 대규모 언어 모델을 훈련함으로써 매우 긴 서열(인간 게놈은 64억 염기쌍)을 사용하여 게놈 진화 패턴을 학습할 수 있는 언어 모델을 만들 수 있습니다. 의료 분야에서 모든 환자 치료 기록을 기반으로 한 진단 예측 모델에는 매우 긴 일련의 컨텍스트가 필요합니다.
생성 AI 및 과학 AI에서 긴 시퀀스 길이의 중요성이 점차 커지고 있지만 기존의 대규모 모델 훈련 시스템과 기본 병렬 기술(데이터, 텐서, 파이프라인, 시퀀스 병렬 처리)은 효율적인 긴 시퀀스 길이를 지원할 수 없습니다. 기존 병렬 접근 방식에는 두 가지 주요 과제가 있습니다. 첫째, 데이터, 텐서, 파이프라인 등 기존 병렬 방법으로는 시퀀스 차원 확장 문제를 해결할 수 없습니다. 둘째, 기존 순차 병렬 방식은 메모리 통신의 비효율성으로 인해 효율성이 충분하지 않습니다. 게다가 기존 방법은 사용 편의성이 부족하고 침입적이고 복잡하며 오류가 발생하기 쉬운 코드 리팩토링이 필요합니다.
이러한 문제를 해결하기 위해 매우 긴 시퀀스 길이의 효율적인 확장성을 지원하기 위한 간단하고 사용하기 쉽고 효율적인 방법인 DeepSpeed-Ulysses(또는 매우 긴 소설인 Ulysses)를 발표하게 된 것을 기쁘게 생각합니다.
DeepSpeed-Ulysses는 시퀀스 차원의 각 샘플을 참여 GPU로 분할합니다. 그런 다음 attention이 계산되기 전에 분할 쿼리(Q), 키(K) 및 값(V)에 대해 전체 통신 작업을 수행하여 각 GPU가 완전한 시퀀스를 수신하지만 attention을 위해서만 Non- 힘 머리의 부분 집합이 겹쳐집니다. 이를 통해 참여하는 GPU가 서로 다른 주의 헤드를 병렬로 계산할 수 있습니다. 마지막으로 DeepSpeed-Ulysses는 또 다른 all-to-all을 사용하여 시퀀스 차원에서 다시 분할하는 동안 어텐션 헤드에 대한 결과를 수집합니다.
DeepSpeed-Ulysses의 주요 기능과 이 블로그에 게시된 구현은 다음과 같습니다.
기존 시스템과 비교하여 시퀀스 길이가 4배 증가하여 100만 개 이상의 토큰으로 시퀀스 학습을 지원합니다 .
통신은 기존 시스템에 비해 10배 이상 감소하여 처리량 이 최대 2.5배 향상 되고 GPU당 175TFlops 이상의 지속적인 처리량(하드웨어 최고치의 54% 이상)을 달성합니다.
완전한 일반 주의: DeepSpeed 시퀀스는 조밀하고 희박한 주의를 병렬로 지원하며 FlashAttention v2와 같은 효율적인 주의 구현과 함께 작동합니다.
대규모 모델 훈련 지원: DeepSpeed 시퀀스 병렬 처리는 큰 시퀀스 길이를 지원할 뿐만 아니라 ZeRO-3과 함께 사용하여 큰 모델 크기를 지원할 수도 있습니다.
사용 및 마이그레이션이 쉽고 기존 교육 프레임워크에 대한 코드 변경 요구 사항이 최소화됩니다.
다음 장에서는 DeepSpeed-Ulysses의 핵심 설계, 통신 복잡성 분석, 실험 평가, 기존 작업과의 비교를 자세히 논의하고 사용성 및 사용 지침을 보여줍니다.
그림 1: DeepSpeed 시퀀스 병렬(DeepSpeed-Ulysses) 설계
그림 1은 DeepSpeed-Ulysses의 핵심 설계를 보여줍니다. 알려진 Transformer 아키텍처와 마찬가지로 이 디자인은 P개의 사용 가능한 장치에 걸쳐 분할된 N개의 입력 시퀀스로 구성됩니다. 각 로컬 N/P 파티션은 쿼리(Q), 키(K) 및 값(V) 임베딩으로 투영됩니다. 다음으로, (QKV) 임베딩은 참여 컴퓨팅 장치 간 고도로 최적화된 all-to-all Collective를 통해 글로벌 QKV 수집을 수행합니다. all-to-all 세트 이후 각 헤드의 어텐션 계산 형식은 다음과 같습니다. 
어텐션 계산 후 또 다른 all-to-all 세트는 어텐션 계산의 출력 컨텍스트 텐서를 병렬 시퀀스(N/P)로 변환합니다. 변환기 모델 계층의 나머지 모듈(MLP MatMul, 계층 정규화 등)의 후속 작업입니다.
DeepSpeed-Ulysses는 누적 통신량이 적고 시퀀스 병렬성이 증가함에 따라 확장성이 향상된다는 점에서 기존의 다른 긴 시퀀스 방법과 다릅니다.
노드 내 NVSwitch 상호 연결 및 노드 간 팻 트리 IB 토폴로지를 갖춘 최신 클러스터에서 집계 메시지 크기가 M인 전체 전송의 경우 P GPU의 각 링크로 전송되는 트래픽은 M/P입니다. 히든 레이어 크기 h, 시퀀스 길이 N, 병렬도 P인 Transformer 모델의 경우 DeepSpeed Sequence Parallel은 어텐션 계산 전 QKV 투영에 대해 총 메시지 크기 3Nh로 all-to-all 작업을 수행하며, 어텐션 계산 후 출력 컨텍스트 투영은 Nh 크기의 또 다른 전체 작업을 수행합니다. 따라서 DeepSpeed 순차 병렬화의 링크당 총 트래픽은 4Nh/P(또는 O(N/P) 복잡도) 입니다 . N과 P가 비례적으로 증가할 때 이 트래픽 양은 일정하다는 점은 주목할 가치가 있습니다.
이에 비해 Megatron-LM과 같은 기존 방식은 P에 관계없이 N이 선형적으로 증가함에 따라 통신량이 선형적으로 증가하여 O(N)의 통신 복잡도를 초래합니다 . 예를 들어 Megatron-LM은 각 Transformer 모델 계층에 대해 Nh 크기의 두 가지 전체 수집 작업과 Nh 크기의 산란 감소 작업을 두 번 수행합니다. 그러나 P >> 1일 때 M 크기의 각각의 집합 및 감소-산란 비용은 M/P가 아닌 여전히 M입니다. 따라서 Megatron-LM 시퀀스 병렬화는 링크당 4Nh 의 트래픽 볼륨을 발생시키며 이는 DeepSpeed 시퀀스 병렬화보다 P배 더 큽니다. 이를 통해 DeepSpeed Sequence Parallel은 매우 긴 시퀀스 훈련을 지원하는 동시에 훨씬 더 높은 훈련 효율성을 달성할 수 있습니다. 우리의 실험적 평가 결과는 이 이론적 분석과 일치합니다.
DeepSpeed 분산 주의 모듈의 구현은 밀도가 높은 버전과 희박한 버전 모두에서 self-attention, cross-attention 및 causal attention과 같은 모든 유형의 attention을 지원할 뿐만 아니라 다양한 긴 시퀀스를 지원할 만큼 충분히 일반적입니다. 로컬 주의 수준 다양한 버전의 FlashAttention과 같은 최적화된 커널.
DeepSpeed-Ulysses의 다용성은 핵심 설계의 모듈식 특성, 즉 주의 중심 순차 병렬 설계에서 비롯됩니다. Attention 계산 전 시퀀스 병렬성은 N/P로 분할하고 Attention 계산은 헤드별 병렬성으로 각 헤드의 Attention은 모두 유지되지만 헤드 수가 적으므로 Attention 계산을 사용할 수 있습니다. 조밀한 주의(dense attention) 및 다양한 형태의 희박한 주의(sparse attention)와 같은 모든 유형의 주의 메커니즘 대체.
DeepSpeed 시퀀스 병렬 처리는 더 긴 시퀀스로 훈련할 때 활성화 메모리 사용량을 줄이지만 모델 상태의 메모리 공간에는 영향을 미치지 않습니다. 따라서 시퀀스 길이가 큰 대규모 언어 모델 훈련을 지원하기 위해 DeepSpeed 시퀀스 병렬 처리와 ZeRO-3의 통합을 구현했습니다.
ZeRO Redundancy Optimizer Stage 3(ZeRO-3)은 대규모 모델 학습을 위한 메모리 최적화 기술입니다. 모델 상태가 데이터 병렬 수준에 걸쳐 복제되는 기존 신경망 데이터 병렬 훈련과 달리 ZeRO-3는 데이터 병렬 수준 간에 모델 상태를 분할하여 메모리 사용을 최적화합니다. 그러나 시퀀스 병렬성을 사용하는 경우 훈련 데이터를 배치(샘플) 및 시퀀스 차원에서 고려할 수 있으며 관련 병렬 그룹을 더 큰 그룹으로 결합하여 ZeRO 병렬성을 달성할 수 있습니다.
따라서 우리는 ZeRO-3 파티셔닝을 데이터 병렬 및 시퀀스 병렬 수준의 조합으로 확장합니다. 즉, DeepSpeed 시퀀스 병렬화에서 ZeRO는 모델 상태를 시퀀스와 데이터 병렬 그룹 간에 분할하고 필요할 때 각 순위 분할(allgather)을 수집합니다. 마찬가지로, 매개변수 업데이트에 대한 데이터 병렬 수준과 시퀀스 병렬 수준 사이의 기울기가 줄어듭니다. ZeRO는 시퀀스와 데이터 차원 전반에 걸쳐 엄청난 메모리 절약을 가능하게 하며, 큰 시퀀스 길이뿐만 아니라 큰 모델까지 확장할 수 있습니다.
우리는 최대 64개의 A100 GPU(40GB 비디오 메모리)를 사용하여 많은 NLP 작업에 사용되는 기본 모델인 GPT에서 DeepSpeed-Ulysses를 평가했습니다. 우리의 평가는 i) 시퀀스 길이 확장성, ii) 조밀한 Attention 처리량 및 기존 시스템과의 비교, iii) Sparse Attention 처리량 및 기존 시스템과의 비교, iv) DeepSpeed 시퀀스 병렬 처리에 대한 융합 연구의 네 가지 측면으로 나뉩니다. 다음으로 각 항목별 평가 결과를 논의하고 제시한다.
첫 번째 실험 세트는 12억 개의 매개변수 GPT 모델에서 시퀀스 길이를 100만 개의 토큰으로 확장하는 것입니다. 이 평가 결과는 그림 2에 나와 있습니다. DeepSpeed 시퀀스 병렬화를 사용하면 GPU 수가 증가함에 따라 시퀀스 길이가 선형적으로 증가할 수 있으며, 시퀀스 길이는 GPU 수와 선형 비례 관계를 유지하여 적절한 GPU 수에서 유사한 컴퓨팅 처리량을 유지합니다.
그림 2: 다양한 시퀀스 길이와 GPU 수를 사용한 DeepSpeed 시퀀스 병렬 확장성 향상 평가.
다음으로 300억 개의 매개변수가 있는 Dense Attention 모델에서 DeepSpeed 시퀀스 병렬성을 평가하고 이를 64개의 A100 GPU에서 Megatron 시퀀스 병렬성과 비교합니다. 이러한 평가 결과는 그림 3에 나와 있습니다.
다양한 시퀀스 길이에서 Megatron-LM과 DeepSpeed 시퀀스 병렬성의 성능을 비교했습니다. 평가를 위해 DeepSpeed Sequence Parallel과 Megatron-LM의 각각 최고의 성능(처리량 또는 TFLOP로 측정)을 이끌어내는 시퀀스 길이-배치 크기 조합을 선택했으며, 이를 최적(배치 크기-시퀀스 길이) 구성이라고 합니다. . DeepSpeed 시퀀스 병렬 처리의 경우 항상 64의 ZeRO 병렬 처리를 사용합니다.
그림 3은 DeepSpeed 시퀀스 병렬성이 동일한 시퀀스 길이에서 항상 Megatron-LM보다 뛰어난 성능을 보인다는 것을 보여줍니다. 또한 DeepSpeed 시퀀스 병렬화는 Megatron-LM보다 더 긴 시퀀스를 실행할 수 있습니다. DeepSpeed 시퀀스 병렬화의 성능 이점은 두 가지 측면에 있습니다:
(1) ZeRO-3의 메모리 최적화와 결합된 DeepSpeed 시퀀스 병렬화는 더 많은 샘플을 수용할 수 있어 처리량이 향상됩니다
. LM 시퀀스 병렬화 전체 수집 통신, DeepSpeed 시퀀스 병렬은 보다 효율적인 전체 통신을 사용합니다.
그림 3: 300억 개의 매개변수 모델에 대한 DeepSpeed 및 Megatron LM 순차 병렬 처리의 Dense attention 평가.
마찬가지로 300억 개의 매개변수 희소 주의 모델에서 DeepSpeed 시퀀스 병렬성을 평가하고 이를 Megatron 시퀀스 병렬성과 비교합니다. 평가 결과는 그림 4에 나와 있습니다. 희소 주의 실험의 결과는 밀도 주의 실험의 결과와 유사합니다. DeepSpeed 순차 병렬 처리의 처리량 성능이 Megatron-LM에 비해 2배 이상 향상되는 것을 관찰했습니다. 메모리를 절약함으로써 ZeRO-3와 결합된 DeepSpeed 시퀀스 병렬 처리는 Megatron-LM보다 4배 더 긴 시퀀스 길이로 확장할 수 있습니다.
DeepSpeed 시퀀스 병렬 처리는 동일한 시퀀스 길이에서 Megatron-LM보다 지속적으로 성능이 뛰어납니다. 실제로 현재 DeepSpeed 처리량은 Local Sparse Attention 구현으로 인해 병목 현상이 발생하므로 시퀀스 길이가 길어질수록 DeepSpeed 처리량은 감소합니다. 우리는 향후 Local Sparse Attention 구현의 성능이 향상됨에 따라 DeepSpeed와 Megatron 간의 성능 격차가 시퀀스 길이가 길어질수록 더욱 커질 것으로 예상합니다.
그림 4: 300억 개의 매개변수 모델에 대한 DeepSpeed 및 Megatron LM 순차 병렬 처리의 희소 주의 평가.
마지막으로 그림 5는 4로 설정된 시퀀스 병렬성을 사용하여 32K 시퀀스 길이의 8개 A100 GPU에서 1억 3천만 개의 매개변수 GPT 모델의 수렴을 보여줍니다. DeepSpeed 시퀀스 병렬화를 위해 다양한 ZeRO 단계를 사용하여 수렴 평가를 수행했습니다. DeepSpeed 시퀀스 병렬화는 긴 시퀀스 Transformer 모델의 학습을 달성하는 데 사용되는 순수 시스템 최적화 기술이므로 학습 모델의 품질에 (부정적인) 영향을 미치지 않으며 그림 5와 같이 실험을 통해 검증되었습니다.
그림 5: 다양한 ZeRO 메모리 최적화 단계를 사용한 DeepSpeed 시퀀스 병렬성의 수렴 평가.
DeepSpeed-Ulysses를 코드에 통합하려면 몇 가지 간단한 코드 변경만 필요합니다. 활성화하는 예는 다음과 같습니다.
순차 병렬 처리(예: Megatron-LM)를 지원하는 다른 라이브러리와 비교할 때 DeepSpeed-Ulysses는 모델 리팩토링이 필요하지 않습니다. DeepSpeed-Ulysses는 Megatron-DeepSpeed 코드 기반과 완전히 통합되고 테스트되었습니다. 이는 이미 이 코드베이스를 사용하여 대규모 언어 모델을 훈련하고 있다면 DeepSpeed-Ulysses를 원활하게 사용하여 매우 긴 시퀀스가 있는 모델을 훈련할 수 있음을 의미합니다. 
DeepSpeed-Ulysses를 발표하게 되어 기쁘게 생각합니다. 해당 코드는 DeepSpeed GitHub를 통해 제공됩니다. 자세한 사용법 튜토리얼은 DeepSpeed 튜토리얼 페이지에서 제공됩니다.
우리는 긴 상황 창이 더 이상 제한이 아닐 때 모든 종류의 혁신을 공동으로 추진하기 위해 모든 형태의 기여와 협력을 환영합니다. DeepSpeed-Ulysses는 대규모 AI 훈련 및 추론을 위한 더 큰 DeepSpeed 생태계의 일부입니다. 모든 DeepSpeed 기술과 혁신에 대한 자세한 내용을 보려면 당사 웹사이트를 방문하고 X(Twitter)(영어, 일본어) 및 중국어 Zhihu를 팔로우하세요.
우리는 대학, 연구소, 기업과 협력할 의향이 있습니다. 이와 같은 요청(및 GitHub에서 수행하기에 적합하지 않은 기타 요청)의 경우 [email protected]으로 직접 이메일을 보내주세요.
Microsoft DeepSpeed 팀의 공식 Zhihu에서 재인쇄됨
편집자丨Shao Kejia
관련 읽기 | 관련 읽기
충격적인 출시, 모두를 위한 초능력 잠금 해제 | DeepSpeed-Chat은 오픈 소스입니다!
Kaiyuanshe 소개
카이위안 소사이어티(Kaiyuan Society)는 2014년 설립되었으며, 오픈소스 기여에 자발적으로 참여하는 개인 회원들로 구성되어 있으며 "기여, 합의, 공동 거버넌스"의 원칙을 바탕으로 형성되었으며 항상 공급업체 중립성의 특성을 유지합니다. , 공공 복지 및 비영리. "국제 통합, 커뮤니티 개발 및 프로젝트 인큐베이션"을 사명으로 하는 "오픈 소스 거버넌스, 오픈 소스 커뮤니티 컨소시엄"을 최초로 사용합니다. Kaiyuan Society는 오픈 소스를 지원하는 지역 사회, 기업 및 정부 관련 부서와 적극적으로 긴밀히 협력하며 "중국을 기반으로 세계에 기여"라는 비전을 가지고 건강하고 지속 가능한 오픈 소스 생태계를 조성하고 중국 오픈 소스를 촉진하는 것을 목표로 합니다. 글로벌 오픈 소스 시스템에서 활발한 플레이어가 되기 위한 소스 커뮤니티 참여 및 기여자.
2017년 Open Source Society는 ASF와 같은 최고의 국제 오픈소스 재단의 거버넌스 모델에 따라 운영되면서 전적으로 개인 회원으로 구성되도록 변모했습니다. 지난 9년 동안 수만 명의 오픈소스 인력을 연결하고, 수천 명의 커뮤니티 회원과 자원봉사자, 수백 명의 국내외 강사를 모았고, 수백 명의 후원자, 미디어, 커뮤니티 파트너와 협력했습니다.
