Amazon Search M5가 AWS Trainium을 사용하여 LLM 교육 비용을 30% 절감한 방법 | 아마존 웹 서비스

수십 년 동안 Amazon은 기계 학습(ML)을 개척하고 혁신하여 고객에게 즐거운 경험을 선사해 왔습니다. Amazon은 초창기부터 도서 추천, 검색, 사기 탐지 등 다양한 사용 사례에 ML을 사용해 왔습니다. 업계의 나머지 부분과 마찬가지로 가속화된 하드웨어의 발전으로 Amazon 팀은 신경망과 딥 러닝(DL)을 사용하여 모델 아키텍처를 추구할 수 있었습니다.

Amazon Search 내의 M5 프로그램은 Amazon을 위한 검색 학습 전략을 보유하고 텍스트, 이미지, 비디오와 같은 다국어, 다중 로케일, 다중 엔터티, 다중 작업 및 다중 모드에 걸쳐 대규모 모델을 구축합니다. M5 프로그램은 비용 최적화에 대한 엄격한 제어를 유지하면서 Amazon 전역의 수백 개의 ML 팀에 범용 임베딩 및 대규모 기반 모델을 제공해 왔습니다. 이를 달성하기 위해 M5 팀은 정기적으로 비용 절감을 위한 새로운 기술을 평가합니다.

많은 ML 조직과 마찬가지로 가속기는 DL 교육 및 추론을 가속화하는 데 주로 사용됩니다. AWS가 첫 번째 릴리스로 특수 제작된 액셀러레이터를 출시했을 때 AWS 인 페렌 시아 2020년 M5 팀은 빠르게 이를 활용하여 프로덕션 워크로드를 보다 효율적으로 배포, 비용을 절감하고 대기 시간을 줄입니다. AWS는 작년에 새로운 서비스를 출시했습니다. AWS 트레이닝 차세대 DL 모델을 개발하고 구축하기 위해 비용당 성능을 최적화하는 가속기입니다. 이 게시물에서는 M5가 어떻게 모델 교육 비용을 30% 절감할 수 있었는지 논의하고 그 과정에서 배운 모범 사례 중 일부를 공유합니다.

Trainium 인스턴스

특수 제작된 액셀러레이터의 발전으로 Amazon은 AWS Inferentia 및 Trainium의 형태로 강력한 액셀러레이터도 제공합니다. 이름에서 알 수 있듯이 이 칩은 각각 추론 및 교육 워크로드의 요구 사항을 초과하도록 최적화되었습니다. 수십억 개의 매개변수 크기에 도달하는 기초 모델의 대규모 훈련을 위해 Trainium Trn1 및 Trn1n 인스턴스 특성상 이상적인 선택입니다. Trn1 인스턴스는 최첨단 기술로 구동됩니다. 뉴런코어-v2, 그리고 엄청난 양의 가속기 컴퓨팅 및 메모리를 보유하고 있습니다. 더 많은 양의 네트워킹 대역폭(1Gbs)을 위해 Trn1,600n 인스턴스를 선택할 수도 있으므로 비용 최적화를 염두에 두고 성능 교육에 이상적으로 적합합니다.

가속기를 사용하려면 이를 지원하는 소프트웨어 계층이 필요합니다. Trn 및 Inf 칩을 사용하면 AWS 뉴런 SDK PyTorch XLA의 도움으로 Amazon에서 특별히 제작한 가속기를 잠금 해제합니다. PyTorch XLA는 PyTorch의 열성 모드를 게으른 모드 그래프 기반 구현으로 변환합니다. 그런 다음 이러한 그래프를 사용하고 추가로 컴파일하여 가속기와 함께 사용합니다. PyTorch Neuron(Neuron SDK의 일부)을 사용하면 PyTorch 사용자가 몇 줄의 코드로 Trainium NeuronCores에서 모델을 훈련할 수 있습니다.

모델 및 워크로드

M5 팀은 Amazon 전역의 다양한 팀이 즐거움을 누릴 수 있도록 지원하기 위해 기본 모델과 범용 표현을 교육하고 배포합니다. Amazon.com 고객. 그러한 모델 중 하나는 수억 개의 훈련 가능한 매개변수가 있는 신경망 아키텍처에 의해 정의된 명시적 또는 암시적 기능 상호 작용을 갖춘 텍스트 인코더 모델과 그 뒤를 잇는 다층 퍼셉트론(MLP)입니다. 이 모델은 수십억 개의 토큰에 대해 훈련되었으며 오프라인 일괄 추론 설정에서 수백만 개의 임베딩을 생성하는 데 사용됩니다. 이러한 임베딩은 고객을 대상으로 하는 Tier-1 Amazon 서비스에 대한 입력입니다.

