브레인 스토리
torchgpipe가 탄생하기까지
2019/08/13 by 이수경, 이흥섭

카카오브레인 AutoML 연구팀 소속 이흥섭 연구원과 정명룡 연구원이 파이토치(PyTorch)[1]GPipe 라이브러리인 ‘torchgpipe’를 대중에 공개했습니다. 이흥섭 연구원을 만나 어떤 시행착오를 거쳐 torchgpipe가 탄생했는지 그 뒷이야기를 들어봤습니다. 이번 글에서는 두 사람이 GPipe에 관심을 가진 계기와 GPipe의 특징, torchgpipe 개발 프로젝트에서 주력한 부분에 관한 이야기를 다루고자 합니다. 


[ 그림 1 ] 카카오브레인의 정명룡 연구원(왼쪽)과 이흥섭 연구원이 기념사진을 촬영하고 있다.


병렬화에 매료되다

torchgpipe 프로젝트 착수에 앞서 이 연구원과 정 연구원은 딥러닝 엔지니어가 갖춰야 할 필수 역량을 빠르게 습득할 수 있는 과제를 수행했습니다. 과제명은 직접 구현한 ResNet에 데이터 병렬화 기법(data parallelism)을 적용해 ILSVRC[2] 2012 데이터셋을 수 시간 내로 훈련하기였습니다.

"일반적인 프로그램을 개발할 때 즉각적인 성공 혹은 실패라는 피드백은 늘 중요했습니다. 그런데 이번 과제에서는 한 아이디어를 검증하기까지 짧게는 수 분, 길게는 수일 걸리는 경우가 많았죠. 시간을 효율적으로 활용하려면 한 번에 여러 아이디어를 동시에 검증할 수밖에 없습니다. 하지만 여러 아이디어에 집중력이 분산되다 보니 잦은 실수가 발생했습니다. 자연스럽게 데이터 병렬화에 강렬한 인상을 받게 됐습니다. 동일 연산을 수행하는 여러 GPU에 데이터셋을 나눠 올려서 데이터 처리 시간을 획기적으로 줄여주니까요.”

두 사람이 다음으로 주목한 기법은 모델 병렬화(model parallelism)입니다. 딥러닝 모델은 크기가 클수록 성능(추론 정확도)이 좋다고 알려져 있는데요, 그렇다고 모델의 크기를 무한정 키울 수는 없습니다. 모델 규모에 비례해 은닉층(hidden layer)과 매개변수가 늘어나는 만큼 더 많은 메모리가 필요해지기 때문이죠. 입력 및 출력 데이터가 클수록 차지하는 메모리 공간도 커집니다.

모델 병렬화는 바로 1개의 GPU 메모리 크기를 초과하는 모델을 여러 부분으로 나누어 각각 서로 다른 GPU에 배정하는 기법을 말합니다. 이를테면 분할된 첫번째 모델은 첫번째 GPU에, 두번째 모델은 두번째 GPU에 올리는 거죠. 좀 더 자세히 설명하자면 다음과 같습니다. 학습에 필요한 메모리가 60GB인 모델이 있습니다. 메모리 크기가 20GB인 1개의 GPU로는 이 모델을 훈련시키지 못하겠죠? 3개의 GPU에 모델을 20GB씩 각각 나눠 올리고 분산 훈련을 시도하는 방법을 고려해볼 수 있을 것입니다.

[ 그림 2 ] 데이터 병렬화, 모델 병렬화의 차이를 시각화한 이미지 © 이흥섭 연구원

