임베딩 매개변수는 1%만 필요하고 하드웨어 비용은 10배 절감되며 오픈 소스 솔루션 단일 GPU는 대규모 권장 모델을 교육합니다.

심층 추천 모델(DLRM)은 동영상 추천, 쇼핑 검색, 광고 푸시 및 기타 트래픽 수익화 서비스 등 인터넷 기업에서 딥 러닝을 적용하는 데 가장 중요한 기술 시나리오가 되어 사용자 경험과 비즈니스 상업성을 크게 향상시켰습니다. 값. 그러나 대규모 사용자 및 비즈니스 데이터, 빈번한 반복 업데이트 요구 사항, 높은 교육 비용 등은 모두 DLRM 교육에 심각한 문제를 야기합니다.

DLRM에서는 먼저 임베디드 테이블(EmbeddingBags)에서 조회를 수행한 후 다운스트림 계산을 완료해야 합니다. 내장된 테이블은 DLRM 메모리 요구 사항의 99% 이상을 차지하는 경우가 많지만 계산에서는 1%에 불과합니다. GPU 온칩 고속 메모리(고대역폭 메모리)와 강력한 컴퓨팅 성능 덕분에 GPU는 DLRM 교육을 위한 주류 하드웨어가 되었습니다. 그러나 추천 시스템에 대한 연구가 심화됨에 따라 임베딩 테이블 크기가 커지고 GPU 메모리가 제한된다는 것은 심각한 모순을 형성합니다. GPU를 사용하여 GPU 메모리 벽의 한계를 극복하면서 초대형 DLRM 모델을 효율적으로 훈련하는 방법은 DLRM 분야에서 해결해야 할 핵심 문제가 되었습니다.

Colossal-AI는 이전에 이종 전략을 사용하여 동일한 하드웨어에서 훈련된 NLP 모델의 매개변수 용량을 수백 배로 늘리는 데 성공했습니다. 최근에는 소프트웨어 캐시를 사용하여 이를 추천 시스템으로 성공적으로 확장했습니다. (캐시) 방식으로 NLP 모델의 매개변수 용량을 수백 배로 늘리고 내장된 테이블을 GPU 메모리에 동적으로 저장합니다. Colossal-AI는 소프트웨어 캐시 설계를 기반으로 향후 입력될 교육 데이터를 관찰하여 소프트웨어 캐시 검색 및 데이터 이동 오버헤드를 줄이기 위해 파이프라인 프리페칭도 추가합니다. 동시에 동기식 업데이트 방식으로 GPU에서 전체 DLRM 모델을 훈련합니다. 이는 널리 사용되는 하이브리드 병렬 훈련 방법과 결합되어 여러 GPU로 확장될 수 있습니다. 실험에 따르면 Colossal-AI는 GPU의 임베딩 매개변수 중 1%만 유지하면 되며 여전히 뛰어난 엔드투엔드 훈련 속도를 유지할 수 있습니다. 다른 PyTorch 솔루션에 비해 그래픽 메모리 요구 사항이 대폭 줄어들고 단일 그래픽 카드로 테라바이트 수준의 추천 모델을 훈련할 수 있습니다. 예를 들어, 91GB의 Embedding Bag을 차지하는 DLRM을 훈련하는 데는 5GB의 비디오 메모리만 필요합니다. 훈련 하드웨어 비용은 약 200,000위안인 A100 2개에서 보급형 그래픽 카드까지 10배로 절감됩니다. 가격이 약 2,000위안인 RTX 3050과 같은 것입니다.

오픈 소스 주소: https://github.com/hpcaitech/ColossalAI

기존 임베딩 테이블 확장 기술