프로덕션 파이프라인을 위한 인프라는 다음을 사용합니다. AWS 배치 과 공정한 공유 대기열 전략, 모델 훈련을 위한 컴퓨팅으로 EFA 지원 다중 노드 trn1.32xlarge 클러스터를 사용합니다. 기능적으로 생산 파이프라인은 PyTorch를 기본 DL 라이브러리로 사용하여 증분 모델 훈련, 훈련된 모델 평가, 훈련된 모델에 대한 오프라인 일괄 추론을 수행합니다.

목표

고객을 기쁘게 하는 것이 가장 중요한 원칙입니다. 파이프라인의 고객 지향 특성을 고려할 때 회귀 없이 모든 SLA(서비스 수준 계약)를 충족하는 것이 중요합니다. 우리는 기존 GPU 생산 파이프라인을 조정하고 이를 Trainium으로 전환하기 위한 두 가지 중요한 승인 기준을 식별했습니다.

• 모델 품질 – 우리 모델의 품질은 고객 경험에 직접적인 영향을 미칩니다. GPU와 Trainium 간의 모델 품질 차이는 0.1% 미만이어야 합니다.
• 훈련 처리량 – 우리는 고객에게 가장 새로운 경험을 제공하기 위해 정기적으로 모델을 반복적으로 교육합니다. 프로덕션 SLA를 충족하려면 사전 정의된 기간(예: 1주) 내에 모델 수렴을 달성해야 합니다.

다음 섹션에서는 이 기준에서 거꾸로 작업하는 여정과 Amazon 규모의 프로덕션 워크로드를 지원하기 위한 학습 내용을 공유합니다.

교육 스크립트

모델 학습을 시작하기 전에 XLA를 준수하도록 학습 스크립트를 변경해야 합니다. 모델의 크기를 고려하여 DDP(분산 데이터 병렬)를 사용하여 모델을 학습합니다. DDP를 사용하면 코드 변경 없이 모델 훈련을 실행하는 데 사용되는 시스템 수를 확장하여 모델 훈련의 처리량을 늘릴 수 있습니다. 우리는 다음에 제공된 지침을 따랐습니다. Neuron PyTorch MLP 교육 튜토리얼 훈련 스크립트에 XLA 관련 구성을 추가합니다. 이러한 코드 변경은 구현이 간단합니다. 다음은 모델 처리량을 크게 향상시킨 연습을 통해 얻은 몇 가지 중요한 기술 학습 내용입니다.

• xm.mark_step() 배치 - xm.mark_step() 느리게 수집된 계산 그래프를 컴파일하고 실행합니다. 호출 중 mark_step 너무 많이 호출하면 더 많은 수의 작은 그래프가 생성되고, 너무 적게 호출하면 소수이지만 큰 그래프가 생성됩니다. 애플리케이션에 따라 모델 학습의 처리량과 구현은 배치에 따라 달라집니다. xm.mark_step(). 우리의 구현은 하나를 배치합니다 xm.mark_step() 정방향 및 역방향 패스 후, 그리고 최적화 단계 후.
• XLA 다중 처리 장치 로더로 데이터 로더 래핑 – 이는 쉽게 놓칠 수 있는 중요한 단계입니다. 다중 처리 장치 로더 torch_xla.distributed.parallel_loader.MpDeviceLoader 처리량 향상을 위해 데이터 로드를 장치 실행과 미리 로드하고 겹치는 옵션을 사용하여 각 XLA 장치에 훈련 데이터를 로드합니다. 장치 로더도 호출합니다. xm.mark_step() 따라서 호스트에서 장치로 데이터를 로드하기 위한 그래프를 작성할 수 있습니다.

Trainium용 편집

전통적으로 GPU를 사용한 모델 개발 주기에는 모델 또는 학습 스크립트를 변경하고 이를 GPU 장치에서 직접 실행하는 작업이 포함됩니다. XLA를 사용하는 Trainium과 같은 가속기는 가속기에서 모델 훈련을 실행하기 전에 추가 단계가 필요합니다. XLA 계산 ​​그래프는 컴파일된 후에만 실행할 수 있습니다. 일반적으로 이 컴파일을 수행하는 방법에는 두 가지가 있습니다. 먼저 모든 그래프를 추적하고 컴파일한 다음 실행하는 AOT(Ahead of Time) 또는 그래프를 추적하고 컴파일하고 실행하는 JIT(Just In Time)입니다. 발생합니다. Neuron SDK는 이 두 가지를 모두 즉시 제공합니다. 일반적으로 AOT 컴파일이 먼저 수행됩니다. 그런 다음 이 컴파일 후에 그래프가 실행됩니다. 새 그래프가 발견되면 Neuron 런타임은 그래프를 실행하기 전에 JIT 컴파일을 호출합니다. AOT 컴파일을 수행하기 위해 Neuron SDK는 다음을 제공합니다. Neuron_parallel_compile, 훈련 스크립트의 시험 실행에서 그래프를 추출하고 병렬 AOT 컴파일을 수행하는 컴파일 유틸리티입니다.