모델 병렬화는 의료영상 분석 태스크 활용에 큰 관심을 둔 AutoML 연구팀의 주 관심사이기도 했습니다. 굉장히 큰 의료 데이터를 모델을 효율적으로 훈련시킬 방법론이라고 생각했던 거죠. 이미지 분류 학습에 쓰이는 데이터셋인 CIFAR가 32×32픽셀, ILSVRC 2012는 평균 473×406픽셀 크기인 반면, 의료영상인 CT는 512×512×512픽셀, 병리학에서 쓰이는 표준 규격은 20만×20만 픽셀 크기입니다. 보통의 학습 이미지와 비교해 그 크기가 훨씬 큰 의료영상을 훈련하는 데는 더 많은 메모리 자원이 필요하죠. 결과적으로, 1개의 GPU(또는 TPU)[3]만으로는 거대한 모델을 훈련하기가 어렵기에 모델을 단순화해 그 크기를 줄이거나 모델 병렬화를 통한 분산 훈련을 시도했습니다.

하지만 딥러닝 모델 학습에서의 모델 병렬화는 GPU를 효율적으로 사용하지 못합니다. 이 문제는 대부분의 딥러닝 모델이 앞쪽 층에서 연산을 마쳐야 뒤쪽 층에서 연산을 이어서 하는 순차적인 구조를 갖추고 있다는 특징에서 비롯됩니다. [그림 3]에서처럼 첫번째 GPU가 작업을 마치기 전까지는 두번째 GPU가 작업을 진행할 수 없죠. 결과적으로는 모델 병렬화에 필요한 3개의 GPU를 확보해 메모리 크기를 3배로 늘려도 연산 속도는 GPU 1대를 쓸 때와 마찬가지입니다.

[ 그림 3 ] 순차적인 구조를 갖는 딥러닝 모델에서는 모델 병렬화만으로는 연산 속도 향상을 기대할 수 없다. © 이흥섭 연구원

두사람은 1회 학습에 활용하는 데이터셋 크기를 작게 줄여 최대한 많은 GPU가 동시에 작업을 수행할 수 있도록 하는 파이프라인 병렬화(pipeline parallelism)에 자연스럽게 눈길을 돌리게 됐습니다. 최근 구글 브레인(Google Brain)이 공개한 GPipe 논문은 이를 효과적으로 풀어내는 데 성공했습니다.

[ 그림 4 ] 파이프라인 병렬화를 시각화한 이미지. 이 기법은 모델을 여러 부분으로 나누어 각각 GPU에 올리며, 1회 학습에 사용하는 데이터셋 크기를 최대한 작게 줄인다. © 이흥섭 연구원

GPipe의 핵심, ‘파이프라인 병렬화’와 ‘체크포인팅’

텐서플로로 구현된 GPipe는 1개의 GPU로 학습시킬 수 없는 거대한 모델을 가능한 효율적으로 훈련하기 위한 모델 학습 메커니즘입니다. GPipe를 도입하면 메모리 사용량이 큰 모델을 여러 GPU에서 병렬 처리해 GPU를 효율적으로 사용하면서 빠른 학습이 가능해집니다. 실제로 모델 크기를 키운 AmoebaNet[4]-B(6, 512)[5]로 ILSVRC 2012 데이터셋을 훈련시킨 결과, 최고 수준의 정확도를 확인할 수 있었습니다([표 1] 참고).

[ 표 1 ] ImageNet ILSVRC 201 검증 데이터셋을 상대로 최신 모델의 성능을 비교해본 결과표. 모델 크기가 클수록 정확도가 높아짐을 확인해볼 수 있다.

이처럼 규모가 큰 모델을 효율적으로 훈련시키는 GPipe는 파이프라인 병렬화와 체크포인팅(checkpointing)이라는 2가지 기법을 조합합니다.

순차 구조를 갖춘 딥러닝 모델을 훈련할 때 모델 병렬화로는 여러 GPU를 효율적으로 사용하지 못한다는 한계를 앞단에서 지적했습니다. 파이프라인 병렬화는 이 한계를 극복합니다. 여러 부분으로 나눈 모델을 각각 GPU에 올리는 방식은 모델 병렬화와 같습니다. 다른 점은 모델 훈련이 여러 GPU에서 일정 시간 동안 동시에 연산될 수 있도록 [그림 5]처럼 미니배치를 여러 개로 쪼갠 마이크로배치(micro-batch)를 학습에 활용한다는 거죠. 마이크로배치를 활용하면 중첩되어 처리되는 작업량이 늘어나 같은 양의 데이터를 더 짧은 시간 안에 처리할 수 있게 됩니다.

