KT, Amazon SageMaker를 이용한 ViT 기반 Food Tag 모델의 학습 시간 단축 여정

KT의 ‘AI Food Tag’는 사진 속 음식의 종류와 영양 성분을 알려 주는 인공지능(AI) 기반 식이 관리 솔루션입니다. KT가 개발한 Vision 모델은 레이블(Label)이 없는 대용량 이미지 데이터로 학습한 사전 학습 모델이 적용되어, 다양한 음식들의 영양 성분과 칼로리 정보를 분석하여 당뇨 등 만성질환 환자의 식단 관리에 도움을 줄 수 있습니다. 이러한 ‘AI Food Tag’ 모델의 학습 성능 향상 및 상용화를 위한 모델 경량화를 위해 AWS 인프라 활용 및 Amazon SageMaker의 도입 여정을 설명하고자 합니다.

과제 소개 및 문제 정의

KT가 개발한 ‘AI Food Tag’ 모델은 Vision Transformers (ViT) 기반으로 학습이 되었으며, 학습 시 정확도를 높이기 위해 기존 Vision 모델 대비 많은 모델 파라미터를 사용하여 정확도를 개선하였습니다. KT는 이 모델의 상용화 시 모델을 효율적으로 배포하기 위한 목적으로 Knowledge distillation 기법을 통한 경량화 모델을 개발하고 있습니다. Knowledge distillation 기법에서 기존 모델을 Teacher 모델로 설정하고, 경량화 모델은 Student 모델로 학습을 하게 되는데, 경량화 모델은 기존 모델의 성능 저하를 최소화하면서 모델 파라미터의 크기를 줄여서 작은 크기의 GPU 메모리에도 모델 배포가 가능하게 됩니다. 일반적으로 아래 같은 구조로 학습을 진행하게 됩니다.

Teacher 모델은 학습을 하지 않지만, Student 모델이 학습하면서 Loss (손실)을 계산하는데 Teacher 모델의 결과를 함께 사용하게 됩니다. 따라서, Teacher 모델과 Student 모델이 모두 GPU에 올라가야 합니다.

KT는 경량화 목적의 Student 모델을 학습하기 위해 내부적으로 2장의 GPU (A100 80GB)를 사용하여 1차 학습을 진행하였으나, 학습 속도는 300 epoch을 학습하는데 약 40일이 소요되었습니다. 이에 학습 속도 향상을 통해 경량화 모델 확보를 가속화하고자 AWS와 협력을 통해 추가적인 모델 학습을 진행하였고, 확장성 높은 AWS 클라우드를 사용하여 더욱 빠르게 모델 학습을 성공적으로 수행하였습니다. 이번 PoC에서는 Amazon SageMaker Training과 함께 데이터 병렬 처리를 위한 SageMaker’s Data Parallelism Library, 그리고 리소스의 병목현상을 확인하기 위해 Amazon SageMaker Debugger와 SageMaker Profiler를 활용하였습니다.

SageMaker를 이용한 분산 학습 환경의 구축

Amazon SageMaker Training은 AWS에서 관리하는 머신러닝 학습 환경으로 머신러닝 학습에 필요한 Linux 및 python 패키지가 설치된 다양한 딥러닝 프레임워크 버전 별로 분산 컴퓨팅을 위한 학습 클러스터를 제공하며, 모델 학습을 실행하고자 하는 데이터 과학자나 ML 엔지니어는 Docker를 관리하거나 학습 클러스터를 구성하는 부담 없이 학습을 실행할 수 있습니다.

1일 워크샵을 통해 AWS의 KT 계정 내에 Amazon SageMaker 기반의 분산 학습 클러스터를 구축하고, KT의 학습 스크립트에 PyTorch SageMaker DDP (Distributed Data Parallel) 라이브러리를 적용하였고, 2개의 ml.p4d.24xlarge를 활용하여 학습 테스트까지 완료할 수 있었습니다. 테스트 과정에서 처음에는 PyTorch DDP의 native한 backend인 NCCL (NVIDIA Collective Communications Library)을 사용하였으나, 다음과 같은 오류가 발생하였습니다.

File "/opt/conda/lib/python3.9/site-packages/torch/autograd/__init__.py", line 197, in backward
Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
RuntimeError: CUDA error: CUBLAS_STATUS_ALLOC_FAILED when calling `cublasCreate(handle)`

위의 오류를 해결하기 위해서는 기존의 batch size를 줄여야 하기에 오히려 모델의 학습 속도가 느려질 수 있어 PyTorch DDP 대신 SageMaker DDP library를 적용하여 다시 실행하였고, 그 결과 동일한 batch size에서도 에러 없이 학습이 시작되는 것을 확인하였습니다. PyTorch DDP에서 SageMaker DDP로 변경하려면 torch_smddp 패키지를 선언하고 backend를 아래와 같이 smddp로 변경하면 됩니다.