AOT 컴파일의 중요한 측면은 훈련 과정에서 새로운 계산 그래프가 생성되지 않도록 하는 것입니다. 새로운 계산 그래프(및 그에 따른 재컴파일)의 소스 중 하나는 모델 훈련 중 훈련 배치의 동적 형태입니다. 정적 모양과 고정 크기 배치를 사용하면 훈련 시간 컴파일이 제거되고 모델 정확도에 영향을 주지 않고 훈련 처리량이 크게 향상된다는 사실을 발견했습니다. 훈련에 이러한 제약을 적용함으로써 우리는 AOT 컴파일 중에 모든 그래프를 추적하는 데 모델 훈련 4~5단계, 모델 검증 XNUMX단계, 모델 체크포인트 XNUMX회만 필요하다는 것을 확인했습니다. Neuron SDK는 지속적으로 발전하고 있으며 앞으로는 동적 모양도 지원할 예정이라는 점을 기억하는 것이 중요합니다.

또한, 컴파일된 그래프는 뉴런 영구 캐시 디스크 또는 아마존 단순 스토리지 서비스 (Amazon S3) 버킷. 이는 모델 아키텍처와 교육 구성이 변경되지 않는 프로덕션 워크로드에 특히 유용합니다. 따라서 컴파일 오버헤드는 한 번만 발생합니다. 캐시를 사용하는 것은 환경 플래그를 설정하는 것만큼 간단합니다.

export NEURON_COMPILE_CACHE_URL="s3://BUCKET/KEY"

Neuron 컴파일러는 또한 세 가지 기능을 제공합니다. 컴파일러 수준 최적화 옵션 (O1, O2, O3)을 사용하여 컴파일 시간과 모델 실행 처리량의 균형을 맞춥니다. O1은 컴퓨팅 그래프의 핵심 최적화를 활성화하고 컴파일 시간을 최소화하며, O3은 컴파일 시간을 늘리면서 향상된 모델 실행 처리량을 제공하고, O2(기본 옵션)는 둘 사이의 균형을 유지합니다. 사용 사례의 경우 O1 최적화를 사용하여 모델 정확도 지표를 변경하지 않고도 컴파일 시간이 86% 감소하는 동시에 기본 최적화(O5)에 비해 처리량이 약 7~2% 감소하는 것을 관찰했습니다. 사용 사례에 따라 다양한 최적화 수준을 선택할 수 있습니다.

요약하면 컴파일에 다음 플래그를 사용했습니다.

NEURON_CC_FLAGS="--target trn1 --auto-cast all --auto-cast-type bf16 --model-type transformer --optlevel O1"

체크포인트 호환성

컴파일이 성공적으로 완료되면 Trainium에서 모델 학습을 진행할 수 있습니다. 앞서 언급했듯이 우리는 모델을 점진적으로 훈련합니다. 즉, 이전에 훈련된 모델 체크포인트를 로드하고 새로운 데이터로 계속 훈련합니다. PyTorch와 PyTorch XLA는 체크포인트 상호 운용성을 통해 가속기 간 원활한 전환을 허용합니다. GPU와 Trainium 사이를 유연하게 이동할 수 있었기 때문에 이전 GPU 모델을 원활하게 로드하고 Trainium 머신에서 교육할 수 있었습니다. 이는 생산 중단 시간이나 모델 정확도 손실 없이 이전에 훈련된 최고의 모델로 모델을 초기화할 수 있도록 하는 데 매우 중요했습니다.

GPU 모델은 표준 PyTorch 모델 저장 유틸리티를 사용하여 저장되었기 때문에 PyTorch 체크포인트 로딩 유틸리티를 사용하여 Trainium 장치에 GPU 모델을 로드할 수 있었습니다.

예를 들어 GPU/CPU에서는 다음 코드를 사용하여 모델을 저장할 수 있습니다.

torch.save(model.state_dict(), PATH)

그런 다음 Trainium에 모델을 다시 로드합니다.