[ 그림 5 ] 1개의 미니배치를 더 작은 N개의 마이크로배치로 잘라낸다. © 이흥섭 연구원

“각 GPU가 이미지 100장짜리 미니배치를 처리하는 데 10초가 걸린다고 가정해 봅시다. 그러면 두번째 GPU는 10초 뒤 작업을 이어서 처리할 수 있을 것입니다. 그렇다면 10장짜리 마이크로배치 10개를 훈련하는 건 어떨까요? 같은 미니배치를 학습하는 상황인데도 두번째 GPU는 9초 더 빠르게 일을 시작할 수 있습니다. 이후 수 초간 두 GPU는 동시에 일을 처리하게 됩니다. 모델 병렬화에선 한 번에 한 GPU만 쓸 수 있었다면, 파이프라인 병렬화로는 여러 GPU를 최대한 동시에 동원할 수 있기에 더 효율적인 학습이 가능해집니다.”

[그림 6]을 보면 마이크로배치의 크기가 작아질수록 다음 GPU가 더 빨리 일을 시작할 수 있어 유휴 시간이 점차 줄어듦을 확인할 수 있습니다. 하지만 마이크로배치의 크기를 무한정 줄일 수는 없습니다. 그 크기가 지나치게 작으면 CPU가 GPU에 충분한 일감을 제공하지 못해서 되려 GPU 사용 효율이 떨어질 수도 있기 때문입니다. 모델에 맞는 적절한 마이크로배치 크기를 찾는 건 개발자의 몫이죠.

[ 그림 6 ] 마이크로배치 크기가 작을수록 GPU의 유휴 시간이 줄어듦을 확인할 수 있다. © 이흥섭 연구원

두번째 기법인 체크포인팅은 은닉층 결과값 일부를 버려서 모델이 차지하는 메모리 크기를 줄이는 방법을 가리킵니다. GPipe는 모델을 나눈 부분에 체크포인트를 적용하고 분리된 모델을 서로 연결하는 은닉층만 메모리에 올립니다. 다시말해 순전파(forward propagation)[6] 단계에서 분리된 모델 내부에서의 은닉층 결과값을 기억하지 않고 버리죠. 이 값은 역전파(backward propagation)[7] 단계에서 입력층 결과값과 계산 로직을 가지고 재계산할 수 있습니다.

결과론적으로는 모든 은닉층 결과값을 메모리에 올려둘 필요가 없기 때문에 모델이 훈련에 사용하는 메모리 공간을 줄일 수 있게 됩니다. 역전파 과정에서 순전파 계산을 한 번 더 진행하기에 모델 학습 속도가 약 25% 느려집니다만, 마이크로배치 수에 반비례하게 메모리 사용량을 줄일 수 있어 속도를 희생할 만한 가치가 있다고 볼 수 있겠습니다.

[ 그림 7 ] 체크포인팅을 적용하면 모델이 차지하는 메모리 크기를 획기적으로 줄일 수 있다. © 이흥섭 연구원


torchgpipe가 주력한 ‘이것’

이런 장점에도 불구하고 GPipe는 여전히 진입 장벽이 높습니다. 논문 속 코드를 구현하려면 관련 전문 지식을 갖춰야합니다. 모델을 여러 GPU에 어떻게 나눠 올릴지, 매개변수(parameter)와 기울기를 언제 어떻게 동기화할지, 어떤 예외사항을 대비해야 할지 등 병렬화에 필요한 모든 요소를 고려해야 합니다. 학습속도와 메모리 용량 측면에서 손해를 보더라도 1개의 GPU만 훈련에 사용하는 경우가 많았던 건 바로 이런 이유 때문입니다.

