NLLB-200 distilled 350M 두근두근 개발일지

배경

졸업 전 캡스톤 디자인 프로젝트에서 생성형AI를 사용한 서비스를 만들고자 하였다. 교수님도 생성형AI가 들어갔으면 하셨고, 무엇보다도 내가 원했다! 최종적으로 생각한 그림은 크롬 익스텐션으로 영문 이미지를 캡쳐하는 즉시 이미지가 한국어로 번역된 이미지로 보여지는 그런 형태의 제품을 고안했다.

 

초기엔 LLM을 QLora를 통해 학습하는 방식을 택했으나, 요즘 나오는 LLM들이 1B은 훌쩍 넘어가는데, 인퍼런스 타임이 최대한 빠릿빠릿해야 사용함에 불편함이 없을 것이다. 그리고 GPU서버를 쓰는데에 비용이 적잖이 들어간다. AI서비스가 GPU서버 없이 운영된다는건 좀 모순이긴 하지만... 뭐 어쩌겠는가... 돈없는 대학생인데... ㅠㅠ 😭😭😭 결국 LLM 포기

 

추가적인 리서치를 통해 Meta에서 공개한 NLLB-200 모델에 대해 알게 되었다. 공교롭게도 우리의 목적과 딱 맞아 떨어졌다! distill된 600M 모델도 제공하고, 번역성능도 상당히 좋다! 바로 사용해도 괜찮을 듯 싶었지만, 아주 영끌을 해서 더더욱 작고 빠르게 만들어 보면 어떨까? 가정용 GPU로도 서빙이 가능한 사이즈로 말이다!! 교수님께서도 경량화에 대한 가능성을 말씀해 주셨고, 무엇보다도 Knowledge Distillation은 내가 한번쯤은 꼭 해보고 싶은 분야이기도 했다!

 

개발

우선 가장 큰 어려움은 학습에 필요한 GPU를 어떻게 충당하느냐는 점이었다. 우리는 매주 발표때마다 학습에 사용될 GPU리소스 수급에 큰 낭항을 겪고 있다는 점을 감정을 들여 호소하였고, 결국 마침내 GPU를 지원받을 수 잇는 방법을 알려주셨다. 학교 산악협력단에서 매달 사용하게될 Colab Pro+비용정도는 지원해 준다는 것이다!! 야호!

교수님 저 가격좀 보세요 ㅠㅠ 저희 돈 없는데...

 

 

그럼에도 자원이 한정적이다. 충분한 실험과 테스트를 해볼 수 없는 상황이다. 모델의 사이즈를 어떻게 줄일 것인가, 하이퍼파라미터는 어떻게 세팅할 것인가, teacher model은 어떤 모델을 써야할까 등등 우리는 매순간 선택을 해야했다. 

 

우선 모델은 600M 모델을 베이스로 dimension은 1024로 유지하되, Layer의 수를 줄이는 방향으로 선택했다. 보통 논문에서도 dimension보다 Layer의 수를 줄이는 방식을 택하는 방식을 취하는 까닭이다. encoder layers 와 decoder layers를 각각 12에서 3으로 줄여서 350M로 만들었다.  존경하는 monologg님도 DistilKoBert를 그렇게 만드셨음!

 

다음으로는 학습에 사용될 데이터를 구해야한다. LLM을 fine tuning 할때에는 AI허브서 제공하는 영한 병렬 데이터를 사용했으나, Distillation은 Teacher모델을 학습할 때와 동일한 데이터를 사용하는 것이 일반적이다. 논문에서는 Wiki 데이터를 Laser3 인코더를 통해 점수가 높은 데이터만 골라서 학습하는 방식을 사용했다. 대용량의 위키 말뭉치를 다 받아서 일일이 인코더에 태워보는 것은 부담이다. 참 까마득한 일이다... 아아 이를 어찌하면 좋을까... 하던 중 위 작업을 누군가가 해서 Huggingface에 올려 두셨다.  당신도 내 존경하는 인물 목록에 추가

 

가장 난감한 골치덩이 2개가 깨끗이 해결된 느낌이다. 그 다음엔 바로 학습에 들어갔다. 논문에서는 Teacher모델로는 NLLB-200 1.3B모델을 사용했다. 54B나 3.3B를 사용할 법도 한데, 200종류의 언어를 다 보는것은 Meta여도 부담인가 보다. 우리 또한 1.3B을 Teacher로 하여 Online distillation을 진행했다. 목적에 맞게 영-한 번역만을 취급했고, 하나 윗단계 모델(600M)의 Score는 32점! 32점에 비비는 모델을 만들고야 말겠노라, 컴피티션 한다는 마음가짐으로 임했다.

두근두근 얼마나 할까

 

약속받은 지원금을 다 털어서 colab pay as you go 를 결제했고 A100 GPU로 학습을 했다. 데이터의 크기가 약 2천만건 쯤 되어서, 100만건 단위로 쪼개어 학습하고 저장하고를 반복했다. colab은 중간에 맘대로 런타임이 꺼져서 꼭 중간 저장을 해주어야 한다. huggingface에 Private으로 사이사이에 push를 해주었다. 중간에 다른 데이터로 학습도 해보았지만, chrF 점수가 오르지 않았다. 아마 데이터 퀄리티가 떨어지는 탓일 것이다. 최종적으로 1500만건의 데이터를 학습했고 결과는 아래와 같다.

 

결과

으으음...

 

최대한 chrF점수를 올려보려고 했지만 24.6에서 더이상 오르지 않았다. 좀더 다양한 테스트를 해보면 좋았겠지만 자원이 충분치 않은 상황이다. 모델을 다르게 줄여 봤다면 어땠을까, 데이터를 좀더 점검해봤다면 달랐을까 등등 오만가지 생각이 다 들지만... Computing Unit이 남아있지 않아서 추가적인 진행을 해볼수 없다. 😭😭😭

 

무엇보다도 600M 녀석이 사이즈에 비해 성능이 지나치게 좋다. 어쩜 1.3B모델과 점수를 비빈단 말인가... 교수님께서 조언해 주시기를 Meta측에서도 300 400 500 700 등 다양한 사이즈로 실험을 해보고 그중에 가장 괜찮은 사이즈가 600M이기에 600M으로 만들어서 냈을 가능성이 높다 하셨다. 

If Nothing last forever, would you be my 'Nothing'? 이런건 잘 못함 ㅠ

 

그래도 악조건 속에서도 그럭저럭 모델을 만들어 냈다. 무엇보다도 사이즈가 1.4GB라서 나의 반려 GPU 1650s 4GB로도 서빙이 가능하다! 초기 목적은 어느정도 달성한 셈이다. 24.6점이긴 하지만 formal한 영문장 정도는 문제 없이 번역한다. 사용할만한 가치가 있을 것 같다. 역시 뭐든 첫시도에 완벽하게 되는건 없는것 같다. 이렇게 점점 경험치를 먹어가며 성장해 가야지!

 

 

https://github.com/newfull5/NLLB-200-Distilled-350M-en-ko

GitHub - newfull5/NLLB-200-Distilled-350M-en-ko: nllb-200 distilled 350M for English to Korean translation

 

https://huggingface.co/dhtocks/nllb-200-distilled-350M_en-ko

dhtocks/nllb-200-distilled-350M_en-ko · Hugging Face