카카오스타일의 Amazon SageMaker 분산 훈련을 활용한 카테고리 자동 분류 시스템 모델 구축 사례

회사/팀 소개

카카오스타일은 모든 사람이 나만의 특별한 스타일을 가지고 있고, 내가 좋아하는 무언가를 발견했을 때의 즐거움이 일상을 더욱 나답게 만든다고 믿습니다. 사용자의 패션 뿐 아니라 일상에서의 모든 스타일을 위해 뷰티, 라이프 카테고리까지 확장하여 사용자의 즐거운 발견을 돕고 있습니다.
카카오스타일의 Vision & NLP(Natural Language Processing) 팀은 패션, 뷰티, 라이프 분야에서 컴퓨터 비전과 자연어 처리 기술을 활용하여 카카오스타일의 서비스를 발전시키고 사용자 경험을 향상시키는 데 초점을 맞추고 있습니다. 상품 메타 정보 분석, 추천 시스템, 리뷰 분석, 패션 트렌드 분석 등을 통해 사용자에게 더 나은 플랫폼 경험을 제공하기 위해 노력하고 있습니다.

프로젝트 소개

Vision & NLP 팀은 지그재그 서비스에 연동되어 자동으로 등록되는 상품들의 카테고리 입력율을 높이고자 이번 프로젝트를 시작했습니다. 카테고리 입력은 상품의 노출영역 등 여러 정책에 사용되기 때문에, 카테고리 입력율이 낮으면 문제가 발생할 수 있습니다.
상품은 상품명, 상품 설명, 썸네일 이미지 등 풍부한 메타정보를 가지고 있습니다. 이 정보를 이용해 지그재그 내 카테고리에 자동으로 분류하는 것이 이번 프로젝트의 목표였습니다.
그러나, 지그재그 앱 내에는 2천만 개 이상의 상품이 등록되어 있으며, 패션 및 라이프 카테고리로의 확장으로 카테고리의 추가 등 카테고리 구조가 변경될 때마다 모델을 갱신해야 합니다. 설계된 모델을 대량의 데이터로 훈련하기 위해 단일 GPU 환경에서는 수 일 이상이 소요됩니다. 이는 컴퓨팅 비용의 문제뿐 아니라 파라미터 튜닝 등 모델의 관리 측면에서도 치명적입니다. 따라서 분산 훈련을 통해 훈련 시간을 줄일 필요가 있었습니다.

분산 훈련 아키텍처

그림 1. 카카오스타일의 분산 훈련 아키텍처

최근 GPT, LLaMA 등 트랜스포머 디코더 기반 언어 모델이 놀라운 문장 생성 성능을 바탕으로 지식 검색, 번역 등 다양한 분야에서 사용되고 있습니다. 그러나, 컴퓨팅 비용, 미세 조정의 어려움, 구문 분석이 필요한 출력 구조 등의 문제로 분류 문제와 같은 태스크에서는 여전히 상대적으로 가벼운 BERT와 같은 트랜스포머 인코더 기반 언어 모델이 널리 사용되고 있습니다. 따라서 카카오스타일에서도 패션 커머스 분야의 언어 분포를 잘 이해할 수 있도록 서비스 내 데이터를 이용하여 BERT 모델을 경량화한 DistilBERT 모델을 훈련했습니다.
BERT 모델은 MLM(Masked Language Model), NSP(Next Sentence Prediction)을 통해 문장 구조와 단어의 의미, 문법 등 언어에 대한 정보를 사전 훈련합니다. 사전 훈련된 BERT 모델은 12개의 트랜스포머 층으로 구성된 모델로, 3억 개 이상의 파라미터를 가지고 있습니다. 하지만 수천만 개 이상의 상품과 수천만 건 이상의 리뷰 데이터로 분류, 개체명 인식 등 다양한 작업을 수행하기에 BERT 모델을 그대로 사용하는 것은 지나치게 많은 비용이 필요합니다.
따라서 BERT를 지식 증류 기법(knowledge distillation)을 이용해 경량화한 DistilBERT 모델을 훈련했습니다. 지식 증류는 작은 모델(학생)이 대형 모델(교사)를 모방하도록 훈련하는 기법으로 훈련된 DistilBERT는 BERT의 절반인 6개의 트랜스포머 층을 사용하면서도 거의 동일한 성능을 보입니다. 그럼에도 불구하고 전체 상품을 이용하여 모델을 훈련하는 것은 여전히 많은 시간이 소요합니다.
Amazon SageMaker는 데이터 병렬 처리 및 모델 병렬 처리를 위한 빌트인 분산 학습 라이브러리를 제공하며, 데이터 병렬 처리 라이브러리인 SageMaker DDP(Distributed Data Parallel)은 훈련 코드를 거의 수정하지 않고 데이터 병렬화를 수행할 수 있습니다. PyTorch뿐 아니라 자연어 처리 분야에서 널리 사용되는 허깅페이스(Hugging Face) 기반의 파이프라인도 손쉽게 분산 훈련이 가능합니다.
모델의 훈련 데이터는 카카오스타일 내에 등록된 모든 상품을 이용하여 구성합니다. 오랜 시간 데이터가 쌓인 패션 카테고리와는 다르게 카테고리 확장 이후 등록된 뷰티, 라이프 상품은 상대적으로 상품 수가 적습니다. 따라서 검증 데이터 세트는 층화 추출법(stratified sampling)을 통해 전체 데이터의 10%로 구성합니다.

