Hugging Face와 Dask를 결합하여 100행 pandas 처리를 211백만행 클라우드 GPU 병렬 처리로 확장

Scaling AI-based Data Processing with Hugging Face + Dask

2024년 10월 9일8분intermediate

AI 요약

Context

Hugging Face의 사전 학습 모델과 대규모 데이터셋(수백 GB~TB)은 접근성이 높지만, AI 작업 확장 시 메모리 부족과 순차 처리로 인한 높은 연산 비용 문제가 발생한다. FineWeb 데이터셋(432GB 디스크, 715GB 메모리)의 경우 단일 머신에서 순차 처리는 금지적으로 느리다.

Technical Solution

Dask DataFrame 도입: pandas와 유사한 API로 TB 규모 데이터셋을 청크로 분할하여 메모리 외 컴퓨팅 수행
map_partitions를 활용한 병렬 추론: compute_scores 함수를 Dask DataFrame의 각 파티션(pandas DataFrame)에 병렬 적용
동적 장치 감지: GPU 가용성(CUDA/MPS) 확인 후 자동으로 CPU/GPU 전환하는 함수 작성으로 로컬과 클라우드 환경 호환성 확보
Parquet 포맷 활용: Hugging Face 기본 포맷인 Parquet으로 열 기반 필터링과 압축 효율화
Coiled 클러스터 오케스트레이션: 100개 워커 노드(g5.xlarge GPU 인스턴스) 구성하여 분산 처리 및 GPU 추론 수행, 결과를 Hugging Face 데이터셋으로 병렬 저장

Impact

100행 처리 시간: M1 Mac GPU에서 약 10초 | 211백만행 처리: 클라우드 100개 GPU 워커 노드에서 병렬 실행

Key Takeaway

Hugging Face 모델과 Dask의 조합은 로컬 프로토타입 코드를 최소한의 수정으로 멀티 노드 GPU 클러스터로 확장 가능하게 한다. 배치 크기(25 → 768) 조정과 장치 감지 로직을 통해 개발-프로덕션 간 일관된 처리 파이프라인을 유지할 수 있다.

실천 포인트

TB 규모의 텍스트 데이터를 Hugging Face 트랜스포머로 분류해야 하는 데이터 엔지니어링 팀에서 Dask의 map_partitions를 사용하면 pandas 코드의 함수 시그니처를 유지한 채 멀티 GPU 병렬 처리로 전환할 수 있다. 동시에 배치 크기와 장치 로직을 코드에 매장하면 로컬 GPU, CPU, 클라우드 GPU 환경 간 마이그레이션 비용을 최소화할 수 있다.

태그

#GPU Inference #Distributed Computing #Dask #Data Processing #Hugging Face

원문 읽기