임베딩 테이블은 이산 정수 특징을 연속 부동 소수점 특징으로 매핑합니다. 벡터, 다음 그림은 DLRM의 임베딩 테이블 훈련 프로세스를 보여줍니다. 먼저 임베딩 테이블의 각 특징에 대해 임베딩 테이블에 해당하는 행을 찾은 후 최대, 평균, 합계 연산 등의 축소 연산을 사용하여 특징 벡터로 변환하고 후속 조밀 신경망에 전달합니다. . DLRM의 임베디드 테이블 훈련 프로세스는 주로 불규칙한 메모리 액세스 작업이므로 하드웨어 메모리 액세스 속도에 의해 심각하게 제한된다는 것을 알 수 있습니다.

산업 등급 DLRM의 임베디드 테이블은 단일 GPU의 최대 비디오 메모리 용량인 수십 GB를 훨씬 초과하는 수백 GB 또는 심지어 TB 수준에 도달할 수 있습니다. 단일 GPU의 메모리 벽을 뛰어넘어 DLRM의 임베디드 테이블 크기를 늘리는 방법에는 여러 가지가 있습니다. 아래 그림에 표시된 GPU 클러스터의 메모리 계층 다이어그램을 예로 들어 몇 가지 일반적인 솔루션의 장점과 단점을 분석해 보겠습니다.

GPU 모델 병렬성: 임베딩 테이블은 여러 GPU의 메모리에 분할되어 분산되며 훈련 중에 GPU 간의 상호 연결 네트워크를 통해 중간 결과가 동기화됩니다. 이 방법의 단점은 임베디드 테이블 분할 부하가 균일하지 않고 확장성 문제를 해결하기 어렵다는 것입니다. 둘째, GPU를 추가하는 데 드는 초기 하드웨어 비용이 높으며, DLRM 훈련 중에 GPU의 컴퓨팅 성능이 완전히 활용되지 않고 HBM 대역폭 이점만 사용되므로 GPU 활용도가 낮아집니다.

CPU 부분 훈련: 임베딩 테이블을 두 부분으로 분할합니다. 한 부분은 GPU에서 훈련되고 다른 부분은 CPU에서 훈련됩니다. 데이터 분포의 롱테일 효과를 활용하여 CPU 계산 비율을 최대한 작게 만들고 GPU 계산 비율을 최대한 크게 만들 수 있습니다. 하지만 배치 크기가 커지면 모든 미니 배치 데이터가 CPU나 GPU에 도달하기 어렵습니다. 데이터가 동시에 CPU나 GPU에 도달하면 이 방법은 처리하기 어렵습니다. 또한 DDR 대역폭과 HBM은 데이터 크기가 하나 다르기 때문에 입력 데이터의 10%가 CPU에서 훈련되더라도 전체 시스템 속도가 최소한 절반으로 느려집니다. 또한 CPU와 GPU는 중간 결과를 전송해야 하기 때문에 상당한 통신 오버헤드가 발생하고 훈련 속도가 더욱 느려집니다. 따라서 연구자들은 이러한 성능 결함을 방지하기 위해 비동기 업데이트와 같은 방법을 설계했습니다. 그러나 비동기 방법은 훈련 결과에 불확실성을 초래할 수 있으며 실제로 알고리즘 엔지니어가 가장 먼저 선택하는 방법은 아닙니다.

소프트웨어 캐시: 모든 학습이 GPU에서 수행되는지 확인합니다. 임베딩 테이블은 CPU와 GPU로 구성된 이기종 공간에 저장될 때마다 소프트웨어 캐시 방식을 통해 필요한 부분이 GPU로 교체됩니다. . 이 방법을 사용하면 증가하는 임베디드 테이블 요구 사항을 충족하기 위해 스토리지 리소스를 저렴하게 확장할 수 있습니다. 또한 계산을 위해 CPU를 사용하는 것과 비교하여 이러한 방식의 전체 훈련 프로세스는 GPU에서 완전히 완료되어 HBM 대역폭 이점을 최대한 활용합니다. 그러나 캐시 쿼리 및 데이터 이동으로 인해 추가적인 성능 손실이 발생합니다.

