본 콘텐츠의 이미지 및 내용은 AI로 생성되었습니다.
본 콘텐츠의 이미지 및 내용을 무단으로 복제, 배포, 수정하여 사용할 경우 저작권법에 의해 법적 제재를 받을 수 있습니다.
이미지 로딩 중...
AI Generated
2025. 11. 28. · 14 Views
GPT Transformer 아키텍처 분석
OpenAI의 GPT 모델이 어떻게 구현되어 있는지 gpt.py 코드를 직접 분석합니다. Attention 메커니즘부터 1.9B 파라미터의 출처까지, 초급 개발자도 이해할 수 있도록 차근차근 살펴봅니다.
목차
- gpt.py 전체 구조 파악
- CausalSelfAttention 구현 분석
- MLP와 LayerNorm
- depth 파라미터의 의미
- 1.9B 파라미터는 어디서 오는가
- 메모리 최적화 기법들
1. gpt.py 전체 구조 파악
김개발 씨는 요즘 GPT에 푹 빠져 있습니다. ChatGPT를 쓸 때마다 "이게 도대체 어떻게 동작하는 거지?"라는 궁금증이 떠나지 않았습니다.
마침내 용기를 내어 Andrej Karpathy의 nanoGPT 코드를 열어보았습니다.
GPT의 전체 구조는 놀랍도록 단순합니다. 토큰 임베딩, 위치 임베딩, 여러 개의 Transformer Block, 그리고 최종 출력을 위한 언어 모델 헤드로 구성됩니다.
마치 레고 블록을 쌓듯이, 동일한 구조의 블록을 반복해서 쌓아올린 것이 GPT의 핵심입니다.
다음 코드를 살펴봅시다.
class GPT(nn.Module):
def __init__(self, config):
super().__init__()
# 토큰을 벡터로 변환하는 임베딩 테이블
self.transformer = nn.ModuleDict(dict(
wte = nn.Embedding(config.vocab_size, config.n_embd),
wpe = nn.Embedding(config.block_size, config.n_embd),
drop = nn.Dropout(config.dropout),
# n_layer개의 동일한 Block을 쌓습니다
h = nn.ModuleList([Block(config) for _ in range(config.n_layer)]),
ln_f = nn.LayerNorm(config.n_embd),
))
# 최종 출력을 어휘 크기로 변환
self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)
김개발 씨는 처음 gpt.py 파일을 열었을 때 당황했습니다. 수백 줄의 코드가 눈앞에 펼쳐졌기 때문입니다.
하지만 선배 박시니어 씨가 옆에서 조언해주었습니다. "전체 그림부터 보세요.
나무를 보기 전에 숲을 먼저 봐야 합니다." GPT의 전체 구조는 의외로 단순합니다. 마치 고층 빌딩을 짓는 것과 비슷합니다.
1층에 로비가 있고, 중간에 똑같은 구조의 사무실 층이 반복되며, 꼭대기에 전망대가 있는 것처럼요. GPT에서 로비에 해당하는 것이 바로 임베딩 레이어입니다.
여기서 텍스트가 처음으로 숫자의 세계로 들어옵니다. "안녕하세요"라는 텍스트는 먼저 토큰으로 쪼개지고, 각 토큰은 고차원의 벡터로 변환됩니다.
wte는 Word Token Embedding의 약자입니다. 50,257개의 어휘 각각에 대해 768차원(또는 그 이상)의 벡터를 할당합니다.
이 벡터는 학습을 통해 단어의 의미를 담게 됩니다. wpe는 Word Position Embedding입니다.
"나는 밥을 먹는다"와 "밥을 나는 먹는다"는 같은 단어를 사용하지만 의미가 다릅니다. 위치 정보가 없으면 모델은 이 차이를 알 수 없습니다.
wpe는 각 위치에 고유한 벡터를 부여하여 순서 정보를 제공합니다. 중간의 반복되는 사무실 층이 바로 Transformer Block입니다.
코드에서 h라는 이름으로 정의되어 있습니다. n_layer 개수만큼 동일한 Block이 쌓입니다.
GPT-2 Small은 12개, GPT-2 Large는 36개, GPT-3는 무려 96개의 블록을 쌓습니다. 각 Block은 두 가지 핵심 구성요소를 가집니다.
Self-Attention과 MLP입니다. 이 두 가지가 번갈아가며 데이터를 처리합니다.
마치 회사에서 회의(Attention)와 개인 작업(MLP)을 반복하는 것과 같습니다. 꼭대기 전망대에 해당하는 것이 lm_head입니다.
Language Model Head의 약자로, 마지막 히든 스테이트를 어휘 크기의 벡터로 변환합니다. 이 벡터에 소프트맥스를 적용하면 다음 토큰의 확률 분포가 됩니다.
흥미로운 점은 wte와 lm_head가 가중치를 공유한다는 것입니다. 이를 weight tying이라고 합니다.
입력과 출력에서 같은 임베딩을 사용하면 파라미터 수를 줄이면서도 성능을 유지할 수 있습니다. 김개발 씨는 이제 전체 구조가 눈에 들어왔습니다.
"생각보다 단순하네요. 같은 블록을 반복해서 쌓는 거였군요!" 박시니어 씨가 고개를 끄덕였습니다.
"맞아요. 이제 각 블록 안을 들여다볼 차례입니다."
실전 팁
💡 - GPT 구조를 이해할 때는 항상 데이터의 shape 변화를 추적하세요
- ModuleDict와 ModuleList는 PyTorch에서 여러 레이어를 깔끔하게 관리하는 방법입니다
2. CausalSelfAttention 구현 분석
박시니어 씨가 김개발 씨에게 질문했습니다. "GPT의 핵심이 뭔지 알아요?" 김개발 씨는 잠시 생각하다가 대답했습니다.
"Attention... 아닌가요?" 박시니어 씨가 미소를 지었습니다.
"정확해요. 그럼 Attention이 실제로 어떻게 구현되어 있는지 볼까요?"
CausalSelfAttention은 GPT의 심장입니다. 각 토큰이 이전 토큰들만 참조할 수 있도록 마스킹된 어텐션을 구현합니다.
Query, Key, Value 세 가지 투영을 통해 토큰 간의 관계를 계산하고, 이를 Multi-Head로 병렬 처리합니다.
다음 코드를 살펴봅시다.
class CausalSelfAttention(nn.Module):
def __init__(self, config):
super().__init__()
# Q, K, V를 한 번에 계산하는 통합 projection
self.c_attn = nn.Linear(config.n_embd, 3 * config.n_embd)
self.c_proj = nn.Linear(config.n_embd, config.n_embd)
self.n_head = config.n_head
self.n_embd = config.n_embd
# 미래를 보지 못하게 하는 삼각형 마스크
self.register_buffer("bias", torch.tril(
torch.ones(config.block_size, config.block_size)
).view(1, 1, config.block_size, config.block_size))
def forward(self, x):
B, T, C = x.size() # 배치, 시퀀스길이, 임베딩차원
# Q, K, V 계산 후 분리
q, k, v = self.c_attn(x).split(self.n_embd, dim=2)
김개발 씨는 CausalSelfAttention 클래스를 열어보았습니다. 이름부터 어려워 보입니다.
하나씩 뜯어보기로 했습니다. 먼저 Causal이라는 단어의 의미부터 알아봅시다.
Causal은 "인과적"이라는 뜻입니다. 원인이 결과보다 먼저 와야 한다는 것이죠.
GPT에서는 이전 단어들만 보고 다음 단어를 예측해야 합니다. 미래의 단어를 미리 보면 안 됩니다.
이것을 비유하자면 시험을 보는 상황과 같습니다. 1번 문제를 풀 때 2번, 3번 문제의 답을 미리 볼 수 없어야 공정한 시험이 됩니다.
GPT도 마찬가지입니다. "나는 사과를 ___"에서 빈칸을 예측할 때, 뒤에 올 단어를 미리 보면 안 됩니다.
이를 구현하는 것이 바로 삼각형 마스크입니다. torch.tril은 lower triangular matrix를 만듭니다.
대각선 아래만 1이고 위는 0인 행렬입니다. 이 마스크를 어텐션 스코어에 적용하면, 각 위치에서 자신보다 뒤의 토큰은 볼 수 없게 됩니다.
다음으로 Self-Attention의 핵심인 Q, K, V를 살펴봅시다. Query, Key, Value의 약자입니다.
도서관에 비유하면 이해하기 쉽습니다. 여러분이 도서관에서 "파이썬 입문서"를 찾는다고 합시다.
여러분의 요청이 Query입니다. 각 책의 제목과 분류가 Key입니다.
실제 책의 내용이 Value입니다. Query와 Key를 비교해서 관련성이 높은 책을 찾고, 해당 책의 Value를 가져오는 것이 어텐션의 원리입니다.
코드에서 c_attn은 입력을 Q, K, V 세 가지로 변환합니다. 효율성을 위해 3 * n_embd 크기로 한 번에 계산한 후 split으로 나눕니다.
별도의 Linear 레이어 세 개를 쓰는 것보다 빠릅니다. Multi-Head라는 개념도 중요합니다.
헤드가 여러 개라는 뜻입니다. 왜 여러 개가 필요할까요?
한 문장에서도 여러 종류의 관계가 있습니다. "철수가 영희에게 사과를 주었다"에서 "주었다"는 "철수"와 주어 관계, "영희"와 간접목적어 관계, "사과"와 직접목적어 관계를 가집니다.
헤드 하나로는 이 모든 관계를 포착하기 어렵습니다. 여러 헤드가 각각 다른 종류의 관계를 학습합니다.
GPT-2 Small은 12개의 헤드를 사용합니다. 768차원의 임베딩을 12개로 나누면 각 헤드는 64차원을 담당합니다.
각 헤드가 계산한 결과를 다시 합쳐서 최종 출력을 만듭니다. c_proj는 어텐션 결과를 원래 차원으로 투영합니다.
여러 헤드의 출력을 통합하고, 다음 레이어로 전달할 준비를 합니다. 김개발 씨가 고개를 끄덕였습니다.
"도서관 비유가 정말 와닿네요. Query로 검색하고, Key로 매칭하고, Value를 가져오는 거군요." 박시니어 씨가 덧붙였습니다.
"맞아요. 그리고 Causal 마스킹으로 미래를 못 보게 막는 거죠."
실전 팁
💡 - register_buffer는 학습되지 않지만 모델과 함께 저장되어야 하는 텐서에 사용합니다
- Flash Attention을 사용하면 메모리 효율성이 크게 개선됩니다
3. MLP와 LayerNorm
김개발 씨가 Attention을 이해하고 나니 자신감이 붙었습니다. "Attention이 전부인 줄 알았는데, MLP도 있네요?" 박시니어 씨가 설명을 시작했습니다.
"Attention이 정보를 모으는 역할이라면, MLP는 그 정보를 처리하는 역할이에요."
**MLP(Multi-Layer Perceptron)**는 Attention이 수집한 정보를 비선형 변환합니다. 4배로 확장했다가 다시 원래 크기로 축소하는 구조입니다.
LayerNorm은 각 레이어의 출력을 정규화하여 학습을 안정시킵니다. 이 둘이 함께 작동하여 깊은 네트워크의 학습을 가능하게 합니다.
다음 코드를 살펴봅시다.
class MLP(nn.Module):
def __init__(self, config):
super().__init__()
# 4배로 확장 (병목 구조)
self.c_fc = nn.Linear(config.n_embd, 4 * config.n_embd)
self.gelu = nn.GELU()
# 다시 원래 크기로 축소
self.c_proj = nn.Linear(4 * config.n_embd, config.n_embd)
self.dropout = nn.Dropout(config.dropout)
def forward(self, x):
x = self.c_fc(x) # 768 -> 3072
x = self.gelu(x) # 비선형 활성화
x = self.c_proj(x) # 3072 -> 768
return self.dropout(x)
Transformer Block 안에는 Attention만 있는 것이 아닙니다. MLP도 중요한 역할을 합니다.
둘의 관계를 이해하면 GPT의 작동 원리가 더 명확해집니다. Attention을 회의에 비유했습니다.
여러 사람의 의견을 듣고 종합하는 과정입니다. 그런데 회의만 한다고 일이 되지는 않습니다.
회의에서 나온 정보를 바탕으로 실제 작업을 해야 합니다. 그것이 MLP의 역할입니다.
MLP는 Feed-Forward Network라고도 불립니다. 구조는 단순합니다.
두 개의 Linear 레이어 사이에 활성화 함수가 있습니다. 하지만 여기에 재미있는 설계가 숨어 있습니다.
첫 번째 Linear 레이어 c_fc는 차원을 4배로 늘립니다. 768차원이 3072차원이 됩니다.
왜 굳이 늘릴까요? 더 넓은 공간에서 복잡한 패턴을 표현하기 위해서입니다.
좁은 방에서는 큰 가구를 배치하기 어렵지만, 넓은 방에서는 자유롭게 배치할 수 있는 것과 같습니다. GELU 활성화 함수가 그 사이에 위치합니다.
ReLU의 부드러운 버전이라고 생각하면 됩니다. ReLU는 0 이하를 무조건 0으로 만들지만, GELU는 확률적으로 부드럽게 처리합니다.
GPT 논문에서 GELU가 더 좋은 성능을 보여서 선택되었습니다. 두 번째 Linear 레이어 c_proj는 다시 원래 차원으로 줄입니다.
3072에서 768로 돌아갑니다. 이런 구조를 병목(bottleneck) 구조라고 합니다.
정보를 압축하는 과정에서 핵심만 남기는 효과가 있습니다. 이제 LayerNorm을 살펴봅시다.
왜 정규화가 필요할까요? 깊은 신경망에서는 레이어를 통과할 때마다 값의 분포가 변합니다.
이를 Internal Covariate Shift라고 합니다. 값이 너무 커지거나 작아지면 학습이 불안정해집니다.
LayerNorm은 각 레이어의 출력을 평균 0, 분산 1로 정규화하여 이 문제를 해결합니다. GPT에서는 Pre-LN 구조를 사용합니다.
LayerNorm을 Attention과 MLP 앞에 배치합니다. 원래 Transformer 논문에서는 Post-LN(뒤에 배치)을 사용했지만, Pre-LN이 학습 안정성이 더 좋다는 것이 밝혀졌습니다.
Block의 전체 흐름을 정리하면 이렇습니다. 입력이 들어오면 먼저 LayerNorm을 거칩니다.
그 다음 Attention을 통과하고, 원래 입력과 더해집니다. 이것이 **잔차 연결(Residual Connection)**입니다.
다시 LayerNorm, MLP를 거치고, 또 잔차 연결을 합니다. 잔차 연결은 왜 필요할까요?
네트워크가 깊어지면 기울기가 사라지는 문제가 발생합니다. 잔차 연결은 기울기가 직접 흐를 수 있는 고속도로를 만들어줍니다.
덕분에 96개 레이어를 쌓아도 학습이 가능합니다. 김개발 씨가 정리했습니다.
"Attention으로 정보를 모으고, MLP로 처리하고, LayerNorm으로 안정시키고, 잔차 연결로 기울기를 전달하는 거네요." 완벽한 요약이었습니다.
실전 팁
💡 - Pre-LN vs Post-LN의 차이를 이해하면 다양한 Transformer 변형을 쉽게 파악할 수 있습니다
- MLP의 4배 확장 비율은 관례적인 것이며, 최근에는 다른 비율을 실험하기도 합니다
4. depth 파라미터의 의미
김개발 씨가 GPT 모델을 학습시켜보려고 config를 살펴보았습니다. n_layer, n_head, n_embd...
여러 파라미터가 있었습니다. "이것들을 어떻게 설정해야 하죠?" 박시니어 씨가 대답했습니다.
"depth와 width의 균형이 중요해요."
GPT의 depth는 Transformer Block의 개수(n_layer)를 의미합니다. width는 임베딩 차원(n_embd)과 헤드 개수(n_head)로 결정됩니다.
같은 파라미터 수라도 depth와 width의 조합에 따라 성능이 달라집니다. 일반적으로 더 깊은 모델이 복잡한 패턴을 학습하는 데 유리합니다.
다음 코드를 살펴봅시다.
# GPT 모델 크기별 설정
GPT_CONFIGS = {
'gpt2': dict(n_layer=12, n_head=12, n_embd=768), # 124M
'gpt2-medium': dict(n_layer=24, n_head=16, n_embd=1024), # 350M
'gpt2-large': dict(n_layer=36, n_head=20, n_embd=1280), # 774M
'gpt2-xl': dict(n_layer=48, n_head=25, n_embd=1600), # 1.5B
}
# depth가 깊어질수록
# - 더 추상적인 특징을 학습
# - 더 긴 의존성을 포착
# - 더 많은 연산량 필요
# n_embd는 항상 n_head로 나누어 떨어져야 함
assert config.n_embd % config.n_head == 0
모델의 크기를 결정하는 핵심 파라미터들이 있습니다. 이것들의 의미와 관계를 이해하면 GPT의 설계 철학이 보입니다.
먼저 n_layer입니다. Transformer Block이 몇 개 쌓이는지를 결정합니다.
이것이 바로 모델의 **depth(깊이)**입니다. GPT-2 Small은 12개, GPT-2 XL은 48개를 사용합니다.
깊이가 중요한 이유는 무엇일까요? 각 레이어는 이전 레이어의 출력을 입력으로 받아 더 추상적인 특징을 추출합니다.
얕은 레이어에서는 단어의 기본적인 의미를 파악합니다. 깊은 레이어로 갈수록 문맥, 뉘앙스, 논리적 관계 등 고수준 특징을 학습합니다.
이미지 인식에 비유하면, 얕은 레이어는 가장자리와 색상을 감지합니다. 중간 레이어는 눈, 코, 입 같은 부분을 인식합니다.
깊은 레이어는 "이것은 고양이다"라는 추상적 개념을 파악합니다. 언어 모델도 비슷한 계층 구조를 가집니다.
다음은 n_embd입니다. 각 토큰을 표현하는 벡터의 차원입니다.
이것이 모델의 **width(너비)**를 결정합니다. GPT-2 Small은 768, GPT-2 XL은 1600을 사용합니다.
차원이 크면 각 토큰에 더 많은 정보를 담을 수 있습니다. 하지만 무작정 크게 하면 메모리와 연산량이 급격히 늘어납니다.
적절한 균형이 필요합니다. n_head는 Attention의 헤드 개수입니다.
앞서 설명했듯이 여러 종류의 관계를 동시에 학습하기 위해 여러 헤드를 사용합니다. 중요한 제약이 있는데, n_embd는 반드시 n_head로 나누어 떨어져야 합니다.
768을 12로 나누면 64, 각 헤드가 64차원을 담당합니다. block_size는 최대 시퀀스 길이입니다.
GPT-2는 1024 토큰까지 처리할 수 있습니다. GPT-4는 32K 또는 128K까지 늘어났습니다.
긴 컨텍스트를 처리하려면 더 큰 block_size가 필요하지만, 메모리 사용량은 시퀀스 길이의 제곱에 비례해서 늘어납니다. vocab_size는 어휘 크기입니다.
GPT-2는 50,257개의 토큰을 사용합니다. BPE(Byte Pair Encoding) 알고리즘으로 만들어진 서브워드 토큰들입니다.
depth와 width 중 무엇이 더 중요할까요? 연구에 따르면 같은 파라미터 수일 때 적당히 깊은 모델이 더 좋은 성능을 보입니다.
너무 얕으면 복잡한 패턴을 학습하지 못하고, 너무 깊으면 학습이 어려워집니다. Scaling Law라는 개념도 있습니다.
모델 크기, 데이터 양, 연산량의 관계를 수식으로 표현한 것입니다. OpenAI의 연구에 따르면, 이 세 가지를 균형 있게 늘려야 성능이 효율적으로 향상됩니다.
김개발 씨가 질문했습니다. "그러면 제가 학습시킬 때는 어떤 설정을 써야 하나요?" 박시니어 씨가 답했습니다.
"GPU 메모리에 맞춰서 시작하세요. 작은 모델로 실험하고, 잘 되면 크기를 키우면 됩니다."
실전 팁
💡 - 처음에는 gpt2 설정으로 시작해서 동작을 확인하고, 점진적으로 크기를 키우세요
- n_embd % n_head == 0 제약을 항상 확인하세요
5. 1.9B 파라미터는 어디서 오는가
"GPT-2 XL이 15억 파라미터라고 하는데, 이게 어디서 나오는 숫자예요?" 김개발 씨의 질문에 박시니어 씨가 종이를 꺼내 계산을 시작했습니다. "직접 계산해보면 이해가 빨라요."
GPT의 파라미터 수는 임베딩, Attention, MLP, LayerNorm의 파라미터를 모두 합한 것입니다. 대부분의 파라미터는 Attention의 QKV 투영과 MLP의 확장/축소 레이어에 있습니다.
1.9B 같은 큰 수가 어떻게 나오는지 직접 계산해보면 모델 구조에 대한 이해가 깊어집니다.
다음 코드를 살펴봅시다.
def count_parameters(config):
n_embd = config.n_embd # 예: 1600
n_layer = config.n_layer # 예: 48
vocab_size = config.vocab_size # 예: 50257
block_size = config.block_size # 예: 1024
# 토큰 임베딩: vocab_size * n_embd
wte = vocab_size * n_embd # 50257 * 1600 = 80.4M
# 위치 임베딩: block_size * n_embd
wpe = block_size * n_embd # 1024 * 1600 = 1.6M
# Attention (QKV + proj): 4 * n_embd^2 per layer
attn = 4 * n_embd * n_embd * n_layer # 4 * 1600^2 * 48 = 491M
# MLP (확장 + 축소): 8 * n_embd^2 per layer
mlp = 8 * n_embd * n_embd * n_layer # 8 * 1600^2 * 48 = 983M
# 총합 (LayerNorm, bias 등 제외한 근사치)
total = wte + wpe + attn + mlp # 약 1.56B
GPT-3가 175B 파라미터라는 말을 들으면 막연히 "크다"고만 느껴집니다. 하지만 직접 계산해보면 그 숫자가 어디서 오는지 명확해집니다.
파라미터가 있는 곳을 하나씩 살펴봅시다. 첫 번째는 **토큰 임베딩(wte)**입니다.
어휘의 각 토큰을 n_embd 차원의 벡터로 매핑합니다. GPT-2 XL 기준으로 50,257 * 1,600 = 약 8천만 개입니다.
생각보다 많습니다. 두 번째는 **위치 임베딩(wpe)**입니다.
각 위치에 n_embd 차원의 벡터를 할당합니다. 1,024 * 1,600 = 약 160만 개입니다.
전체에서 차지하는 비중은 작습니다. 세 번째는 Attention 레이어입니다.
각 블록마다 c_attn과 c_proj가 있습니다. c_attn은 n_embd를 3 * n_embd로 변환합니다.
n_embd * 3 * n_embd = 3 * n_embd^2입니다. c_proj는 n_embd를 n_embd로 변환합니다.
n_embd^2입니다. 합치면 4 * n_embd^2 per layer입니다.
네 번째는 MLP 레이어입니다. c_fc는 n_embd를 4 * n_embd로 변환합니다.
4 * n_embd^2입니다. c_proj는 4 * n_embd를 n_embd로 변환합니다.
역시 4 * n_embd^2입니다. 합치면 8 * n_embd^2 per layer입니다.
여기서 중요한 패턴이 보입니다. n_embd의 제곱에 비례한다는 것입니다.
n_embd를 2배로 늘리면 파라미터는 4배가 됩니다. 그래서 width를 늘리는 것이 depth를 늘리는 것보다 파라미터 증가가 급격합니다.
GPT-2 XL로 계산해봅시다. n_layer=48, n_embd=1600입니다.
임베딩: 50,257 * 1,600 + 1,024 * 1,600 = 약 82M Attention: 4 * 1,600^2 * 48 = 약 491M MLP: 8 * 1,600^2 * 48 = 약 983M 합계는 약 1.56B입니다. 공식 발표인 1.5B와 비슷합니다.
LayerNorm과 bias까지 포함하면 더 정확해집니다. MLP가 Attention의 2배라는 점이 눈에 띕니다.
4배 확장 때문입니다. 실제로 GPT에서 가장 많은 파라미터를 차지하는 것은 Attention이 아니라 MLP입니다.
최근에는 MoE(Mixture of Experts) 방식으로 MLP 파라미터를 효율적으로 사용하는 연구가 많습니다. 여러 개의 MLP 중 일부만 선택적으로 활성화하는 방식입니다.
파라미터는 많지만 실제 연산량은 줄일 수 있습니다. 김개발 씨가 감탄했습니다.
"직접 계산해보니까 어디에 파라미터가 많은지 확실히 알겠어요." 박시니어 씨가 덧붙였습니다. "모델을 최적화할 때 이 분포를 알면 어디를 건드려야 할지 감이 옵니다."
실전 팁
💡 - numel() 함수로 실제 파라미터 수를 확인할 수 있습니다: sum(p.numel() for p in model.parameters())
- weight tying을 사용하면 wte와 lm_head의 파라미터를 공유하여 메모리를 절약합니다
6. 메모리 최적화 기법들
김개발 씨가 실제로 GPT를 학습시켜보려고 했습니다. 하지만 바로 문제에 부딪혔습니다.
"CUDA out of memory 오류가 계속 나요!" 박시니어 씨가 웃으며 말했습니다. "대형 모델 학습의 첫 번째 관문이에요.
메모리 최적화 기법들을 알아봅시다."
대형 언어 모델 학습에서 메모리는 가장 큰 병목입니다. Gradient Checkpointing은 메모리를 아끼기 위해 중간 활성화를 저장하지 않고 역전파 시 다시 계산합니다.
Mixed Precision Training은 FP16을 사용하여 메모리를 절반으로 줄입니다. Flash Attention은 어텐션 계산의 메모리 복잡도를 O(N^2)에서 O(N)으로 줄입니다.
다음 코드를 살펴봅시다.
# Mixed Precision Training with autocast
from torch.cuda.amp import autocast, GradScaler
scaler = GradScaler()
for batch in dataloader:
optimizer.zero_grad()
# FP16으로 forward pass 수행
with autocast():
logits = model(batch['input_ids'])
loss = F.cross_entropy(logits.view(-1, vocab_size),
batch['labels'].view(-1))
# 스케일링된 backward pass
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()
# Flash Attention 사용 (PyTorch 2.0+)
# torch.nn.functional.scaled_dot_product_attention 자동 사용
GPU 메모리는 유한합니다. RTX 3090이 24GB, A100이 80GB입니다.
하지만 GPT-3 수준의 모델은 수백 GB의 메모리가 필요합니다. 어떻게 학습시킬 수 있을까요?
첫 번째 기법은 Gradient Checkpointing입니다. 체크포인팅이라고도 합니다.
일반적인 학습에서는 forward pass 중에 모든 중간 활성화를 저장합니다. backward pass에서 기울기를 계산할 때 필요하기 때문입니다.
문제는 레이어가 깊을수록 저장할 활성화가 많아진다는 것입니다. Gradient Checkpointing은 발상을 전환합니다.
중간 활성화를 저장하지 않습니다. 대신 backward pass에서 필요할 때 forward pass를 다시 계산합니다.
메모리는 아끼지만 연산량이 늘어납니다. 보통 연산량이 30% 정도 증가합니다.
두 번째 기법은 Mixed Precision Training입니다. 혼합 정밀도 학습이라고도 합니다.
기본적으로 신경망은 FP32(32비트 부동소수점)를 사용합니다. 하지만 많은 연산에서 FP16(16비트)으로도 충분합니다.
FP16을 사용하면 메모리 사용량이 절반으로 줄고, 텐서 코어를 활용하면 연산 속도도 빨라집니다. 문제는 FP16의 표현 범위가 좁다는 것입니다.
기울기가 너무 작으면 언더플로우가 발생합니다. 이를 해결하기 위해 Loss Scaling을 사용합니다.
Loss를 크게 스케일업한 후 기울기를 계산하고, 가중치 업데이트 전에 다시 스케일다운합니다. PyTorch의 GradScaler가 이 과정을 자동으로 처리합니다.
autocast 컨텍스트 안에서는 자동으로 FP16 연산이 적용됩니다. 세 번째 기법은 Flash Attention입니다.
2022년에 발표된 혁신적인 알고리즘입니다. 표준 Attention은 Q, K를 곱해서 N*N 크기의 어텐션 행렬을 만듭니다.
시퀀스 길이 N이 길어지면 메모리가 N^2에 비례해서 늘어납니다. 8K 시퀀스면 6천4백만 개의 원소를 저장해야 합니다.
Flash Attention은 이 행렬을 명시적으로 만들지 않습니다. 대신 타일 단위로 쪼개서 계산하고, GPU의 빠른 SRAM을 활용합니다.
메모리 복잡도가 O(N)으로 줄어들고, 속도도 2-4배 빨라집니다. PyTorch 2.0부터는 scaled_dot_product_attention 함수가 Flash Attention을 자동으로 사용합니다.
별도의 라이브러리 설치 없이 바로 활용할 수 있습니다. 네 번째로 Gradient Accumulation도 중요합니다.
배치 크기가 클수록 학습이 안정적이지만, 메모리가 부족하면 큰 배치를 사용할 수 없습니다. Gradient Accumulation은 작은 배치로 여러 번 forward/backward를 수행하고, 기울기를 누적한 후 한 번에 업데이트합니다.
효과적으로 큰 배치와 같은 효과를 얻습니다. 마지막으로 DeepSpeed와 FSDP 같은 분산 학습 프레임워크가 있습니다.
모델을 여러 GPU에 나눠서 학습합니다. ZeRO 최적화는 옵티마이저 상태, 기울기, 파라미터를 GPU 간에 분산하여 메모리 사용량을 대폭 줄입니다.
김개발 씨가 하나씩 적용해보았습니다. Mixed Precision만으로도 배치 크기를 2배로 늘릴 수 있었습니다.
"이제 학습이 되네요!" 박시니어 씨가 뿌듯하게 웃었습니다. "메모리 최적화는 대형 모델 학습의 필수 기술이에요."
실전 팁
💡 - torch.cuda.memory_summary()로 메모리 사용 현황을 상세히 확인할 수 있습니다
- 학습 시작 전에 torch.cuda.empty_cache()로 캐시를 비우세요
- bfloat16은 FP16보다 수치 안정성이 좋아서 최신 GPU에서 권장됩니다
이상으로 학습을 마칩니다. 위 내용을 직접 코드로 작성해보면서 익혀보세요!
댓글 (0)
함께 보면 좋은 카드 뉴스
Helm 마이크로서비스 패키징 완벽 가이드
Kubernetes 환경에서 마이크로서비스를 효율적으로 패키징하고 배포하는 Helm의 핵심 기능을 실무 중심으로 학습합니다. Chart 생성부터 릴리스 관리까지 체계적으로 다룹니다.
보안 아키텍처 구성 완벽 가이드
프로젝트의 보안을 처음부터 설계하는 방법을 배웁니다. AWS 환경에서 VPC부터 WAF, 암호화, 접근 제어까지 실무에서 바로 적용할 수 있는 보안 아키텍처를 단계별로 구성해봅니다.
AWS Organizations 완벽 가이드
여러 AWS 계정을 체계적으로 관리하고 통합 결제와 보안 정책을 적용하는 방법을 실무 스토리로 쉽게 배워봅니다. 초보 개발자도 바로 이해할 수 있는 친절한 설명과 실전 예제를 제공합니다.
AWS KMS 암호화 완벽 가이드
AWS KMS(Key Management Service)를 활용한 클라우드 데이터 암호화 방법을 초급 개발자를 위해 쉽게 설명합니다. CMK 생성부터 S3, EBS 암호화, 봉투 암호화까지 실무에 필요한 모든 내용을 담았습니다.
AWS Secrets Manager 완벽 가이드
AWS에서 데이터베이스 비밀번호, API 키 등 민감한 정보를 안전하게 관리하는 Secrets Manager의 핵심 개념과 실무 활용법을 배워봅니다. 초급 개발자도 쉽게 따라할 수 있도록 실전 예제와 함께 설명합니다.