[병렬 처리]. CPU 모델 , 문장 임베딩 멀티 프로세싱으로 10배 빠르게 처리하기

jongmin-oh·2023년 5월 11일

문제 발생:

오늘 새로운 모델을 테스트 하기 위해, 600만개의 문장을 임베딩 해야했고,
GPU를 사용해서 30분만에 임베딩 했다, 그래서 테스트를 해보니 성능이 괜찮았다.

하지만 여기서 끝이 아니라 서비스를 위해서는 모델을 양자화 시켜야 한다.
**속도가 5배 이상 차이나기 때문

또 ONNX 양자화 모델은 CPU 전용으로 만들었기 때문에(*서버가 CPU 서버이다)
양자화 된 모델의 성능 또한 테스트하려면 CPU로 600만 개의 문장을 임베딩 해야했다.

기본적으로 돌려보았더니 약 2시간 10분 가량 걸렸다.

멀티 프로세싱

바로 Htop을 실행해서 CPU 상태를 확인했더니, 하나의 CPU만 사용하는 게 아니라 여러 CPU를 사용하는 것 처럼 보였지만, 결국 하나의 CPU를 사용하는 것과 다름 없었다.(속도 계산이 딱 맞아떨어진다 ㅠㅠ)

그래서 전처리 할때 사용했던 멀티 프로세싱 코드를 조금 수정하여 CPU 임베딩에도 사용해보도록 했다.

추가적으로 멀티 프로세싱을 쉽게 해주는 Python 모듈 Ray도 있는데 테스트 해본 결과 시간이 좀 처럼 줄지 않았다, Ray는 더 공부해봐야곘다...

해결

모델 및 데이터 로드

# 모델 불러오기
from sentence_transformers import SentenceTransformer
model = SentenceTransformer(MODEL_PATH)
model.max_seq_length = MODEL_MAX_LENGTH

# 문장 불러오기
df = pd.read_parquet(DATA_PATH)
sentence = df['sentence'] # 600만개 문장

ONNX 모델 사용

import onnxruntime as rt
from transformers import AutoTokenizer

sess = rt.InferenceSession("onnx_models/model_uint8.onnx.onnx", providers=['CPUExecutionProvider'])
tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)

멀티 프로세싱 코드

import numpy as np
from concurrent.futures import ProcessPoolExecutor
import multiprocessing
import itertools
from tqdm import tqdm

# 가지고 있는 CPU중 50%만 사용하겠다 , 나의 경우 24 * 0.5 = 12개
n_cpu = int(multiprocessing.cpu_count() * 0.5)


def split_embedding(start, end):
	"""
    문장 임베딩을 수행하는 함수를 정의합니다. 
    이 함수는 시작 인덱스(start)와 끝 인덱스(end)를 받아서 해당 범위 내의 문장들의 임베딩을 계산하고, 
    리스트로 반환합니다.
    """
    return [ sentence_embedding(i) for i in tqdm(sentence[start:end])]
    
def func_multi():
	"""
    멀티프로세싱 함수를 정의합니다. 
    이 함수는 문장 리스트(sentence)를 n_cpu개의 코어로 분할하고,
    각 코어에서 split_embedding 함수를 병렬로 실행하여 문장 임베딩을 계산합니다.
    그리고 코어별로 계산된 임베딩을 numpy array로 반환합니다.
    """
    global sentence
    full_len = len(sentence) # 데이터 수
    process_index = int(full_len / n_cpu) # 분할 범위
    
    rng_list = [(i + 1) * process_index for i in range(n_cpu)] # 분할 범위 리스트 생성
    if rng_list[0] != 0:  # 맨앞에 0 추가 ( 0부터 시작해야함 )
        rng_list.insert(0, 0)
    if rng_list[-1] < full_len: # 맨 뒤는 총 데이터수 추가
        rng_list.append(full_len)
    
    with ProcessPoolExecutor(max_workers = n_cpu) as executor:
        embedding_list = list(executor.map(split_embedding, rng_list[0:-1], rng_list[1:]))
    
    result = list(itertools.chain.from_iterable(embedding_list)) # 코어별로 나눠진 리스트를 1차원으로 합침
    result = np.array(result)
    
    
    return result
    

# 멀티프로세싱 실행
quantizer_embeddings = func_multi()