임베딩 테이블을 위한 우수한 소프트웨어 캐시 솔루션이 이미 있지만 fbgemm과 같은 맞춤형 EmbeddingBags 커널 구현을 사용하거나 타사 딥 러닝 프레임워크를 사용하는 경우가 많습니다. Colossal-AI는 기본 PyTorch를 기반으로 커널 수준 변경을 수행하지 않으며 즉시 사용 가능한 소프트웨어 Cache EmbeddingBags 구현 세트를 제공하며 DLRM 교육 프로세스를 더욱 최적화하고 프리패치 파이프라인을 제안합니다. 캐시 오버헤드를 더욱 줄입니다.

Memory Hierarchy

Colossal-AI의 임베디드 테이블 소프트웨어 Cache

Colossal-AI는 사용자가 자신의 모델에서 사용할 수 있도록 소프트웨어 캐시를 구현하고 이를 nn.Module로 패키징합니다. . DLRM의 임베딩 테이블은 일반적으로 여러 개의 Embedding으로 구성된 EmbeddingBags이며 CPU 메모리에 상주합니다. 메모리 공간의 이 부분을 CPU 가중치라고 합니다. EmbeddingBags 데이터의 작은 부분은 훈련에 사용될 데이터를 포함하는 GPU 메모리에 저장됩니다. 메모리 공간의 이 부분을 CUDA 캐시 가중치라고 합니다. DLRM 훈련 중에 먼저 이 반복을 위해 미니 배치에 입력된 데이터에 해당하는 내장 테이블의 행을 결정해야 합니다. 일부 행이 GPU에 없으면 CPU 가중치에서 CUDA로 전송되어야 합니다. 캐시된 가중치. GPU에 공간이 충분하지 않으면 LFU 알고리즘을 사용하여 캐시에 액세스한 기록 빈도를 기준으로 가장 적게 사용된 데이터를 제거합니다.

캐시 검색을 구현하려면 도움이 되는 몇 가지 보조 데이터 구조가 필요합니다. 캐시된_idx_map은 CPU 가중치와 CUDA 캐시 가중치의 행 번호 및 빈도 간의 대응 관계를 저장하는 1차원 배열입니다. GPU에서 액세스되는 해당 행의 정보입니다. CPU 가중치 크기에 대한 CUDA 캐시 가중치 크기의 비율을 캐시_ratio라고 하며 기본값은 1.0%입니다.

캐시는 CUDA Weight의 데이터를 조정하기 위해 각 반복 전에 특히 3단계로 실행됩니다.

1단계: CPU 인덱스: CPU Weight

입력 미니 배치의 input_ids와 cashed_idx_map의 교차점을 가져와 해당 라인을 찾아야 합니다. CPU Weight Line 번호에서 CPU에서 GPU로 이동해야 하는 번호입니다.

2단계: GPU 인덱스: CUDA에서 제거할 수 있는 행 찾기 사용 빈도에 따른 가중치

이를 위해서는 캐시_idx_map과 input_ids의 차이 세트 뒤의 부분을 낮은 것부터 순서대로 가져가야 합니다. top-k(최대 k개 숫자 사용) 작업에 따라 높게 설정됩니다.

3단계: 데이터 이동:

CUDA 캐시 가중치의 해당 행을 CPU 가중치로 이동한 다음, CPU 가중치의 해당 행을 CUDA 가중치로 이동합니다.

데이터 전송 모듈은 CUDA 캐시 가중치와 CPU 가중치 간의 양방향 데이터 전송을 담당합니다. 비효율적인 라인별 전송과 달리 캐시를 먼저 사용하고 중앙 집중식 전송을 사용하여 PCI-e 대역폭 활용도를 향상시킵니다. 메모리에 흩어져 있는 임베디드 행은 소스 장치의 로컬 메모리에서 연속적인 데이터 블록으로 수집되고, 블록은 CPU와 GPU 간에 전송되어 대상 메모리의 해당 위치로 분산됩니다. 데이터를 블록 단위로 이동하면 PCI-e 대역폭 활용도가 향상될 수 있으며 병합 및 분산 작업에는 CPU 및 GPU에 대한 온칩 메모리 액세스만 포함되므로 오버헤드가 그리 크지 않습니다.