두 연구원은 딥러닝 개발자가 모델 개발에만 집중할 수 있도록 기술 구현에 따르는 복잡한 작업을 간소화하면 좋겠다고 의견을 모았습니다. 누구나 쉽게 가져다 사용할 수 있는 GPipe 제작 프로젝트에 착수하게 된 배경이죠. 파이토치를 주로 쓰는 AutoML 연구팀을 위해 텐서플로가 아닌 파이토치용 라이브러리로 제작했습니다. 지난 2월부터 대략 4개월간 진행된 오픈소스 프로젝트 torchgpipe에 관한 정보는 여기서 확인해보실 수 있습니다.

[ 그림 8 ] 이흥섭 연구원은 “이 그림은 벨기에의 초현실주의 화가 르네 마그리트의 대표작인 ‘이미지의 배반’에서 차용한 torchgpipe 프로젝트의 로고“라고 설명했다.

이흥섭 연구원과 정명룡 연구원은 파이토치 코드와 조화를 이루면서도 GPipe의 핵심인 파이프라인 병렬화와 체크포인팅을 효과적으로 구현할 방법을 집요하게 고민했습니다. 다양한 시도 끝에 CUDA 프로파일링 도구[8](NVIDIA Nsight Systems)를 통해 GPU가 쉴 새 없이 모델을 빠르게 처리하는 이상적인 타임라인을 확인해볼 수 있었습니다. 이는 GPU가 본연의 능력대로 열심히 일하고 있음을 뜻합니다.

이 연구원은 “미니배치에 해당하는 입력 데이터를 여러 마이크로배치로 분할하고, 마이크로배치에 해당하는 여러 출력 데이터를 다시 미니배치로 모아서 출력하면 됐기에 파이프라인 병렬화 구현은 어렵지 않았다”고 설명합니다.

반면 체크포인팅은 구현이 쉽지 않았습니다. 역전파 중 예기치 않은 동작이 발견됐는데 자체 코드를 분석하는 것만으로는 어디에서 문제가 발생했는지 알 수 없었죠. 이에 두 사람은 파이토치의 자동미분(autograd) 엔진 코드 분석을 통해 반직관적인 동작을 파악했습니다. 파이프라인 병렬화와 체크포인팅을 동시에 적용하면 특정 마이크로배치를 역전파할 때 다른 마이크로배치도 함께 처리됐던 것입니다.

좀 더 쉽게 설명하자면 이렇습니다. A-B-C라는 마이크로배치가 있다고 가정해봅시다. 모델 역전파시 자동미분 엔진에 “C를 역전파해, B를 역전파해, A를 역전파해’’라는 명령어가 순서대로 실행되길 기대했죠. 그런데 실제로는 첫번째 명령어(C를 역전파해)가 실행될 때 B와 A도 함께 역전파되는 현상이 발생했습니다.

두 사람은 파이토치 커뮤니티와 심층적으로 논의한 끝에 마이크로배치 간의 관계를 정의하는 것으로 문제를 해결할 방법을 찾았습니다. 위 예시에 적용해보면 C를 처리하는 동안 B와 A는 처리되지 못하도록, 마찬가지로 B가 처리되는 동안 A가 처리될 수 없도록 [그림 9]처럼 계산 그래프를 정의했습니다. 그 결과, 체크포인트가 정상적으로 작동하지 않는 문제를 해결할 수 있었습니다.

[ 그림 9 ] 마이크로배치간 의존 관계를 설정하면 역전파 때 체크포인트 간 실행 순서를 강제할 수 있다. © 이흥섭 연구원

한편, 파이프라인 병렬화를 분석할 때 GPU의 대기시간이 뒤쪽으로 갈수록 늘어나는 현상이 발견됐습니다. 앞쪽 GPU가 처리한 순전파 결과값을 뒤쪽 GPU의 입력값으로 동기화하는 작업은 두 GPU의 연산 작업이 끝나고 나서야 이뤄집니다. 만약 한 GPU가 일을 먼저 끝낸다면 다른 GPU가 일을 끝낼 때까지 대기하는데요, 이 대기 시간은 수십~수백 ms 정도에 불과할 정도로 대단히 짧습니다. 다만 이런 흐름이 쌓이면 모델 뒤쪽을 처리하는 GPU가 쉬는 시간이 늘어나게 되죠. 많은 GPU를 사용하는 파이프라인 병렬화에서는 큰 손해가 아닐 수 없습니다.

