Amazon SageMaker 모델 병렬 및 데이터 병렬 라이브러리를 사용한 분산 교육 및 효율적인 확장

특히 2022년 5.32월 ChatGPT가 출시된 이후 대규모 언어 모델(LLM)을 위한 분산 딥 러닝 분야에서 엄청난 발전이 있었습니다. LLM은 수십억 또는 심지어 수조 개의 매개변수로 크기가 계속해서 증가하고 있지만 종종 그렇지 않습니다. 메모리 제한으로 인해 GPU와 같은 단일 가속기 장치 또는 ml.p3xlarge와 같은 단일 노드에 적합합니다. LLM을 교육하는 고객은 수백 또는 수천 개의 GPU에 워크로드를 분산해야 하는 경우가 많습니다. 이러한 규모로 훈련을 가능하게 하는 것은 분산 훈련에서 여전히 어려운 과제로 남아 있으며, 대규모 시스템에서 효율적으로 훈련하는 것도 똑같이 중요한 문제입니다. 지난 몇 년 동안 분산 교육 커뮤니티에서는 이러한 문제를 해결하기 위해 XNUMXD 병렬 처리(데이터 병렬 처리, 파이프라인 병렬 처리, 텐서 병렬 처리)와 기타 기술(예: 시퀀스 병렬 처리, 전문가 병렬 처리)을 도입했습니다.

2023년 XNUMX월, 아마존은 출시를 발표했습니다. SageMaker 모델 병렬 라이브러리 2.0 (SMP)는 대규모 모델 훈련에서 최고의 효율성을 달성합니다. SageMaker 분산 데이터 병렬 처리 라이브러리 (SMDDP). 이번 릴리스는 1.x의 중요한 업데이트입니다. 이제 SMP가 오픈 소스 PyTorch와 통합되었습니다. 완전히 샤딩된 데이터 병렬 (FSDP) API는 대규모 모델을 훈련할 때 친숙한 인터페이스를 사용할 수 있도록 하며 다음과 호환됩니다. 변압기 엔진 (TE), 처음으로 FSDP와 함께 텐서 병렬성 기술을 잠금 해제합니다. 릴리스에 대해 자세히 알아보려면 다음을 참조하세요. Amazon SageMaker 모델 병렬 라이브러리는 이제 PyTorch FSDP 워크로드를 최대 20% 가속화합니다..

이번 포스팅에서는 성능상의 이점을 살펴보겠습니다. 아마존 세이지 메이커 (SMP 및 SMDDP 포함) 및 라이브러리를 사용하여 SageMaker에서 대규모 모델을 효율적으로 교육하는 방법을 알아보세요. 최대 4개 인스턴스의 ml.p24d.128xlarge 클러스터에 대한 벤치마크와 Llama 16 모델에 대한 bfloat2을 사용한 FSDP 혼합 정밀도를 통해 SageMaker의 성능을 시연합니다. SageMaker의 거의 선형적인 확장 효율성을 시연한 후 최적의 처리량을 위해 각 기능의 기여도를 분석하고 텐서 병렬성을 통해 최대 32,768까지 다양한 시퀀스 길이를 사용하는 효율적인 교육으로 마무리됩니다.

SageMaker를 사용한 선형에 가까운 확장

LLM 모델의 전체 교육 시간을 줄이려면 노드 간 통신 오버헤드를 고려할 때 대규모 클러스터(수천 개의 GPU)로 확장할 때 높은 처리량을 유지하는 것이 중요합니다. 이 게시물에서는 SMP와 SMDDP를 모두 호출하는 p4d 인스턴스에서 강력하고 선형에 가까운 확장(고정된 전체 문제 크기에 대해 GPU 수를 변경하여) 효율성을 보여줍니다.

이 섹션에서는 SMP의 거의 선형적인 확장 성능을 보여줍니다. 여기서는 고정 시퀀스 길이 2, 집단 통신을 위한 SMDDP 백엔드, TE 지원, 글로벌 배치 크기 7백만, p13d 노드 70~4,096개를 사용하여 다양한 크기(4B, 16B 및 128B 매개변수)의 Llama 4 모델을 훈련합니다. . 다음 표에는 최적의 구성 및 교육 성능(초당 모델 TFLOP)이 요약되어 있습니다.