아래는 SageMaker DDP를 사용하기 위해 모델을 정의하는 훈련 스크립트 train.py 예시입니다. 보다 자세한 사용 방법은 Modify a PyTorch Training Script와 참고 자료의 핸즈온을 확인하기 바랍니다.

import os
import sys
import torch
import torch.distributed as dist
from torch.nn.parallel.distributedimport DistributedDataParallel as DDP
from transformers import AutoTokenize, AutoModelForSequenceClassification
import smdistributed.dataparallel.torch.torch_smddp
 
# Environment variables set by torch.distributed.launch or torchrun
world_size =int(os.environ['WORLD_SIZE'])
rank = int(os.environ['RANK'])
local_rank = int(os.environ['LOCAL_RANK'])
 
# initialize the process group
dist.init_process_group(backend='smddp', rank=rank, world_size=world_size)
torch.cuda.set_device(local_rank)
device = torch.device("cuda", local_rank)
 
tokenizer = AutoTokenizer.from_pretrained('distilbertmodel')
model = AutoModelForSequenceClassification.from_pretrained('distilbertmodel')
model = DDP(
    model.to(device),
    device_ids=[local_rank],
    output_device=local_rank
).to(local_rank)

훈련 스크립트 train.py를 작성하면, SageMaker SDK를 이용하여 SageMaker 훈련 Job을 정의하고 훈련을 시작합니다. 다음 예시는 8개의 NVIDIA V100 GPU를 사용하는 ml.p3.16xlarge 인스턴스를 이용하여 모델을 훈련하는 방법을 보여줍니다. 보다 자세한 사용 방법은 Launch a SageMaker Distributed Training Job Using the SageMaker Python SDK를 확인하기 바랍니다.

import sagemaker
from sagemaker.debugger import ProfilerRule, rule_configs
from sagemaker.pytorch import PyTorch

role = sagemaker.get_execution_role()
rules = [
    ProfilerRule.sagemaker(rule_configs.ProfilerReport()),
]

distribution = {
    'smdistributed':{
        'dataparallel':{
            'enabled': True
        }
    }
}

job_name = 'ddp-training'
instance_type = 'ml.p3.16xlarge'  # 'ml.p3.16xlarge', 'ml.p3dn.24xlarge', 'ml.p4d.24xlarge', 'local_gpu'
instance_count = 2    # 사용할 인스턴스 개수
max_wait = None
max_run = 48*60*60    # 최대 48시간 훈련

hyperparameters = {
    'epochs': 5,
    'batch_size': 64,
    'test_batch_size': 128
}

sagemaker_session = sagemaker.Session()
source_dir = 'src'

image_uri = "763104351884.dkr.ecr.ap-northeast-2.amazonaws.com/pytorch-training:1.10.2-gpu-py38-cu113-ubuntu20.04-sagemaker"
estimator = PyTorch(
    entry_point='train.py',
    source_dir=source_dir,
    role=role,
    image_uri=image_uri,
    instance_count=instance_count,
    instance_type=instance_type,
    distribution=distribution,
    rules=rules,
    max_wait=max_wait,
    max_run=max_run,
    hyperparameters=hyperparameters,
    sagemaker_session=sagemaker_session
)

estimator.fit(
    job_name=job_name,
    wait=False
)

SageMaker DDP는 ml.p3.16xlarge, ml.p3dn.24xlarge, ml.p4d.24xlarge 인스턴스 및 로컬 GPU를 이용하여 훈련이 가능합니다. 카카오스타일의 카테고리 분류 모델은 2개의 ml.p3.16xlarge 인스턴스를 이용하여 총 16장의 GPU를 사용하며, 과도한 비용 청구 방지를 위해 최대 48시간 동안 훈련을 수행합니다.
분산 훈련을 위한 개발 환경은 도커 이미지(Docker Image)를 통해 불러올 수 있습니다. SageMaker는 PyTorch 버전 등 자주 사용되는 환경에 맞춰 최적화된 빌트인 도커 이미지를 제공합니다. 사용 가능한 이미지 목록은 Deep Learning Containers 레포지토리에서 확인할 수 있습니다.

분산 훈련 팁

SageMaker DDP 기반 분산 훈련은 매우 간단한 방법으로 구현이 가능하지만, 예상치 못한 오류나 의도치 않은 방향으로 훈련이 진행될 수도 있습니다. 따라서, 아래의 분산 훈련 팁들을 적용해 보세요.

