LLM Trend Note (6) LLM.int8(), LoRA, Prefix LM, Sparse transfomer, Sparse attention, Model parallelism, Data parallelism

지난 노트에서 보았던 그림을 다시 한번 보도록하겠습니다.

어떠신가요?
전체 흐름이 처음보다 더 눈에 들어오시나요?

OpenAI, Google, Deepmind, Meta, HuggingFace, EleutherAI, SalesForces, BigScience
그리고 그림엔 빠져 있는 Anthropic까지
Foundation Model 이라고 명명되는 LLM을 향한 유수 테크기업들의 결과물들을 빠르게 살펴보면서
여러분들은 어떤 생각이 드셨는지 궁금합니다.

Emergent Abilities라는 다소 거창해 보이는 개념을 짚어보면서
single GPU로 학습하는 상황이 대부분일 사람들에게
LLM 연구라는 것이 어쩌면 다소 멀게만 느껴질 수 있을거라고도 말씀드린 바 있었습니다.

이번 노트를 작성하면서 발견한 트윗 하나도 그런 심정을 잘 대변하고 있는 것 같습니다.

현재 시점에서 가장 강력한 Emergence를 보여주는 LLM은
instruction tuning된 LLM과 RLHF가 적용된 LLM 입니다.
각각 2번 스텝의 FLAN과 5번 4번 스텝의 instructGPT를 통해 살펴본 내용들이었습니다.

이러한 모델들의 단점은 오직 "크기"입니다.
FLAN은 13B 짜리 LaMDA를 사용했고 StackLLaMA의 경우 6.7B짜리 LLaMA를 사용했었죠.
4번 스텝에서 봤던 것처럼 만약 reference model을 initial model의 사본 그대로 사용한다면
단일 GPU에서 PPO를 돌릴 시에 최소 2개 이상의 모델이 로드가 되어야 합니다.
모델의 크기가 커지면 단일 GPU에 원하는 크기의 배치로 학습을 최적화시키는 일이 점점 어려워질 수밖에 없습니다. 참고

StackLLaMA의 경우 위 문제를
reference model과 active모델 간에 레이어를 공유하고 헤드만 학습시키는 방법으로 해소했습니다.

구체적으로,
StackLLaMA에서는 저렴한 computing cost로 LLM을 fine-tuning 할 수 있는 최신 학습기법을 사용했습니다.
모델 가중치를 8bit로 로드하고,
A Gentle Introduction to 8-bit Matrix Multiplication for transformers at scale using Hugging Face Transformers, Accelerate and bitsandbytes
논문 : LLM.int8(): 8-bit Matrix Multiplication for Transformers at Scale

PLM의 가중치를 고정시키는 대신
어텐션 레이어에 훈련 가능한 rank decomposition 행렬을 주입해
연산비용을 획기적으로 낮춰주는 LoRA(Low-Rank Adaptation of Large Language Models) 와 같은 기법들입니다.
(참고 링크 1)
(참고 링크 2)

허깅페이스에서는 위 기법들을 망라한 PEFT library를 제공하고 있습니다.
PEFT: Parameter-Efficient Fine-Tuning of Billion-Scale Models on Low-Resource Hardware
(깃헙 참고 링크)

위와 같은 기법들의 메커니즘을 이해하고 보다 경제적이며 효과적인 기법들로 발전시키는 연구는
LLM을 다루기 위한 수백대의 GPU가 달린 서버를 필요로 하지 않습니다.
Amazon의 Sagemaker로 176B짜리 BLOOMZ를
single GPU에서 fine-tuning하는 방법이 궁금하시다면 이 링크를 참고해보세요.

또한 모델 아키텍쳐 자체에 대한 보다 심도 깊은 연구의 필요성도 계속해서 요구될 것입니다.
앞서 Instruction tuning을 설명하면서 살펴본 FLAN에 쓰인 LaMDA는 디코더 기반 트랜스포머 모델입니다.
단, 어텐션 기법은 T5에 사용된 아키텍쳐를 그대로 사용했습니다.
T5의 어텐션 레이어에서 주목해볼만한 특징은 Prefix LM 구조입니다.

Vanilla Transformer에서는 마스킹과 관련해 두 종류의 어텐션만 사용합니다.
fully-visible attention과 causal attention이지요. 이에 더해 T5의 어텐션은 새로운 어텐션 매커니즘이 추가됩니다.
입력 시퀀스의 앞쪽 몇개의 토큰만 fully-visible attention을 수행하고
나머지 뒤쪽 토큰들에는 causal attention을 수행할 수 있는 아키텍쳐 이지요.
instruction tuning이 어떻게 이뤄지는지 생각해보면
FLAN 연구진이 왜 T5 아키텍쳐에 기반한 모델인 LaMDA를 사용했는지 추론해 볼 수 있습니다.

좀 더 깊숙이 들어가면 어텐션 연산 자체에 대한 연구도 필요합니다.

소프트웨어 측면에서 셀프 어텐션 연산은 일반적으로 O(n^2) 의 시간복잡도를 요구하지만
이를 O(log n) 로 단축시키는 새로운 어텐션 메커니즘이 있습니다.
구글에서 발표한 Self-attention Does Not Need O(n^2) Memory논문을 참고하세요.
어텐션 연산에 소요되는 메모리를 절약할 수 있다면,
여유분의 메모리를 다른 곳에 더 사용할 수 있는 새로운 어텐션 아키텍쳐의 도입이 가능해집니다.
예를 들어, 모든 토큰에 대해 어텐션을 수행하는 게 아니라 특정 토큰에 선택적으로 어텐션을 하는 식으로 말이죠.
이러한 기법을 통칭해 sparse transformer(또는 sparse attention) 이라고 합니다.