import torch_xla.core.xla_model as xm
xla_device = xm.xla_device()
model = MyModel(*args, **kwargs)
model.load_state_dict(torch.load(PATH))
model.to(xla_device)

마찬가지로 다음 코드를 사용하여 Trainium에 모델을 저장할 수 있습니다.

import torch_xla.core.xla_model as xm
# automatically moves the data to CPU for the master device
xm.save(model.state_dict(), PATH) 

그리고 GPU/CPU에 모델을 다시 로드합니다.

model = MyModel(*args, **kwargs)
model.load_state_dict(torch.load(PATH))
model.to(device) # can be any device

실제로 모델 훈련에 DDP를 사용하기 때문에 모델 로딩은 이전 체크포인트를 훈련하는 데 사용된 머신 수에 구애받지 않습니다. 이를 통해 코드 변경이나 모델 교육에 대한 부작용 없이 Trn1 플릿을 수평으로 확장할 수 있습니다. 이러한 PyTorch 기반 체크포인트는 AWS Inferentia2 또는 기타 가속기의 추론 사용 사례에 대해 직접 사용되거나 토치 스크립트로 작성될 수도 있습니다.

운영 안정성

프로덕션 환경에서 워크로드를 실행하려면 여러 SLA를 충족해야 한다는 점은 아무리 강조해도 지나치지 않습니다. 우리의 사용 사례에서는 모델 품질 및 교육 처리량 SLA와 별도로 생산 파이프라인이 운영상 안정적이어야 합니다. 즉, 모델 교육, 평가 및 추론 중에 가동 중지 시간과 중단이 최소화되어야 합니다.

기존 GPU 기반 파이프라인과 마찬가지로 파이프라인의 운영 안정성을 높이기 위해 다양한 메커니즘을 추가했습니다. 모델 훈련을 시작하기 전에 여러 건전성 테스트를 실행하여 기계의 상태를 평가합니다. 이러한 테스트에는 일반적으로 가속기 장치의 상태를 확인하기 위한 간단한 텐서 작업이 포함됩니다. 분산 교육의 경우 테스트를 실행하여 인스턴스 간의 집단 통신도 확인하는 것이 중요하다는 사실을 관찰했습니다. 우리는 NCCOM 테스트 스위트 이를 달성하기 위해 Neuron SDK에서 All-Gather, All-Reduce 및 Reduce-Scatter와 같은 다양한 작업을 실행합니다.

우리가 언급한 제안을 따른 후에도 기본 가속기에 관계없이 모든 파이프라인에서 일시적인 문제가 불가피하다는 것을 관찰했습니다. 학습 파이프라인에서 복원력을 구축하려면 재시도 메커니즘을 구축하여 이러한 잠재적인 문제를 해결하는 것이 좋습니다. 우리는 사용 AWS Batch 자동 재시도 모델 학습 중에 일시적인 오류가 발생한 작업을 재시도합니다. 훈련이 끝날 무렵 오류가 발생하면 다시 시작하는 데 비용이 많이 들 수 있습니다. 이 문제에 대응하기 위해 우리는 이전에 훈련된 모델 체크포인트를 로드하고 해당 지점부터 훈련을 계속하도록 훈련 스크립트를 조정했습니다. 이 기능을 사용하면 최소한의 오버헤드로 실패한 훈련 작업을 적극적으로 다시 시작할 수 있습니다.

이러한 탄력성 메커니즘을 적용하여 Trn98.5의 워크로드에 대해 기존 GPU 파이프라인 성공률과 비슷한 1%의 성공률을 달성할 수 있었습니다.

결과

모델의 정확성을 검증하기 위해 동일한 GPU 체크포인트에서 두 모델을 초기화하고 하나는 Trainium에서, 다른 하나는 유사한 GPU에서 훈련했습니다. 두 모델 모두 동일한 훈련 하이퍼파라미터로 훈련되었습니다. 지표 계산에 사용되는 데이터세트는 홀드아웃 데이터세트이며, N 글로벌 단계마다 이 데이터세트에 대한 모델의 정확도를 평가합니다. X축은 전역 단계이고 Y축은 모델 정확도입니다. 다음 그래프의 각 지점에서 모델 정확도의 차이는 0.1% 미만으로 관찰되었습니다.

또한 모델 훈련의 비용 효율성을 평가하기 위해 모델 수렴에 도달하는 데 걸리는 벽시계 시간을 비교하는 것을 선호합니다. 우리는 이것이 토큰당 비용, 달성된 FLOPS/달러 및 기타 요소와 같은 측정값과 비교하여 비용 절감에 대한 보다 실용적인 관점을 제공한다고 믿습니다. trn1.32xl과 비교 가능한 훈련 시간을 고려하면 아마존 엘라스틱 컴퓨트 클라우드 (Amazon EC2) 인스턴스에서는 Trainium이 모델 수렴 비용을 최대 30% 더 저렴하게 제공하는 것으로 나타났습니다.