Colossal-AI는 크기가 제한된 버퍼를 사용하여 CPU와 GPU 간에 데이터를 전송합니다. 최악의 경우 모든 입력 ID가 캐시를 놓치고 많은 수의 요소를 전송해야 합니다. 버퍼가 너무 많은 메모리를 차지하는 것을 방지하기 위해 버퍼 크기가 엄격하게 제한됩니다. 전송되는 데이터가 버퍼보다 ​​큰 경우 전송은 여러 번 완료됩니다.

Cached EmbeddingBag Workflow

소프트웨어 캐시 성능 분석

위의 캐시 Step1 및 Step2 작업은 모두 메모리 액세스 집약적입니다. 따라서 GPU HBM의 대역폭을 활용하기 위해 GPU에서 실행되고 딥러닝 프레임워크로 캡슐화된 API를 사용하여 구현됩니다. 그럼에도 불구하고, 캐시 작업의 오버헤드는 테이블 임베딩을 위한 GPU에서의 훈련 작업에 비해 특히 중요합니다.

예를 들어 199초 동안 지속되는 학습 작업에서 캐시 작업의 오버헤드는 99초로 전체 컴퓨팅 시간의 거의 50%을 차지합니다. 분석 결과, Cache의 주요 오버헤드는 주로 Step1과 Step2에서 발생했습니다. 아래 그림의 기본 위치는 이때 캐시 오버헤드의 시간 분석을 보여줍니다. 캐시 1단계와 2단계의 빨간색과 주황색 단계가 전체 캐시 오버헤드의 70%를 차지합니다.

캐시 작업 시간 분석

위의 문제가 발생하는 이유는 전통적인 캐시 전략이 다소 "근시안적"이며 현재 미니에 따라서만 캐시를 조정할 수 있기 때문입니다. - 배치 상황이므로 큰 문제입니다. 쿼리 작업에 시간의 일부가 낭비됩니다.

캐시 파이프라인 프리패치

캐시 오버헤드를 줄이기 위해 Colossal-AI는 "미래 지향적" 캐시 메커니즘을 설계했습니다. Colossal-AI는 첫 번째 미니 배치에서만 캐시 작업을 수행하는 대신 나중에 사용할 여러 미니 배치를 미리 가져오고 통합된 방식으로 캐시 쿼리 작업을 수행합니다.

아래 그림과 같이 Colossal-AI는 통합 캐시 작업을 위해 프리페칭을 사용하여 여러 미니 배치 데이터를 병합하고, 파이프라인 접근 방식을 사용하여 데이터 읽기 및 계산의 오버헤드를 중첩합니다. 이 예에서 프리페치 미니 배치 번호는 2입니다. 훈련을 시작하기 전에 미니 배치 0,1 데이터를 디스크에서 GPU 메모리로 읽은 다음 캐시 작업을 시작한 다음 두 미니 배치의 순방향 및 역방향 전파와 매개 변수 업데이트를 수행합니다. 동시에 미니 배치 2, 3에 대한 초기 데이터 읽기가 수행될 수 있으며, 이 오버헤드는 계산과 겹칠 수 있습니다.

기본 캐시 실행 방법과 비교한 그림 [캐시 작업의 시간 분석]은 프리페치 8 미니 배치의 캐시 시간 분석과 기준을 비교합니다. 총 훈련 시간은 201초에서 120초로 줄어들었고, 그림에 표시된 캐시 단계 작업 시간의 비율도 크게 줄었습니다. 캐시 작업을 독립적으로 수행하는 각 미니 배치와 비교할 때 각 부분의 시간, 특히 캐시 작업의 처음 두 단계가 단축되는 것을 볼 수 있습니다.

요약하자면, 캐시 파이프라인 프리페치는 두 가지 이점을 제공합니다.

a. 희석 캐시 인덱스 오버헤드