이와 관련한 자세한 내용이 잘 요약되어 있는 논문 두 편을 소개해 드립니다.
ㄱ. Efficient Transformers: A Survey
ㄴ. Generating Long Sequences with Sparse Transformers

위 Efficient Transformers에 나온 sparse transfomer 아키텍쳐의 구현은
facebookresearch에서 공개한 라이브러리 xformers에서 찾아볼 수 있습니다.

이전 스텝에서 언급한 최신 LLM인 LLaMA의 경우에도
모델의 학습 속도를 개선하기 위해 위 기법들을 모델에 적용했다고 합니다.
(LLaMA 논문 2.4 Efficient implementation 참고하세요.
(LLaMA의 깃헙 공개 레포짓에는 attention구현에 xformer가 사용되어 있지 않습니다.
비공개 모델에 사용되었단 뜻일까요? 여러분이 직접 구현해 볼 수도 있을 것 같습니다:)

어텐션 연산의 하드웨어 측면에서의 개선도 있었습니다.
LLaMA 연구진들은 참고했다고 밝힌 스탠포드에서 발표한
FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness논문은
어텐션 연산이 이뤄지는 물리적 공간인 GPU의 HBM(high bandwidth memory)와
SRAM 간의 메모리 read/writes 수를 줄일 수 있는 새로운 어텐션 메커니즘입니다.
flash attention은 xformer에 third party라이브러리로 공개되어 있습니다.

결과적으로 LLaMA는 위와 같은 기법을 활용해 모든 어텐션 가중치를 저장하지 않고,
causal language modeling 시 마스킹 되는 부분에 대한
key/query score를 계산하지 않는 멀티헤드어텐션을 사용함으로써
메모리 공간을 절약하고 어텐션 속도도 향상시켰습니다.

여기서 끝이 아닙니다.
만약 multiple GPU 사용이 가능한 환경이라면
우리들에게 필요한 것은 LLM을 보다 효율적으로 분산학습시킬 수 있는 분산학습 알고리즘이 될 것입니다.
Chain of thought Prompting 기법을 처음으로 소개했던 PaLM 모델이 기억나시나요?
input으로 엄청난 길이의 prompt를 입력시키고 이를 모델이 이해하게 만들기 위해선
모델 아키텍쳐의 응용뿐만 아니라,
그렇게 거대해진 모델과 데이터들을 parallel할 수 있는 학습기법이 반드시 필요합니다.

위 그림에서 4x4 그리드는 16개의 코어를 나타냅니다.
음영처리된 사각형들은 해당 코어에 포함된 데이터(토큰들로 이뤄진 배치) 또는 모델 가중치 입니다.

첫번째 행의 다섯개 그림은 모델 가중치가 코어간에 분할되는 방식을 보여줍니다.
예컨대 Data parallelism은 서로 다른 코어에 모델 복제본들이 분할되어 있음을 보여줍니다.
Model parallelism은 코어 간에 하나의 모델이 걸쳐 있음을 보여줍니다.
Model and Data parallelism은 위 두 방식이 결합되어 있음을 보여주고
Expert and Data parallelism은 MoE 방식을 마지막 그림은 앞선 네가지 방식을 합친 parallelism을 도식화하고 있습니다.

두번째 행의 그림은 데이터 배치가 코어 간에 분할되는 방식을 보여주는 그림입니다.
따라서 여기서 Data parallelism과 Model parallelism은 첫번째 행에 상응하는 그림이 역전되어 있죠.

어떤 전략을 선택하느냐에 따라 각 코어는 동일한 배치를 가질수도 있고 서로 다른 배치를 가질수도 있습니다.
이에 따라 data pipeline을 어떻게 짜야 할지,
gradient update를 위한 코어간 통신전략은 어떻게 해야 할지,
각 상황에 맞는 코딩을 해줘야 합니다.
이전 스텝에서 살펴본 StackLLaMA역시 분산학습 전략을 제공하고 있습니다.

PPO_trainer.py 스트립트와 torch.distributed 패키지의 functional.py 스크립트에서
all_reduce 메서드를 찾아보세요.
all_reduce는 어떤 알고리즘인가요?

지금까지 살펴본 내용들의 공통점이 보이시나요?

foundation model의 emergent abilities를 시작으로
model scale과 data size 그리고 RLHF로 이어졌던 이야기 끝에서 소개해 드린 한 엔지니어의 다소 의기소침한 트윗으로
이번 스텝을 시작했습니다. 마지막 스텝 내용을 보시고 난 후 여러분은 그 트윗에 어떤 댓글을 달고 싶어 지셨나요?

우리는 첫번째 노트에서
low resource language로 분류되는 한국어로 언어모델을 만드는 데 따르는 어려움을
LLM에 사용가능한 토큰 총량의 문제와 결부하여 생각해보았습니다.

LLM Trend Note2에서는
한국어 데이터를 사용해 ChatGPT의 RLHF를 간접적으로 구현한 모델을 직접 학습시켜봄으로써
foundation model과 emergent abilities 문제를
한국어 언어모델과 한국어 데이터의 관점으로 접근해보는 시간을 가져보도록 하겠습니다.