https://arxiv.org/abs/2402.16828
Training Neural Networks from Scratch with Parallel Low-Rank Adapters 라는 내용의 논문을 발견해 소개드립니다. 제목에 써 놓은 대로 트랜스포머 기반 모델을 처음부터 lora 만 사용해서 학습시키는 기법입니다.
기존에는 전체 파라메터를 학습시키고 나서 나중에 파인튜닝할 때 메모리 요구량을 줄이고 학습속도를 높이기 위해서 lora 를 사용하는 것이 일반적이었습니다. 풀 파라메터 학습은 상당히 방대한 량의 컴퓨팅 자원과 시간이 필요하기 때문에 극소수의 대기업수준에서만 수행할 수 있었고, 소규모의 연구자들은 pretrained 모델을 받아다가 한정된 자원으로 파인튜닝하기 위해서 lora 기법을 쓰는 식으로 영역이 나뉘어져 있었습니다.
lora 는 트랜스포머 모델이 구조적으로 일정한 데이터 (주어진 전체 데이타셋중 이번 batch 에 주어진 데이타) 에 대해서 그것을 담당하는 (주어진 데이타의 패턴을 인식하고, 그 결과에 대해 In context learning 을 수행하는 등의 작업) 파라메터는 전체 파라메터중 아주 일부만이 실제로 관련되어 있다라는 관찰에 기반한 것입니다. 그러한 특징은 여러가지로 응용될 수 있는데, 파인튜닝하는 데이타가 전체 모델의 데이타중 아주 일부에 속한다면 모델의 전체 공간중 부분공간만을 떼어서 학습시키는 식으로 응용할 수도 있고 (일반적인 lora), 최근에 발표된 GaLore 처럼 파인튜닝의 프로세스에서 역전파되는 파라메터들의 그레디언트값은 전체 그레디언트 공간의 부분공간으로 충분히 근사해서 표현할 수 있다라는 식으로 응용될 수도 있습니다. 특히 GaLore 와 FSDP 의 조합을 이용해서 4090 * 2 GPU 구성으로 70B 모델의 파인튜닝이 가능해진다라는 돌파구가 마련되기도 하였습니다.
하지만 파인튜닝만으로는 아쉽습니다. AI 모델을 처음부터 만들 수 있는 능력이 더 많은 사람들의 손에 주어질 수는 없는 것일까요? Lora 에서 경험했던 ai 민주화의 가능성이 프리트레이닝의 범주로까지 확장될 수는 없는걸까요?
얼마전에는 nvidia 의 2024 GTC 키노트가 있었습니다. 키노트를 보신 분들은 느끼셨겠지만 블랙웰 GPU 는 이제 더 이상 그래픽 칩이 아니라 엄청난 규모의 플랫폼이고, 우리가 이걸 구축할테니 당신들은 돈을 내고 허락받고 쓰시면 된다라는, 업계천하통일 선언과도 같은 일성이었습니다. 어차피 GPU 한장으로는 할 수 있는 일이 제한적이니 nvlink 기술을 극한으로 끌어올려서 수만개의 B100 칩을 서로 연결하고 그것을 지탱할 수 있는 엄청난 인프라로 AI 공장을 만들어서 보급하겠다는 모습이 너무나도 생생하게 표현되었습니다. 칩 하나가 전체 데이터센터로 규모가 scale-up 되는 모습입니다.
저에게 이 시점에서 떠오르는 것은 1990년대 중반에 인터넷과 www 가 태동하던 시절의 모습입니다. 솔라리스 서버와 오라클 db 같은 고가의 장비에 의해서 컴퓨터가 전 세계에 정보를 전파하고, 국제전화비가 엄청나게 비싸던 시절에 공짜로 외국에 메일을 보낼 수 있던 충격은 대단했습니다. 많은 사람들이 닷컴이란 것에 관심을 갖기 시작했고, 어떻게 하면 우리 회사도 인터넷 서버라는 것을 만들 수 있는걸까? 라고 관심을 갖던 모습은 30년이 지난 지금 AI 열풍에 비견할만 했었습니다.
인터넷 서비스를 만들기 위해서는 비싼 서버와 비싼 소프트웨어가 당연히 필요할 것이라고 생각하던 시점에 모두의 예상을 깬 것은 구글이었습니다. 그들의 검색엔진은 엄청난 성능을 자랑했는데, 반전이었던 것은 비싼 컴퓨터와 비싼 DB 가 아니라 그냥 보통 리눅스와 mysql db 로 구성된 인스턴스들을 수평적으로 많이 분산시켜서 만들었다는 점이었습니다. MS 나 오라클 같은 회사들은 서버 시장을 어떻게든 장악하고, 리눅스를 비롯한 오픈소스를 밟아죽여버리고 싶었지만 결국 최종 승자는 리눅스와 오픈소스 소프트웨어 들이 되었습니다.
대규모 언어모델을 프리트레이닝시키는 작업을 엄청나게 비싼, 수직적으로 확장된 고가의 머신이 아니라, 수평적으로 확장된, 오픈소스와 일용품화된 소비자용 머신들이 서로 느슨하게 연결되어 수행할 수 있게 된다면, 그것은 nvidia 나 openai 가 가지고 싶어하는 가장 강력한 MOAT (해자, 방어막) 을 부술 수 있는 열쇠가 될 것입니다. 그리고 지금 소개할 LTE 는 개별적인 머신이 각각 lora 를 학습하고 가끔씩 합쳐주는 식으로 llm 을 프리트레이닝을 가능하게 해줄 수 있습니다.
그동안 lora 를 이용해서 분산 학습을 하려는 시도들이 없었던 것은 아닙니다. 다만 그런 시도들은 어디까지나 lora 로 파인튜닝 학습하는 것을 목표로 했기 때문에 이미 사전학습된 모델을 함께 갖고 시작하는 것을 전제로 했습니다.
분산은 일단 둘째치고, 사전학습부터 lora 로만 학습하는 것이 과연 가능한 것인지부터 생각해봅시다.
lora 의 크기는 rank r로 표현됩니다. 여기서의 rank 는 서로 독립인 기저 벡터들의 갯수를 의미합니다. 텐서의 랭크와는 다른 개념입니다.
보통 사전학습의 대상이 될 웨이트 W 는 
같은 식으로 m*n 크기의 행렬로 나타냅니다. 로라의 경우 r 는 m 이나 n 중 작은값보다 더 작게 잡아야 의미가 있습니다.
정확하진 않지만 대충 비슷하게 표현하자면 이런 느낌입니다
위의 Wv 나 Wo 는 4*16 짜리 행렬의 형태를 띄고 있습니다. 이게 전체 파라메터 공간의 크기라고 한다면 칠해져 있는 4*4 짜리 네모는 랭크4 짜리 부분공간에 해당합니다.
로라 행렬은 정사각행렬이고, 그 축의 크기인 r 은 위 파라메터 크기인 m*n 의 m 이나 n 보다 작거나 같아야 합니다. 즉 4*16 짜리 파라메터 공간에서 lora 가 차지할 수 있는 공간은 4*4 가 한계입니다. 이 경우 4*4 는 full rank Lora 가 됩니다.
4*4 와 4*16 의 크기 차이 때문에 lora 는 파인튜닝 용도로만 활용되었던 것이고 사전학습용으로는 고려의 대상이 아니었습니다. 이것을 극복하려면 어떻게 해야 할까요?
멀티 헤드 로라라는 것을 사용합니다. 아래 수식을 보면 알 수 있는데, 전체 가중치 행렬 W 를 여러개의 AB (로라 행렬) 의 가중치합으로 표현한 것입니다. 즉 위쪽의 16*4 짜리 전체 행렬을 4*4 짜리 행렬 4 개를 이어붙인 식으로 표현했다고 보면 됩니다.
이런 방식을 사용해서 튜닝을 하면 풀파라메터 튜닝의 성능을 따라갈 수 있을까요? 아래 그래프와 같이 풀랭크 lora 는 멀티 헤드 로라를 쓰지 않고도 이미 튜닝 능력은 충분합니다
수렴속도가 조금 느리긴 하지만 로라의 랭크수를 최대한 늘린 상태에서는 성능을 따라갈 수 있습니다.
이제 남은 것은 학습입니다. 여러대의 머신을 동일한 세팅으로 초기화하고, 각 머신이 전체데이타셋을 분할해서 나눠가서 (pull) 각자 학습을 한 다음에 그걸 commit 하고 메인 브랜치에서 병합하는 과정을 반복하는 것으로 프리트레이닝이 가능합니다!
기존의 분산 학습과 가장 큰 차이점은 머신들간에 서로 동기화해야 할 정보량이 훨씬 적다는 것입니다. 전체 가중치값을 동기화하는 것이 아니라 로라 행렬만큼만 메인 머신으로 보내주면 됩니다 (몇기가 vs 몇백메가 수준) 그리고 그 동기화의 간격도 하나의 micro batch step 마다 할 필요 없이 더 많이 묵혔다가 동기화해도 됩니다.
동기화는 위 오른쪽의 수식에서 보는 것처럼 각 로라의 값들을 더해서 평균을 내 주는 것으로 충분합니다.
이후 논문에서 소개하는 내용은 가중치의 초기화값을 어떻게 할 것이냐라던가 학습률의 조정 같은 세부사항을 논하고 있습니다. 이 부분은 관심있는 분들은 직접 보시면 될 것같고, 제가 이 논문을 소개하면서 가장 강조하고 싶었던 것은, 지금 30년전에 인터넷 혁명이 일어났던 것처럼 ai 혁명이 일어나고 있는 시점에서 기술적 인프라적 분기점이 우리 앞에 놓여있다는 점입니다.
전기와 발전 기술같은 것은 처음에 기술이 알려지면서 각자 발전기를 돌리는 식으로 활용이 되다가 나중에는 규모의 경제를 이룰 수 있도록 국영기업이나 대규모 사업자들 중심으로 산업이 만들어지고 각 가정이나 공장은 필요한만큼 한전 같은 회사와 계약을 하고 받아오는 식으로 발전을 하게 되었습니다. 반면 인터넷 같은 것은 전기와는 또다른 느낌의 공공재로 발전하면서 지금도 망중립성 문제, 비용문제 등으로 팽팽한 줄다리기가 이어지고 있습니다만, 적어도 한명의 해커가 나만의 서비스를 만들어서 세상에 선보이는 데에는 비용이 거의 들지 않는 방향으로 가고 있습니다.
AI 의 경우에는 어떻게 될 것이라고 생각하시나요? 전 적어도 이번 논문이나 galore 같은 기술이 역사의 축을 약간씩 바꾸고 있다고 생각합니다.
Ai 언어모델 로컬 채널 이용규정