프리페치의 가장 확실한 이점은 1단계와 2단계의 오버헤드를 줄여서 이 2단계 작업이 전체 훈련 프로세스의 5% 미만을 차지하게 한다는 것입니다. [캐시 작업의 시간 분석]에서 볼 수 있듯이 8개의 미니 배치 데이터를 프리페치함으로써 프리페치를 하지 않은 기준에 비해 캐시 쿼리의 오버헤드가 크게 줄어듭니다.

b. CPU-GPU 데이터 이동 대역폭 증가

더 많은 데이터를 집중시켜 CPU-GPU 전송 대역폭을 최대한 활용하여 데이터 전송 세분성을 향상시킵니다. 위의 예에서는 CUDA->CPU 대역폭이 860MB/s에서 1477MB/s로 증가하고, CPU->CUDA 대역폭이 1257MB/s에서 2415MB/s로 증가하여 성능 향상이 거의 두 배로 늘어납니다.

Easy to use

추천 모델 구축 시 다음과 같은 코드만 초기화하면 임베딩 테이블 용량이 대폭 늘어나 TB급 대용량 추천이 가능합니다. 저렴한 비용으로 모델 훈련을 해보세요.

Bashfrom colossalai.nn.parallel.layers.cache_embedding import CachedEmbeddingBag
emb_module = CachedEmbeddingBag(num_embeddings=num_embeddings,embedding_dim=embedding_dim,mode="sum"include_last_offset=True,sparse=True,_weight=torch.randn(num_embeddings, embedding_dim),warmup_ratio=0.7,cache_ratio = 0.01,)

성능 테스트

NVIDIA A100 GPU(80GB) 및 AMD EPYC 7543 32코어 프로세서(512GB) 하드웨어 플랫폼에서 Colossal-AI는 Meta의 DLRM 모델을 테스트 대상으로 사용하여 초대형 데이터를 사용합니다. Cretio 1TB를 설정하고 Meta의 dlrm_datasets는 데이터 세트를 테스트 모델로 생성합니다. 실험에서는 모든 임베딩 테이블을 GPU에 저장하는 PyTorch 학습 속도를 기준으로 사용했습니다.

Cretio 1TB

Cretio 1TB 임베딩 테이블에는 총 177944275개의 행이 있고 임베딩 희미함=128로 설정되었으며 임베딩 테이블 메모리 요구 사항은 91.10GB입니다. 가장 고급형 NVIDIA A100 80GB라도 단일 GPU 메모리에 모든 EmbeddingBags를 저장하는 메모리 요구 사항을 충족할 수 없습니다.

그러나 Colossal-AI를 사용하면 여전히 단일 GPU에서 학습이 완료됩니다. 캐시 비율이 0.05일 때 메모리 소비는 5.01GB에 불과하며 이는 직접적으로 약 18배까지 확장될 수 있습니다. 단일 GPU의 테라바이트 수준 추천 시스템. 학습 속도 측면에서 보면 아래 그림과 같이 다양한 배치 크기에서 100M 샘플을 학습하는 데 지연이 있음을 보여줍니다. 녹색 Prefetch1은 프리페치를 하지 않은 지연이고, 파란색 Prefetch8은 프리페치를 사용한 지연입니다(prefetch mini-batch=8). 프리페치 파이프라인 최적화가 전체 성능 향상에 중요한 역할을 한다는 것을 알 수 있습니다. 그림에서 각 열의 어두운 부분은 Cache 오버헤드입니다. Prefetching을 사용한 후 Cache 오버헤드는 전체 훈련 시간의 15% 이내로 제어됩니다.

다중 GPU 확장성