[ 그림 10 ] GPU간 복사 대기에 락스텝 개념을 적용하기 전과 후의 GPU 타임라인을 비교했다. © 이흥섭 연구원

“이 문제를 가지고 몇 날 며칠 씨름했지만 이렇다 할 성과를 내지 못했습니다. 그러던 와중에 파이토치 1.1 버전과 함께 공개된 파이프라인 병렬화 튜토리얼 문서를 찬찬히 살펴보게 되었죠. 문서를 읽고 나니 멀티스레딩을 활용한 파이프라인 병렬화 구현이 최선이 아닐 수도 있다는 생각이 들었습니다.”

두 사람은 다른 GPU의 연산 작업에 영향을 주지 않으면서도 각각 빠르게 작동하게 만들고자 GPU마다 스레드를 할당하는 멀티스레딩 방식으로 파이프라인 병렬화를 구현했습니다. 그런데 튜토리얼 문서에서는 한 스레드로 여러 GPU를 처리하는 싱글스레딩 방식으로도 파이프라인 병렬화를 충분히 구현할 수 있다고 설명합니다. 어차피 CPU는 순차적으로 명령만 내리고 GPU는 이 명령에 따라 비동기적으로(asynchronous) 작업을 처리할 테니 말이죠. 하지만 학습 조건마다 결과가 다를 것이라고 본 두 사람은 어느 방식이 파이프라인 병렬화 구현에 더 효과적인지 알아보고자 비교 실험을 진행했습니다.

그 결과, 마이크로배치의 크기가 지나치게 작은 상황일 때 싱글스레딩 방식의 순차적 명령어 실행 덕분에 학습 속도가 저하될지라도, GPU간 값을 복사할 때 뒤쪽 GPU의 연산 작업이 지체되는 현상이 발생하지 않았습니다. 두 사람은 이 패턴을 보고 락스텝(lockstep) 개념을 떠올렸습니다. 스타크래프트(Starcraft)를 비롯한 멀티플레이어 게임에 흔히 사용되는 기법인 락스텝은 매 프레임마다 모든 플레이어의 게임 상태를 동기화하는 방법을 가리킵니다. 클라이언트 상에서 게임 상태가 맞춰질 때까지 진행을 멈추고(lock), 맞춰지면 게임을 진행하는(step) 거죠.

GPU 간 복사 대기에 락스텝 개념을 적용하면 다음과 같습니다. 마이크로배치를 활용한 파이프라인 병렬화에서는 빨리 일을 끝낸 GPU가 다른 GPU가 일을 마칠 때까지 대기합니다. 학습속도에 영향을 미치는 건 제일 느린 GPU뿐이니 이렇게 한다고 해서 속도가 느려지지도 않고요. 오히려 제일 느린 GPU의 복사 대기 시간이 줄어들어 학습속도는 빨라집니다.

[ 그림 11 ] 락스텝 개념을 적용하면 GPU 간 복사에 불필요한 대기시간이 발생하지 않는다. © 이흥섭 연구원

“앞쪽 GPU가 연산하는 도중에 결과값을 복사하라는 명령어가 실행되면 복사가 즉시 처리될 수 없어서 뒤쪽 GPU는 출발지 GPU가 일을 마칠 때까지 기다립니다. 그 결과 뒤쪽 GPU에서 다음 연산 작업이 약간 지연됩니다. 멀티스레딩 상황에서는 순차적으로 진행되어야 할 명령어의 순서가 엉켜서 이런 현상이 높은 빈도로 발생했습니다. 하지만 싱글스레딩에서는 두 GPU가 연산하지 않을 때 복사를 수행하는 명령어를 입력하는 게 보장됩니다. 이 덕분에 GPU 사이 입력값-결과값을 동기화 하는 데 낭비되는 시간이 사라졌죠.”