모형 크기	노드 수	TFLOP*	SDP*	티*	짐을 내리다*	확장 효율성
7B	16	136.76	32	1	N	100.0%
	32	132.65	64	1	N	97.0%
	64	125.31	64	1	N	91.6%
	128	115.01	64	1	N	84.1%
13B	16	141.43	32	1	Y	100.0%
	32	139.46	256	1	N	98.6%
	64	132.17	128	1	N	93.5%
	128	120.75	128	1	N	85.4%
70B	32	154.33	256	1	Y	100.0%
	64	149.60	256	1	N	96.9%
	128	136.52	64	2	N	88.5%

*주어진 모델 크기, 시퀀스 길이 및 노드 수에서 다양한 sdp, tp 및 활성화 오프로딩 조합을 탐색한 후 전역적으로 최적의 처리량 및 구성을 보여줍니다.

앞의 표에는 샤딩된 데이터 병렬(sdp) 수준(일반적으로 전체 샤딩 대신 FSDP 하이브리드 샤딩을 사용하며 자세한 내용은 다음 섹션 참조), 텐서 병렬(tp) 수준 및 활성화 오프로딩 값 변경에 따른 최적의 처리량 수치가 요약되어 있습니다. SMDDP와 함께 SMP에 대한 거의 선형적인 확장을 보여줍니다. 예를 들어 Llama 2 모델 크기가 7B이고 시퀀스 길이가 4,096인 경우 전체적으로 97.0, 91.6, 84.1개 노드에서 각각 16%, 32%, 64%(128개 노드 기준)의 확장 효율성을 달성합니다. 확장 효율성은 다양한 모델 크기에 걸쳐 안정적이며 모델 크기가 커질수록 약간 증가합니다.

SMP 및 SMDDP는 2,048 및 8,192와 같은 다른 시퀀스 길이에 대해서도 유사한 스케일링 효율성을 보여줍니다.

SageMaker 모델 병렬 라이브러리 2.0 성능: Llama 2 70B

LLM 커뮤니티의 빈번한 최신 성능 업데이트와 함께 모델 크기는 지난 몇 년 동안 계속해서 증가해 왔습니다. 이 섹션에서는 고정 모델 크기 2B, 시퀀스 길이 70, 글로벌 배치 크기 4,096만 개를 사용하여 Llama 4 모델에 대한 SageMaker의 성능을 보여줍니다. 이전 표의 전역적으로 최적 구성 및 처리량(SMDDP 백엔드, 일반적으로 FSDP 하이브리드 샤딩 및 TE 포함)과 비교하기 위해 다음 표는 분산 백엔드(NCCL 및 SMDDP)에 대한 추가 사양을 사용하여 다른 최적 처리량(잠재적으로 텐서 병렬 처리 포함)으로 확장됩니다. , FSDP 샤딩 전략(전체 샤딩 및 하이브리드 샤딩) 및 TE 활성화 여부(기본값).

모형 크기	노드 수	테플롭스				TFLOP #3 구성			기준선에 비해 TFLOP 개선
.	.	NCCL 전체 샤딩: #0	SMDDP 전체 샤딩: #1	SMDDP 하이브리드 샤딩: #2	TE를 사용한 SMDDP 하이브리드 샤딩: #3	SDP*	티*	짐을 내리다*	#0 → #1	#1 → #2	#2 → #3	#0 → #3
70B	32	150.82	149.90	150.05	154.33	256	1	Y	-0.6 %	0.1%	2.9%	2.3%
	64	144.38	144.38	145.42	149.60	256	1	N	0.0%	0.7%	2.9%	3.6%
	128	68.53	103.06	130.66	136.52	64	2	N	50.4%	26.8%	4.5%	99.2%

*주어진 모델 크기, 시퀀스 길이 및 노드 수에서 다양한 sdp, tp 및 활성화 오프로딩 조합을 탐색한 후 전역적으로 최적의 처리량 및 구성을 보여줍니다.