일부 데이터 세트 및 로컬 모드로 디버깅: 분산 훈련은 대부분 대량의 데이터 세트나 대규모의 모델을 이용하여 훈련이 이루어지므로, 한 번의 오류로도 많은 비용 손실을 야기합니다. 따라서 데이터 세트의 일부만을 이용하여 훈련 작업이 오류 없이 완전히 종료되는지 테스트하는 것을 추천드립니다. 그리고 로컬 모드를 적용하면 디버깅을 더 빠르게 수행할 수 있습니다. 로컬 모드는 훈련 인스턴스를 프로비저닝하지 않고 훈련 인스턴스와 동일한 컨테이너로 로컬 환경에서 모델 개발과 디버깅을 수행하는 모드로 train_instance_type='local_gpu'로 설정하면 됩니다.
혼합 정밀도(Mixed Precision): 대형 모델은 많은 GPU 메모리가 필요하므로 미니배치 크기를 늘리기 어렵습니다. 하지만 파라메터 연산에 16비트 부동 소수점을 사용하되, 수치 안정성을 위해 특정 파라메터는 32비트 부동 소수점을 사용하는 혼합 정밀도를 사용하면 미니배치 크기를 늘리면서 훈련 시간을 절감할 수 있습니다.
훈련 데이터 준비 시간 개선: 대부분의 데이터 분산 훈련은 대량의 데이터 세트를 이용합니다. 만약 데이터 세트가 S3 버킷에 저장되어 있다면 데이터를 훈련 인스턴스에 내려받는 과정도 많은 시간이 소요됩니다. 이럴 때에는 훈련 데이터를 스트리밍하는 FastFile 모드나 Amazon FSx for Lustre 파일 시스템을 이용하여 훈련 데이터를 준비하는 시간을 줄여보세요. 자세한 내용은 SageMaker 개발자 문서(Access Training Data)를 참조하기 바랍니다.
체크포인트 및 관리형 스팟 훈련(managed spot training)으로 비용 절감하기: 특정 주기(예: 매 에폭)마다 중간 훈련 결과를 /opt/ml/checkpoints로 저장합니다. 훈련 Job을 정의할 때 아래 코드 스니펫과 같이 스팟 인스턴스 옵션을 활성화하면 훈련 비용을 최대 90%까지 절감할 수 있는 SageMaker 관리형 스팟 훈련 기능을 사용할 수 있습니다. 훈련 작업이 중단되면 자동으로 훈련 인스턴스의 체크포인트 파일이 S3의 checkpoint_s3_uri 경로로 복사되고 새로운 스팟 인스턴스에서 훈련이 재개될 때 S3의 체크포인트가 인스턴스의 스토리지로 자동으로 복사됩니다.

chkpt_s3_path = f's3://[YOUR-BUCKET]/checkpoints-{curr_time}'
use_spot_instances = False
max_run = 48*60*60 # 최대 48시간 훈련

if use_spot_instances:
 max_wait = 72*60*60 # 최대 72시간 훈련: spot instance waiting + max runtime
else:
 max_wait = None

sm_estimator = PyTorch(
    ...
    use_spot_instances = use_spot_instances,
    max_wait = max_wait,
    max_run = max_run, 
    checkpoint_s3_uri = chkpt_s3_path,
    checkpoint_local_path ='/opt/ml/checkpoints'
)

결론

Amazon SageMaker의 DDP를 이용하여 대량의 데이터로도 빠르게 모델을 훈련하고 실험할 수 있었습니다. 로컬 환경에서의 분산 훈련은 개발환경을 구축하는 것 만으로도 큰 비용이 발생합니다. 단 몇 줄의 코드 수정만으로 쉽게 분산 훈련을 구현한다는 것은 훈련 시간의 감소 그 이상의 가치가 있습니다.
단일 V100 GPU를 사용하는 환경에서 보름 이상 걸리는 훈련을 16개의 V100 GPU를 제공하는 환경에서 SageMaker DDP를 이용한 최적화된 데이터 분산 훈련으로 30시간만에 마칠 수 있었습니다.
카테고리 분류 작업은 비교적 간단한 작업으로, 데이터 전처리 방법에 따라 모델의 성능이 크게 좌우되는 경우가 많습니다. 데이터 전처리 후 긴 시간 동안 모델의 학습 결과를 기다리는 대신 데이터를 전처리하고 준비하는 데에 더 많은 시간을 사용하여 모델을 훈련할 수 있었습니다.
아울러, 모델을 학습하며 맞닥뜨린 여러 문제에 적극적으로 도움 주신 AWS에 감사의 마음을 전합니다.

참고 자료

AWS Summit Korea 2021: Amazon SageMaker를 통한 대용량 모델 훈련 방법 살펴보기
AWS Summit Korea 2022: 실전! 대용량 데이터를 이용한 초거대 모델 학습 환경 만들기
AWS re:Invent 2022: Train and host foundation models with PyTorch on AWS (AIM404)
SageMaker Distributed Training Step-by-Step: https://github.com/daekeun-ml/sm-distributed-training-step-by-step

AWS 기술 블로그