8192를 전역 배치 크기로 사용하고, 8개의 GPU 카드에서 테이블 단위 샤딩을 EmbeddingBags로 사용하여 DLRM을 병렬로 훈련하고, 100M 샘플을 훈련합니다. 이때 Prefetch 크기를 4로 설정하고, ColossalAI-mem-cr0.05는 캐시 비율=0.05, ColossalAI-mem-cr0.5=0.5로 설정합니다. 아래 그림은 다양한 GPU 시나리오에 대한 훈련 대기 시간을 보여줍니다. 1개의 GPU를 사용할 때 PyTorch OOM을 훈련할 수 없다는 점을 제외하면 PyTorch와 Colossal-AI의 훈련 시간은 다른 경우와 비슷합니다. 4개와 8개 GPU를 사용해도 성능이 크게 향상되지 않는 것을 볼 수 있습니다. 1. 결과를 동기화하려면 엄청난 통신 오버헤드가 필요하기 때문입니다. 2. 테이블 방식 샤딩은 샤딩 부하의 불균형을 초래합니다. 또한 임베딩 테이블 학습 확장성을 확장하기 위해 여러 GPU를 사용하는 것이 그리 좋지 않음을 보여줍니다.

아래 그림은 비디오 메모리 사용량을 보여줍니다. 비디오 메모리 사용량은 카드마다 다릅니다. 여기에는 최대 비디오 메모리 값이 표시됩니다. GPU를 하나만 사용할 경우 Colossal-AI의 소프트웨어 Cache 방식만 학습할 수 있으며, 여러 카드가 병렬로 차지하는 메모리도 여러 배로 크게 줄어듭니다.

Meta Research의 합성 데이터 세트 dlrm_datasets는 업계 임베디드 테이블의 교육 액세스 동작을 모방하므로 연구에서 추천 시스템과 관련된 소프트웨어 및 하드웨어 설계에 대한 테스트 참조로 자주 사용됩니다. 5억 행의 임베디드 테이블 항목을 하위 데이터 세트로 선택하고 테스트를 위해 256GB와 128GB의 EmbeddingBags 2개를 구성합니다.

PyTorch는 비디오 메모리 부족으로 인해 단일 카드 A100에서 훈련할 수 없습니다. 이와 대조적으로 Colossal-AI의 소프트웨어 캐시는 GPU 메모리 요구 사항을 크게 줄여 최대 256GB의 임베딩 테이블을 훈련할 수 있을 만큼 충분하며 TB 수준까지 추가 확장이 가능합니다. 또한 파이프라인 프리페칭은 프리페치 횟수가 32개일 경우 프리페칭을 하지 않은 경우에 비해 총 시간이 60% 단축되며 GPU 스토리지에 대한 수요도 증가하지 않습니다.

One More Thing

대형 모델 시대를 위한 종합 딥러닝 시스템인 Colossal-AI는 효율적인 다차원 자동 병렬 처리, 이종 병렬 처리, 메모리 관리, 대규모 최적화 라이브러리, 적응형 작업 스케줄링 등을 통해 AI 대형 모델 교육 및 추론을 효율적이고 신속하게 배포하고 AI 대형 모델 적용 비용을 절감합니다.

Colossal-AI 관련 솔루션은 자율 주행, 클라우드 컴퓨팅, 소매, 의료, 칩 및 기타 산업 분야의 유명 제조업체에서 성공적으로 구현되었으며 널리 호평을 받았습니다.

Colossal-AI는 오픈 소스 커뮤니티 구축에 중점을 두고 중국어 튜토리얼 제공, 사용자 커뮤니티 및 포럼 개설, 사용자 피드백에 대한 효율적인 커뮤니케이션 및 반복 업데이트를 수행하고 PaLM, AlphaFold 및 OPT와 같은 최첨단 애플리케이션을 지속적으로 추가합니다.

Colossal-AI는 본래 오픈소스였기 때문에 GitHub 및 Papers With Code 핫리스트에서 여러 차례 세계 1위를 차지했으며, 수만 명의 스타 오픈소스 프로젝트와 함께 국내외에서 주목을 받았습니다. 별!

프로젝트 오픈 소스 주소: https://github.com/hpcaitech/ColossalAI

위 내용은 임베딩 매개변수는 1%만 필요하고 하드웨어 비용은 10배 절감되며 오픈 소스 솔루션 단일 GPU는 대규모 권장 모델을 교육합니다.의 상세 내용입니다. 자세한 내용은 PHP 중국어 웹사이트의 기타 관련 기사를 참조하세요!