SMP 및 SMDDP의 최신 릴리스는 기본 PyTorch FSDP, 확장되고 더욱 유연한 하이브리드 샤딩, 변환기 엔진 통합, 텐서 병렬 처리 및 최적화된 전체 수집 작업을 포함한 여러 기능을 지원합니다. SageMaker가 LLM에 대한 효율적인 분산 교육을 달성하는 방법을 더 잘 이해하기 위해 SMDDP 및 다음 SMP의 점진적인 기여를 살펴보겠습니다. 핵심 기능:

FSDP 전체 샤딩을 통해 NCCL보다 SMDDP 향상
FSDP 전체 샤딩을 하이브리드 샤딩으로 대체하여 통신 비용을 줄여 처리량 향상
텐서 병렬 처리가 비활성화된 경우에도 TE를 통해 처리량을 더욱 향상
낮은 리소스 설정에서는 활성화 오프로딩을 통해 높은 메모리 부족으로 인해 실행 불가능하거나 매우 느린 교육을 활성화할 수 있습니다.

FSDP 전체 샤딩: NCCL보다 SMDDP 향상

이전 표에 표시된 것처럼 모델이 FSDP로 완전히 샤딩되면 NCCL(TFLOPs #0) 및 SMDDP(TFLOPs #1) 처리량은 32개 또는 64개 노드에서 비슷하지만 NCCL에서 SMDDP로 50.4%라는 엄청난 향상이 있습니다. 128노드에서.

더 작은 모델 크기에서는 SMDDP가 통신 병목 현상을 효과적으로 완화할 수 있기 때문에 더 작은 클러스터 크기부터 시작하여 NCCL에 비해 SMDDP의 일관되고 중요한 개선이 관찰됩니다.

통신 비용 절감을 위한 FSDP 하이브리드 샤딩

SMP 1.0에서 우리는 샤딩된 데이터 병렬 처리, Amazon 사내에서 제공하는 분산 교육 기술 마이크 기술. SMP 2.0에서는 FSDP 전체 샤딩의 경우와 같이 모든 학습 GPU 대신 모델이 GPU 하위 집합 간에 샤딩될 수 있도록 하는 확장 가능하고 보다 유연한 하이브리드 샤딩 기술인 SMP 하이브리드 샤딩을 도입합니다. GPU당 메모리 제약 조건을 충족하기 위해 전체 클러스터에 걸쳐 샤딩할 필요가 없는 중간 크기 모델에 유용합니다. 이로 인해 클러스터에는 둘 이상의 모델 복제본이 있고 각 GPU는 런타임에 더 적은 수의 피어와 통신하게 됩니다.

SMP의 하이브리드 샤딩을 사용하면 메모리 부족 문제가 없는 가장 작은 샤드 수준부터 전체 클러스터 크기(전체 샤딩과 동일)까지 더 넓은 범위에 걸쳐 효율적인 모델 샤딩이 가능합니다.

다음 그림은 단순화를 위해 tp = 1에서 sdp에 대한 처리량 의존성을 보여줍니다. 이전 표의 NCCL 또는 SMDDP 전체 샤딩에 대한 최적의 tp 값과 반드시 동일할 필요는 없지만 수치는 상당히 비슷합니다. NCCL과 SMDDP 모두에 적용 가능한 128노드의 대규모 클러스터 크기에서 전체 샤딩에서 하이브리드 샤딩으로 전환하는 것의 가치를 명확하게 검증합니다. 더 작은 모델 크기의 경우 하이브리드 샤딩을 통한 중요한 개선은 더 작은 클러스터 크기에서 시작되며 그 차이는 클러스터 크기에 따라 계속 증가합니다.

TE를 통한 개선

TE는 NVIDIA GPU에서 LLM 교육을 가속화하도록 설계되었습니다. p8d 인스턴스에서 지원되지 않기 때문에 FP4을 사용하지 않음에도 불구하고 p4d에서 TE를 사용하면 여전히 상당한 속도 향상을 볼 수 있습니다.

SMDDP 백엔드로 교육된 MiCS 외에도 TE는 텐서 병렬 처리가 비활성화된 경우에도(텐서 병렬 수준은 128) 모든 클러스터 크기(유일한 예외는 1개 노드의 전체 샤딩)에서 처리량을 일관되게 향상시킵니다.

더 작은 모델 크기 또는 다양한 시퀀스 길이의 경우 TE 부스트는 약 3~7.6% 범위에서 안정적이고 사소하지 않습니다.

낮은 리소스 설정에서 활성화 오프로딩

낮은 리소스 설정(적은 수의 노드 지정)에서 활성화 검사점이 활성화되면 FSDP는 높은 메모리 압력(또는 최악의 경우 메모리 부족)을 경험할 수 있습니다. 메모리로 인해 병목 현상이 발생하는 시나리오의 경우 활성화 오프로드를 켜는 것은 잠재적으로 성능을 향상시키는 옵션입니다.

예를 들어 이전에 살펴본 것처럼 모델 크기가 2B이고 시퀀스 길이가 13인 Llama 4,096는 활성화 체크포인트를 사용하고 활성화 오프로딩 없이 최소 32개의 노드로 최적으로 훈련할 수 있지만 16개로 제한되면 활성화 오프로딩으로 최고의 처리량을 달성합니다. 노드.

긴 시퀀스로 훈련 가능: SMP 텐서 병렬성

긴 대화와 맥락을 위해서는 더 긴 시퀀스 길이가 필요하며 LLM 커뮤니티에서 더 많은 관심을 받고 있습니다. 따라서 다음 표에서는 다양한 긴 시퀀스 처리량을 보고합니다. 표는 2~2,048의 다양한 시퀀스 길이를 사용하여 SageMaker에서 Llama 32,768 교육을 위한 최적의 처리량을 보여줍니다. 시퀀스 길이가 32,768이면 글로벌 배치 크기가 32만 개인 4개 노드에서는 기본 FSDP 교육이 불가능합니다.

.	.	.	테플롭스
모형 크기	시퀀스 길이	노드 수	기본 FSDP 및 NCCL	SMP 및 SMDDP	SMP 개선
7B	2048	32	129.25	138.17	6.9%
	4096	32	124.38	132.65	6.6%
	8192	32	115.25	123.11	6.8%
	16384	32	100.73	109.11	8.3%
	32768	32	NA	82.87	.
13B	2048	32	137.75	144.28	4.7%
	4096	32	133.30	139.46	4.6%
	8192	32	125.04	130.08	4.0%
	16384	32	111.58	117.01	4.9%
	32768	32	NA	92.38	.
*: 최대	.	.	.	.	8.3%
*: 중앙값	.	.	.	.	5.8%

클러스터 크기가 크고 고정된 글로벌 배치 크기가 제공되는 경우 일부 모델 교육은 기본 제공 파이프라인 또는 텐서 병렬 처리 지원이 부족하여 기본 PyTorch FSDP를 사용하여 실행 불가능할 수 있습니다. 앞의 표에서 전역 배치 크기가 4만 개, 노드가 32개, 시퀀스 길이가 32,768인 경우 GPU당 유효 배치 크기는 0.5입니다(예: 배치 크기가 2인 tp = 1). 텐서 병렬성.

결론

이 게시물에서는 NCCL에 대한 SMDDP 향상, 전체 샤딩 대신 유연한 FSDP 하이브리드 샤딩, TE 통합 및 텐서 병렬 처리 활성화와 같은 여러 주요 기능에 기여하는 p4d 인스턴스에서 SMP 및 SMDDP를 사용한 효율적인 LLM 교육을 시연했습니다. 긴 시퀀스 길이. 다양한 모델, 모델 크기 및 시퀀스 길이를 사용하여 다양한 설정에 대해 테스트한 후 SageMaker에서 최대 128개의 p4d 인스턴스까지 선형에 가까운 강력한 확장 효율성을 보여줍니다. 요약하자면, SageMaker는 LLM 연구원 및 실무자를 위한 강력한 도구로 계속해서 사용되고 있습니다.

자세한 내용은 다음을 참조하십시오. SageMaker 모델 병렬 처리 라이브러리 v2, 또는 다음 주소로 SMP 팀에 문의하세요. sm-model-parallel-feedback@amazon.com.