import smdistributed.dataparallel.torch.torch_smddp

dist.init_process_group(backend='smddp',
                            rank=args.rank,
                            world_size=args.world_size)

SageMaker Debugger/Profiler를 이용한 학습 속도 저하의 원인 분석

학습 시 상세한 속도 측정을 하기 위해 data loader, forward와 backward 함수 별 iteration단위로 측정하였습니다.

1 iter time - dataloader : 0.00053 sec, forward : 7.77474 sec, backward: 1.58002 sec
2 iter time - dataloader : 0.00063 sec, forward : 0.67429 sec, backward: 24.74539 sec
3 iter time - dataloader : 0.00061 sec, forward : 0.90976 sec, backward: 8.31253 sec
4 iter time - dataloader : 0.00060 sec, forward : 0.60958 sec, backward: 30.93830 sec
5 iter time - dataloader : 0.00080 sec, forward : 0.83237 sec, backward: 8.41030 sec
6 iter time - dataloader : 0.00067 sec, forward : 0.75715 sec, backward: 29.88415 sec

각 iteration에 대한 표준 출력에서의 시간을 살펴보면, backward의 수행 시간이 iteration마다 심하게 변동하는 것을 확인하였습니다. 이러한 일관되지 않은 학습 속도의 원인을 찾기 위해 먼저 Amazon SageMaker Debugger를 활용하여 리소스 병목 현상을 파악하고자 했습니다. SageMaker Debugger는 SageMaker의 학습 클러스터에서 학습 작업을 디버깅하고, 클러스터의 CPU, GPU, network, I/O 등의 리소스 상태를 설정한 초 단위로 확인할 수 있습니다.

SageMaker Debugger에서 얻을 수 있는 다양한 정보 중에서 CPU 사용률과 인스턴스 별 CPU/GPU 사용률 heatmap을 확인하였습니다. CPU 활용률에서는 일부 CPU가 100%로 사용되고 있음을 알 수 있었고, heatmap에서도 동일한 상황을 파악할 수 있었습니다. 색깔이 짙은 경우가 사용률이 높은 것을 의미하는데, 학습의 전체 과정에서 많은 CPU core의 사용률이 높은 것을 알 수 있으며, 시간이 지날수록 GPU 사용률도 일정하지 않는 모습을 확인할 수 있었습니다. 결국 느린 학습 속도의 원인 중 하나가 CPU 병목 현상이라고 의심하였고, KT의 학습 스크립트 상에서 CPU 병목을 유발하는 코드를 검토하였습니다. 가장 의심스러운 부분은 Dataloader에서 num_workers 값이 큰 것을 확인하였는데, CPU 사용률을 줄이기 위해 이 값을 0 또는 1로 변경하였습니다. 그리고 나서 학습 작업을 재수행하여 SageMaker Debugger에서 결과를 확인하였습니다.

단순히 num_workers를 변경함으로써 CPU 사용률이 크게 감소하고 GPU 사용률이 전반적으로 증가하는 것을 확인할 수 있었습니다. CPU 사용률 개선 및 확인 이후, GPU에서 최적화할 수 있는 부분을 찾기 위해 SageMaker Profiler을 사용하여 테스트를 추가로 진행하였습니다. SageMaker Profiler는 GPU와 CPU 사용률 지표, GPU/CPU 의 커널 소비 등에 대한 추적 기능을 학습 스크립트 내 함수 단위로 제공하여 사용률에 대한 가시성을 통해 최적화 기회를 확인할 수 있도록 합니다. 이를 통해 사용자는 어떤 함수가 리소스를 많이 사용하고 있는지 파악할 수 있습니다. 먼저 SageMaker Profiler를 사용하려면 아래와 같이 SageMaker SDK를 사용하여 학습 클러스터를 생성하는 함수에 ProfilerConfig를 추가해야 합니다.

from sagemaker import ProfilerConfig, Profiler
from sagemaker.debugger import (ProfilerRule, rule_configs)

rules=[ProfilerRule.sagemaker(rule_configs.ProfilerReport())]
profiler_config = ProfilerConfig(profile_params = Profiler(cpu_profiling_duration=3600))
from sagemaker.pytorch import PyTorch