결론

DL 워크로드에 대한 다양한 가속기를 평가할 때 고려해야 할 요소가 많이 있습니다. 가장 중요한 것 중 일부는 모델 품질, 처리량, 비용 및 가용성입니다. 선택한 가속기에 따라 모델 품질과 처리량이 희생되지 않도록 하는 것이 가장 중요합니다.

Annapurna Neuron 팀과의 파트너십 및 협력 덕분에 Amazon Search M5 팀은 Trainium으로 이전하여 비용을 최대 30% 절감할 수 있었습니다. 팀은 Trainium을 사용하여 시장에 있는 유사한 가속기와 모델 품질 및 처리량 동등성을 달성할 수 있습니다. 체크포인트 상호 운용성과 XLA 지원을 통한 최소한의 코드 변경을 통해 M5는 워크로드에 대해 여러 가속기 중에서 선택할 수 있었습니다. 이를 통해 M5 팀은 Trainium의 대규모 컴퓨팅 성능을 활용하고 Amazon.com 고객을 만족시킬 수 있는 가속기 독립적 솔루션을 구축할 수 있었습니다. 운영 관점에서 Trainium은 Amazon 규모의 Tier-1 서비스를 지원할 수 있는 능력이 입증되었습니다. M5 팀은 최저 비용으로 Amazon에 최고의 모델을 제공하기 위해 계속해서 더 많은 워크로드를 Trainium으로 이전하고 있습니다.

요약하면, M5 팀은 일련의 가속기에 Trainium을 추가하여 비용 효율적인 프로덕션 등급 ML 교육을 수행할 수 있었습니다. ML 워크로드를 위해 특별히 제작된 Amazon 실리콘의 이점을 활용하려면 Trainium 및 AWS Inferentia와 같은 기타 Neuron 디바이스를 살펴보는 것이 좋습니다. 다음과 같은 다양한 모델이 포함된 많은 튜토리얼 중 하나로 쉽게 시작해보세요. Trainium에서 이용 가능한 Llama 2.

저자에 관하여

아비난단 파트니 Amazon Search의 수석 소프트웨어 엔지니어입니다. 그는 확장 가능한 분산 딥 러닝 교육 및 실시간 추론을 위한 시스템 및 도구 구축에 중점을 둡니다.

제임스 파크 Amazon Web Services의 솔루션 아키텍트입니다. 그는 Amazon.com과 협력하여 AWS에서 기술 솔루션을 설계, 구축 및 배포하며 AI 및 기계 학습에 특히 관심이 있습니다. 그는 여가 시간에 새로운 문화, 새로운 경험을 찾고 최신 기술 동향을 파악하는 것을 즐깁니다. 당신은 그를 찾을 수 있습니다 링크드인.

제리 만닐 Amazon Search의 소프트웨어 엔지니어입니다. 그는 분산된 교육 인프라의 효율성, 견고성 및 확장성을 개선하는 일을 하고 있습니다.

켄 수 Amazon Search의 소프트웨어 엔지니어입니다. 그는 교육 효율성과 확장 가능한 분산 교육 워크플로를 개선하는 데 주력하고 있습니다. 일 외에는 하이킹과 테니스를 좋아합니다.

RJ Amazon의 엔지니어입니다. 그는 훈련을 위한 분산 시스템용 시스템을 구축 및 최적화하고 ML 추론의 대기 시간을 줄이기 위해 시스템 채택을 최적화하는 작업을 하고 있습니다. 업무 외적으로 그는 음식 조리법을 구축하기 위해 Generative AI를 사용하는 방법을 모색하고 있습니다.

• SEO 기반 콘텐츠 및 PR 배포. 오늘 증폭하십시오.
• PlatoData.Network 수직 생성 Ai. 자신에게 권한을 부여하십시오. 여기에서 액세스하십시오.
• PlatoAiStream. 웹3 인텔리전스. 지식 증폭. 여기에서 액세스하십시오.
• 플라톤ESG. 탄소, 클린테크, 에너지, 환경, 태양광, 폐기물 관리. 여기에서 액세스하십시오.
• PlatoHealth. 생명 공학 및 임상 시험 인텔리전스. 여기에서 액세스하십시오.
• 출처: https://aws.amazon.com/blogs/machine-learning/how-amazon-search-m5-saved-30-for-llm-training-cost-by-using-aws-trainium/