이 차이를 비교한 끝에 두 연구원은 멀티스레딩 방식에도 락스텝을 적용하면 멀티스레딩의 장점(마이크로배치가 작을 때도 학습 속도가 저하되지 않는다)과 싱글스레딩의 장점(복사 동기화에 따른 GPU 대기 시간이 발생하지 않는다) 모두를 취할 수 있다는 결론을 내렸습니다. 그 결과, [도표 2]에서처럼 마이크로 배치 크기와 관계없이 GPU 개수에 비례해 학습 처리 속도가 늘어나는 결과를 확인해볼 수 있었습니다. 이 결과는 깃허브에서 확인해볼 수 있습니다.

이흥섭 연구원과 정명룡 연구원은 분류 모델을 대상으로 torchgpipe의 효율성 검증을 완료한 상태로, 그밖에 다양한 모델에도 적용해서 효율성을 확인할 계획입니다. 마지막으로 이 연구원은 “GPipe에 적용된 기법을 잘 모르는 상황에서도 torchgpipe를 이용해 기존의 제약을 뛰어넘는 모델 연구가 활발해지면 좋겠다”고 말했습니다.


참고
[1] 파이썬(Python) 기반 딥러닝 프레임워크
[2] 마이크로소프트 리서치(Microsoft Research)가 개발한 ResNet은 범용적으로 쓰이는 분류 모델이다.
[3] 이 문서에서는 모델 훈련을 처리하는 장치(GPU 또는 TPU)를 가리키는 단어를 GPU로 통일했다.
[4] 구글이 AutoML을 통해 찾은, 최고 성능을 내는 모델 중 하나
[5] AmoebaNet-X (L, F)에서 X는 모델 구조를 뜻한과. L과 F는 각각 층의 깊이와 필터 개수에 비례하는 속성을 갖춘다.
[6] 모델에 데이터를 입력해 여러 개의 은닉층을 거쳐서 출력값을 내는 방법
[7] 목적값과 실제 출력값의 차이인 오차를 출력층-은닉층-입력층으로 거슬러 올라가며 가중치를 업데이트하는 방법
[8] GPU가 언제 어떻게 일하는지를 측정해 이를 타임라임으로 시각화해주는 도구
이 글을 쓴 사람들
samantha.lee
이수경 | 글,정리
지난 2016년 3월 알파고와 이세돌 9단이 펼치는 세기의 대결을 취재한 것을 계기로 인공지능 세계에 큰 매력을 느꼈습니다. 인공지능을 알고 싶어하는 사람들을 위한 콘텐츠를 쓰고자 카카오브레인에 합류했습니다. 현재는 분야별 인공지능 전문가와 함께 기술 콘텐츠를 생산하는 재미에 푹 빠져 있습니다. 인공지능을 만드는 사람들의 이야기와 인공지능이 바꿀 미래 사회에 대한 글은 누구보다 쉽고, 재미있게 쓰는 사람이 되고 싶습니다.
hans.lee
이흥섭 | 기술감수
즐겨 쓰는 서비스에 인공지능이 접목된 것을 보며 우리 삶 속에 인공지능이 깊숙이 들어왔음을 알게 됐습니다. 이 미래 기술에 더 알고 싶다는 지적 욕구를 충족하고자 게임개발자로서의 삶을 뒤로하고 카카오브레인에 합류했습니다. 현재는 동료들이 실험 결과를 지루하게 기다리지 않도록 인공지능 모델의 학습 시간을 획기적으로 줄이는 데 집중하고 있습니다. 새롭게 발견한 값진 아이디어를 인공지능 커뮤니티와 아낌없이 나누며 그 생태계 발전에 크게 기여하고 싶습니다.