region_name = 'us-west-2'
image_uri=f'763104351884.dkr.ecr.{region_name}.amazonaws.com/pytorch-training:2.0.0-gpu-py310-cu118-ubuntu20.04-sagemaker'
estimator = PyTorch(
                    entry_point='train.py',
                    source_dir='src',
                    role=role,
                    image_uri=image_uri,
                    instance_count=4,
                    instance_type='ml.p4d.24xlarge',
                    distribution={'smdistributed': {'dataparallel': {'enabled': True}}},
                    profiler_config=profiler_config,
                    hyperparameters=hyperparameters,
                    sagemaker_session=sagemaker_session,
                   )

또한, 학습 스크립트에서 SageMaker Profiler 패키지를 선언하고, KT의 학습 스크립트 내 상세한 파악이 필요한 코드 앞에 SageMaker Profiler의 annotate함수를 추가로 선언하면 됩니다. 아래는 학습 스크립트에서 SageMaker Profiler에 대해 선언해야 하는 코드들의 예시를 보여주고 있습니다.

import smppy

SMProf = smppy.SMProfiler.instance()
config = smppy.Config()
    
config.profiler = {
    "EnableCuda": "1",
}

SMProf.configure(config)
SMProf.start_profiling()
…
with smppy.annotate("Forward"):
    student_out = student_model(inp)
with smppy.annotate("Backward"):
   loss.backward()
…
SMProf.stop_profiling()

위의 코드를 추가한 학습 스크립트를 이용하여 학습 작업을 다시 실행하면 아래와 같이 GPU kernel이 소비한 연산에 대한 정보를 얻을 수 있게 됩니다. KT의 학습 스크립트의 경우에는 1 epoch을 실행하여 아래 결과를 얻을 수 있었습니다.

SageMaker Profiler의 결과 중에서 GPU 커널을 가장 많이 소비하는 상위 5개 연산을 확인한 결과, 기존 KT 학습 스크립트의 경우 GEMM (General matrix multiplication)이 동작하는 행렬곱 연산에서 가장 많은 시간이 소비된다는 것을 알 수 있었습니다.

학습 시간 절감을 위한 방법 적용

행렬곱의 실행 시간을 절감하기 위한 다양한 방법을 검토하였고, 이 중 2가지의 PyTorch 함수를 적용해 보았습니다.

Shard Optimizer States with ZeroRedundancyOptimizer

ZeroRedundancyOptimizer는 Optimizer 상태를 샤딩 (Sharding) 하여 DDP (Distributed Data Parallel)의 프로세스당 메모리 사용량을 줄이는 DeepSpeed/ZeRO 의 기법을 PyTorch DDP에서 활용할 수 있게 제공하고 있습니다. DDP는 backward 패스에서 동기화된 gradient를 이용하여 모든 Optimizer의 복제본이 반복되면서 동일한 파라미터와 gradient 값으로 작동하게 되는데, ZeroRedundancyOptimizer는 모든 모델 파라미터를 가지고 있는 대신 서로 다른 DDP 프로세스에 대해서만 샤딩을 통해 각 optimizer 상태를 유지하도록 하여 메모리 사용량을 줄이는 효과를 얻습니다.

    student_optimizer = ZeroRedundancyOptimizer(
            student_model.parameters(),
            optimizer_class=torch.optim.AdamW, 
            lr=initial_lr
        )

사용 방법은 기존의 Optimizer를 optimizer_class에 두고 나머지 모델 파라미터와 learning rate을 파라미터로 하여 ZeroRedundancyOptimizer를 선언하면 됩니다.

Automatic Mixed Precision

AMP (Automatic Mixed Precision)은 일부 연산에 대해 torch.float32 데이터 타입을 적용하고 일부 연산은 torch.bfloat16 또는 torch.float16를 이용하여, 혼합된 정밀도 값들을 편리하고 빠르게 연산하면서 메모리 공간을 줄이는데 사용합니다. 특히, torch.bfloat16은 torch.float32 의 지수 (exponent) 비트와 동일하기 때문에 일반적으로 가수 (fraction) 비트보다 지수 (exponent) 비트에 더 민감한 딥러닝 모델의 연산 시 손실을 최소화하면서 빠르게 학습을 할 수 있습니다. torch.bfloat16은 ml.p4d.24xlarge, ml.p4de.24xlarge, ml.p5.48xlarge 인스턴스와 같이 A100 GPU 이상의 Nvidia 아키텍처 (Ampere, Hopper)에서만 동작합니다.

  with torch.cuda.amp.autocast(dtype="torch.bfloat16"):
      teacher = teacher_model(input_data)
                student = student_model(input_data)
                loss = loss(teacher, student, target)
                loss.requires_grad_(True)
            loss.backward()
            student_optimizer.step() 
            student_optimizer.zero_grad(set_to_none=True)

AMP를 적용하려면, 위의 코드와 같이 torch.cuda.amp.autocast를 학습 스크립트에 선언한 다음, dtype을 torch.bfloat16로 선언하면 됩니다.

SageMaker Profiler에서 적용 결과 확인

학습 스크립트에서 2개의 함수를 적용하여 1 epoch 동안 다시 학습 작업을 수행한 결과, SageMaker Profiler에서 GPU 커널 소비에 대한 상위 5개 연산 시간을 확인하였습니다.

2개의 PyTorch 함수를 적용하기 전에 가장 상위에 있던 GEMM 연산은 GPU 커널 소비에 대한 상위 5개 연산에서 없어졌으며, 일반적으로 분산 학습에서 발생하는 ReduceScatter 연산으로 대체된 것을 볼 수 있습니다.

KT 경량화 모델의 학습 속도 결과

추가적으로, 2개의 PyTorch 함수를 적용하면서 얻게 된 메모리 절감 효과를 고려하여 학습 batch size를 128 정도 추가로 늘릴 수 있었으며, 이에 따라 최종 batch size는 기존 1024에서 1152를 사용할 수 있게 되었습니다. 이후 Student 모델의 학습은 1일 동안 210 epoch을 실행할 수 있었고, 기존 KT 내부 학습 환경과 Amazon SageMaker 간의 학습 시간과 속도 향상은 아래 표로 정리하였습니다.

[표] 기존 KT 내부 학습 환경 vs. Amazon SageMaker

학습 환경	학습 GPU	GPU 수	학습 시간 (hour)	Epoch	hour/epoch	hour/epoch 감소율
KT의 내부 학습 환경	A100 (80GB)	2	960	300	3.20	29
Amazon SageMaker	A100 (40GB)	32	24	210	0.11	1

만일 위 결과를 A100 (40GB) GPU 1장을 사용했다고 가정하면, KT의 내부 학습 환경에서는 1 epoch 당 12.8 시간이 걸리고, Amazon SageMaker는 1 epoch 당 3.66 시간이 소요되는 것을 알 수 있습니다. 결과적으로, 동일한 GPU를 사용하더라도 Amazon SageMaker에서 기존 KT의 내부 학습 환경보다 3.5배의 빠른 속도로 학습 결과를 얻을 수 있었고, 실제 PoC에서는 Amazon SageMaker의 확장성을 통해 4개의 ml.p4d.24xlarge (32장 A100 GPU)를 사용하여 기존보다 29배 빠른 학습 속도로 학습 작업을 완료할 수 있었습니다.

마무리

KT R&D의 ‘AI Food Tag’ 모델에 대한 PoC는 AWS의 AI/ML Concierge Program을 통해 KT와 AWS 간의 협업한 결과이며, KT 융합기술원 AI2XL연구소의 박상민 팀장님(비전AI서빙기술팀)은 “최근 Vision 분야에도 Transformer 기반의 모델들이 많아지면서 모델의 파라미터가 증가함에 따라 필요한 GPU 메모리가 늘어나고 있습니다. 해당 이슈에 대해 해결을 위해 KT는 경량화 기술을 사용하고 있고, 한번 학습하는데 한 달 정도의 많은 시간이 소요되고 있었습니다. 우리는 이번 AWS와의 PoC에서 Amazon SageMaker Profiler와 Debugger를 통해 어떤 부분의 리소스가 병목이 있는지 확인할 수 있었고, 이를 해결한 다음 Amazon SageMaker의 Distributed Data Parallelism 라이브러리를 사용하여 4개의 ml.p4d.24xlarge 인스턴스에서 최적화된 학습 스크립트로 약 하루의 시간으로 학습을 마칠 수 있었습니다.”와 같이 학습 성능 개선 효과에 대한 Feedback을 주셨습니다.

본 KT R&D의 Vision 모델에 대한 협업을 바탕으로, 향후 KT R&D 내 다양한 AI/ML 연구 과제를 대상으로 AWS의 확장성, 탄력성 있는 GPU 자원 제공과 Amazon SageMaker를 통한 모델 개발 및 서비스 생산성 향상을 위해 지속적인 협업 기회를 마련할 예정입니다.

Amazon SageMaker의 관련 기능에 대해 자세히 알아 보려면 다음을 확인하십시오.

사공락

사공락 선임연구원은 KT 연구소에서 Vision AI에서 다양한 분야에 대해 연구/개발을 수행하면서 주로 얼굴 쪽을 관련으로 얼굴 속성(성별/ 안경,모자 등 착용 유무) / 얼굴 인식 기술 등에 대해 진행을 해왔었습니다. 현재는 커져가는 Vision model들에 대한 경량화 기술 관련해서 수행하고 있습니다.

AWS 기술 블로그