제로부터 시작한 Llama, 울지 않고 논문을 구현하는 방법

 (blog.briankitano.com)
2P by GN⁺ | ★ favorite | 댓글 1개
  • Brian Kitano는 TinyShakespeare로 축소판 Llama를 직접 만들며, 논문 구현은 작은 모델에서 출발해 부품을 하나씩 교체하고 매번 학습·평가해야 안전하다고 정리함
  • 데이터 분할, 배치 생성, 손실 평가, 생성 함수 같은 검증 보조 함수를 먼저 갖추고, 단순 모델로 컴파일과 학습 가능 여부를 확인한 뒤 Llama 구성요소를 추가함
  • RMSNorm, RoPE, SwiGLU를 순서대로 붙이면서 텐서 shape, 수식 성질, 어텐션 맵으로 각 레이어가 기대대로 동작하는지 확인함
  • RoPE 어텐션에서 causal mask를 빼면 검증 손실은 0.16까지 내려가지만 생성 품질은 나빠졌고, 미래 토큰을 보는 정보 누수가 원인이었음
  • 최종 축소 Llama는 4개 블록, 약 237만 파라미터로 검증 손실을 약 1.0까지 낮췄으며, gradient 흐름과 학습률 스케줄까지 함께 점검해야 함

작게 시작해 반복적으로 확신 쌓기

  • 논문 구현은 작은 모델에서 시작해 구성요소를 하나씩 바꾸고, 바꿀 때마다 학습과 평가를 반복하는 방식이 핵심임
  • 먼저 모델을 정량적으로 확인할 보조 함수들을 준비함
    • 데이터 분할
    • 학습 루프
    • 손실 시각화
    • 검증 손실 평가
  • 논문 구성요소를 모두 한 번에 옮기기보다, 이미 구현 경험이 있는 단순하고 빠른 모델로 생성 결과를 보는 정성 평가 함수도 마련함
  • 텐서 레이어는 .shape, assert, plt.imshow로 확인하고, 처음부터 행렬곱 최적화에 들어가기보다 손으로 기대 결과를 검산한 뒤 torch 함수로 효율화함
  • 배치 크기, 시퀀스 길이, 임베딩 차원을 바꿔가며 테스트해야 하며, 한 크기에서만 맞는 코드는 추론 시점에 깨질 수 있음

데이터셋과 기본 설정

  • 구현 대상은 Meta AI의 Llama를 크게 축소한 버전이며, 학습 데이터는 TinyShakespeare임
  • Llama는 1.4T 토큰으로 학습하지만, 여기서는 약 111만 문자 규모의 TinyShakespeare를 사용함
  • 원래 Llama는 SentencePiece byte-pair encoding 토크나이저를 쓰지만, 이 구현은 단순 문자 단위 토크나이저를 사용함
    • vocabulary size는 65
    • 데이터셋이 작아 메모리 저장 방식은 별도로 최적화하지 않음
  • MASTER_CONFIG 딕셔너리로 vocab_size, batch_size, context_window, d_model 같은 모델 설정을 관리함
    • 상수와 매직 넘버를 줄이고 코드를 읽기 쉽게 하기 위한 목적임
  • get_batches 함수는 데이터를 train 80%, val 10%, test 10%로 나누고, 임의 시작점에서 입력 x와 한 글자 뒤 라벨 y를 생성함

기본 모델로 컴파일과 학습 확인

  • 첫 모델은 임베딩과 간단한 feed-forward 네트워크로 구성된 SimpleBrokenModel임
    • nn.Embedding
    • Linear
    • ReLU
    • Linear
  • 논문 구현에서 모델이 “동작한다”는 말은 두 조건을 모두 만족해야 함
    • 컴파일: 텐서 shape이 레이어 사이에서 맞음
    • 학습: 손실이 실제로 내려감
  • evaluate_loss 함수는 train과 val split에서 10번 배치를 샘플링해 평균 손실을 계산함
  • SimpleBrokenModel은 1000 epochs 학습 후 검증 손실이 3.94 수준이었고, 초기 cross-entropy 4.17에서 거의 줄지 않음
  • 원인은 F.cross_entropy에 이미 softmax 처리된 값을 넣은 것이었음
    • PyTorch의 F.cross_entropy는 정규화되지 않은 logits를 직접 받음
    • softmax를 제거한 SimpleModel은 검증 손실을 2.51 수준까지 낮춤
  • 이후 generate 함수를 추가해 모델이 만든 문자를 직접 확인했고, 기본 모델은 완전하지 않지만 검증 손실이 내려가는 상태가 됨

Llama 구성요소 1: RMSNorm

  • Llama는 원래 Transformer와 비교해 세 가지 주요 아키텍처 수정을 사용함
    • RMSNorm pre-normalization
    • Rotary embeddings
    • SwiGLU activation function
  • 원래 Transformer는 BatchNormalization을 쓰지만, Llama는 벡터를 centering하지 않고 variance로 스케일링하는 RMSNorm을 사용함
  • 원래 Transformer가 attention layer 출력에 normalization을 적용하는 post-normalization 방식인 반면, Llama는 입력에 먼저 적용하는 pre-normalization 방식을 사용함
  • 구현한 RMSNorm은 입력 shape을 (batch, seq_len, d_model)로 가정함
  • RMSNorm 결과는 레이어 norm이 레이어 원소 수의 제곱근이 되는 성질로 테스트함
    • assert
    • row-wise comparison
    • torch.allclose
  • 기본 모델에 RMSNorm을 추가한 SimpleModel_RMS는 검증 손실을 2.5015 수준으로 소폭 낮춤

Llama 구성요소 2: RoPE와 causal mask

  • RoPE는 Transformer용 위치 인코딩 방식이며, 토큰 위치를 embedding 회전으로 표현함
  • get_rotary_matrix는 context window와 embedding dimension에 대해 위치별 회전 행렬을 생성함
  • RoPE 구현은 다음 성질로 테스트함
    • 위치 m, n에서 회전한 두 벡터의 내적이 상대 위치 n-m 회전과 일치해야 함
  • RoPEAttentionHead는 w_q, w_k, w_v를 만들고, query와 key에 RoPE 회전을 적용한 뒤 F.scaled_dot_product_attention을 사용함
  • 학습 시점과 추론 시점의 텐서 shape 차이를 주의해야 함
    • 학습 시에는 (config['batch_size'], config['context_window'], config['d_model'])처럼 설정과 맞는 경우가 많음
    • 추론 시에는 (1, 1, config['d_model']) 같은 단일 예시를 처리할 수 있음
    • forward 안에서는 모델 설정값이 아니라 입력에서 얻은 shape을 기준으로 인덱싱해야 함
  • causal mask 없이 RoPE multi-head attention을 추가한 모델은 검증 손실이 0.1623까지 급격히 낮아졌지만, 생성 결과는 OOOO..., IIII...처럼 좋지 않았음
  • attention map을 확인하자 모든 위치가 모든 위치를 참조하고 있었고, 다음 토큰 예측에서 미래 토큰을 보는 정보 누수가 발생함
  • F.scaled_dot_product_attention에 is_causal=True를 적용한 RoPEMaskedAttentionHead로 바꾸자 미래에 해당하는 upper triangular attention이 거의 0이 됨
  • causal mask 적용 후 검증 손실은 2.0815가 되었고, 더 오래 학습하면 1.8985까지 내려감

Llama 구성요소 3: SwiGLU와 블록 쌓기

  • Llama는 ReLU 비선형성을 SwiGLU activation function으로 교체함
  • 구현한 SwiGLU는 Swish-gated linear unit이며, 두 linear 변환과 learnable beta 파라미터를 사용함
  • feed-forward 부분에 SwiGLU를 넣은 RopeModel은 파라미터 수가 592,706개였고 검증 손실은 1.8963 수준임
  • 이후 LlamaBlock을 만들어 다음 구성을 하나의 블록으로 묶음
    • RMSNorm pre-normalization
    • masked RoPE multi-head attention
    • residual connection
    • RMSNorm pre-normalization
    • SwiGLU feed-forward
    • residual connection
  • 최종 Llama 모델은 n_layers=4로 설정하고, OrderedDict 기반 nn.Sequential로 4개 LlamaBlock을 쌓음
  • 최종 모델 파라미터 수는 2,370,246개이며, 학습 결과는 다음과 같음
    • 초기 4-layer 학습 후 검증 손실 1.5532
    • 10,000 epochs로 더 학습 후 검증 손실 1.1479
    • 추가 학습 후 검증 손실 0.9997
    • test split의 한 배치 손실은 1.2358

생성 결과와 디버깅 점검

  • 최종 모델은 Shakespeare 형식과 유사한 이름, 줄바꿈, 단어 조각을 만들지만 실제 문장 품질은 제한적임
  • cross-entropy 손실은 토큰 선택 관점에서 직관화할 수 있음
    • 초기 손실 4.17은 vocabulary size 65에서 무작위 선택에 가까움
    • 손실 1.08은 무작위로 약 2.9개 토큰 중 고르는 것과 같은 수준으로 해석함
  • gradient 흐름은 show_grads 함수로 확인함
    • 각 파라미터에서 절댓값이 작은 gradient 비율을 계산함
    • 대부분의 파라미터 gradient가 0에 가깝지 않으면 흐름이 괜찮은 상태임
  • 원래 Llama는 Cosine Annealing learning schedule을 쓰지만, 이 구현에서는 실험 결과가 더 나빴음
  • Cosine Annealing 실험에서는 매우 낮은 tolerance에서도 attention bias가 신호를 거의 받지 못했고, 이유는 확실하지 않아 실제 구현에서는 단순하게 시작하는 편이 안전함

함께 보면 좋은 글 β

GN⁺   [-]
Hacker News 의견들

  • SwiGLU 구현에 버그가 있어 보임: 참고 논문에서는 feed-forward network의 beta가 학습 가능한 값이 아니라 상수이고, FFnSwiGLU = Swish1...로 둠
    https://arxiv.org/pdf/2002.05202.pdf의 식 6 기준임
    공식 llama 구현에서도 상수 beta가 제거되어 있음: https://github.com/facebookresearch/llama/blob/main/llama/mo...
    블로그 로그의 "feedforward.1.beta', 0.0" 라인들을 보면 학습 중 beta가 0으로 퇴화했는데, 원래는 상수 1이어야 함

    • 트랜스포머 신경망을 정확히 구현하기가 얼마나 어려운지 보여줌. 여러 단계에서 실수할 수 있고, 보통은 “원래보다 약간 성능이 나쁜” 정도로만 드러나서 확실히 알기 어려움
      네트워크가 의도했든 아니든 변경에 적응하는 경우도 많고, 학습 후에는 여러 아키텍처 변형이 비슷하게 동작하기도 해서 원본과 꼭 일치해야 하는지 애매한 경우도 있음
      이런 실수를 찾는 방법으로는 참조 구현과 출력값을 정확히 맞추는 방식이 있음. HuggingFace의 tiny-random 모델처럼 무작위 가중치라도 출력이 정확히 같아야 하고, 다르면 버그 신호가 됨
      다만 이 방식은 추론 중 생기는 버그에만 잘 먹히고, 데이터 처리·최적화기·학습 중에만 발생하는 문제는 잡기 더 어려움
    • 트랜스포머에서 편향값은 대체로 잘 안 맞는 편이라고 봄
      개인적으로는 자기회귀적이고 ODE 같은 성질 때문이라고 생각하지만, 확신할 정도는 아님

  • 작업은 훌륭하지만, 초기 SimpleBrokenModel과 SimpleModel에는 꽤 많은 낭비 연산이 있음. embedding 65 -> 128, linear 128 -> 128, ReLU, linear 128 -> 65 순서인데, 앞의 두 층 사이에 비선형성이 없고 둘 다 선형이라 두 번째 선형층은 사실상 쓸모가 없음
    이 모델은 결국 고전적인 단일 은닉층 MLP와 같고, FLOPS 기준으로 128*128=16k 연산을 전체 128*128+65*128=24k 중 낭비하는 셈임

    • 비선형성을 아직 익히는 중인 사람이 나뿐만은 아닌 듯함. 여기서 가장 좋은 수정은 임베딩과 첫 선형층 사이에 ReLU나 SwiGLU를 넣는 것인지, 아니면 선형층을 그냥 삭제하는 것인지 궁금함
      임베딩층은 토큰 인덱스를 임베딩 벡터로 바꾸는 특수 구조라 제거할 수는 없을 것 같음

  • 전반적으로 기본 원칙을 잘 보여줌. 특히 “.shape를 종교적으로 써라. assert와 plt.imshow는 친구다”는 말이 좋고, shape의 사전·사후 조건은 항상 assert해야 함
    bear나 typeguard가 데코레이터로 이런 검사를 지원하는지도 궁금함
    다만 “작고 단순하고 빠른 모델을 고르고 정성적으로 평가하는 헬퍼를 만들라”는 부분은 정량 평가를 뜻하는 게 아닌가 싶음. 그래야 더 고급 기법과 비교할 수 있는 수치 기준선이 생김
    논문의 구성 요소를 하나씩 구현하라는 조언도 더 정확해야 함. 논문은 보통 여러 변경을 한꺼번에 시도한 뒤 제거 실험으로 각 요소의 기여를 보여주기 때문에, 핵심 아키텍처 변경부터 시작하고 제거 실험에서 영향이 큰 순서대로, 의존관계를 지키며 원자적 변경마다 평가하는 편이 더 낫다고 봄

    • bear나 typeguard 대신 https://peps.python.org/pep-0646/ 덕분에 일부는 Python 타입 주석으로 직접 밀어 넣을 수 있음
      예를 들면 ndarray[float, Dim1, *Shape] 같은 식으로 축별 shape를 타입에 표현하고, axis 값에 따라 반환 shape를 오버로드할 수 있음
    • PyTorch는 잘 모르겠지만 마지막으로 확인했을 때는 아니었고, Jax는 bear / typeguard를 통해 행렬 shape의 기초적인 런타임 검사를 지원함
      그래도 Python은 Julia만큼 좋기는 어려워 보임. Julia의 타입 시스템은 행렬 크기가 맞는지 훨씬 쉽게 보장할 수 있음

  • SwiGLU를 ReLU 대신 쓰는 원칙이 무엇인지 궁금함. 저자들이 가능한 비선형 함수를 그냥 다 시도해 본 건지, 더 깊은 이유가 있는지 모르겠음

    • 많은 연구가 그렇듯, 엄밀한 연구로 뒷받침된 명확한 설명이 없다면 아마도 멋져 보이는 한 줄짜리 변경들을 무작위로 언덕 오르기식 탐색하다가 논문 작성과 제거 실험을 시작할 시간이 됐을 때 멈췄을 가능성이 큼

  • bearblog가 DDoS를 맞는 중이라 저장소를 남김: https://github.com/bkitano/llama-from-scratch

  • AI를 배우는 입장에서 글에 나온 용어를 간단히 정리해 봄. 토큰은 텍스트 조각을 나타내는 정수 식별자이고, LLM에서는 제한된 어휘 크기 안에서 자주 쓰이는 문자 조각을 묶어 씀
    손실 함수는 예측과 정답의 차이를 재는 값이고 낮을수록 좋음. PyTorch는 텐서와 신경망을 다루는 라이브러리이며, 텐서는 스칼라·벡터·행렬을 포함하는 다차원 수 배열임
    신경망은 가중치와 편향을 가진 뉴런 연결 구조이고, 선형층은 모든 입력과 출력이 연결된 단순한 구조임. ReLU는 Math.max(0, x) 같은 활성화 함수로, 선형층만 쌓으면 결국 선형 함수 하나와 같아지기 때문에 비선형성을 넣어 학습 가능성을 높임
    그래디언트는 학습 중 모델을 더 정확하게 만들기 위해 계산되는 수치 변화량이고, 배치 정규화는 흐르는 숫자를 조정해 학습을 돕는 방법임. 위치 인코딩은 토큰들의 상대적 위치를 벡터로 알려줌
    Python의 @ 연산자는 __matmul__의 별칭으로 행렬 곱셈에 쓰임. 에폭은 데이터셋 전체를 한 번 학습하는 것이고, 배치는 매개변수 갱신 전에 한 번에 넣는 데이터 개수임
    어텐션은 LLM을 작동하게 하는 핵심으로, 입력 토큰들을 병렬 처리해 중간 텐서를 만든 뒤 출력 토큰들을 생성하는 데 사용함

    • 분야 밖에서는 “Karpathy”가 무엇을 뜻하는지 모를 수 있음. Andrej Karpathy를 “과학 커뮤니케이터이자 연구자”처럼 맥락과 함께 소개하면, 그의 글이나 영상을 참고하라는 뜻이 더 분명해짐
    • 토큰은 단순히 텍스트 조각의 정수 식별자라기보다, 자체로 유용할 만큼 흔한 단어 조각에 가깝다고 보는 편이 초보자에게도 더 정확함
      예를 들어 writing, written, writer에 공통으로 들어가는 writ가 하나의 토큰이 될 수 있고, writer는 writ와 er로 토큰화될 수 있음
      임베딩은 이런 토큰들을 고유한 숫자 표현으로 바꾸는 단계임
    • 선형 함수들을 합성하면 다시 선형 함수가 됨. 그래서 전부 선형이면 여러 층을 쌓아도 하나를 제외한 층은 낭비가 되며, 이를 피하려면 비선형성이 필요함
    • Karpathy의 영상 시리즈와 accompanying repo 외에, 학습 여정에서 특히 도움이 된 자료나 책이 더 있는지 궁금함
    • 배치 정규화가 정확히 무엇을 하고, 어떻게 도움이 되는지 궁금함

  • 모델의 기존 구현과 체크포인트가 있다면, 자신의 구현이 맞는지 확인하는 가장 효과적인 방법은 그 체크포인트를 불러와 출력값을 비교하는 것임
    출력이 맞지 않으면 대개 세부 구현을 잘못한 것이고, 각 층을 체계적으로 따라가며 실제 차이를 찾을 수 있음. 그러다 기존 구현의 이상한 점을 발견할 수도 있음
    이건 모델 자체에 대한 이야기이고, 학습은 별개의 축임. 그래도 하이퍼파라미터를 어느 정도 비슷하게 맞췄다면 모델 구현이 정확할 때 대체로 괜찮아짐

  • 논문 읽는 법과 해당 논문의 내용 모두 좋고, Karpathy의 Makemore 시리즈도 추천함

  • 요약 조언들이 아주 좋고, 텐서 shape를 assert하라는 조언은 어떤 일반적인 선형대수 라이브러리에도 적용된다고 봄. 복잡한 선형대수 코드를 짤 때는 작은 단계로 가고 방어적으로 코딩하는 것이 매우 중요함
    주류 언어에서 선형대수를 프로그래밍하는 일은 컴파일 타임 shape 검사가 없어서 끔찍함. 텐서의 shape는 타입의 일부가 되어야 하고, 3x4와 3x4를 전치 없이 곱하려 하면 컴파일 자체가 안 되어야 함
    긴 계산을 돌린 뒤 차원 불일치 연산에서 실패하는 건 정말 최악임
    PyTorch 텐서도 장치가 정적으로 타입화되어야 한다고 봄. 지금은 CPU 메모리의 텐서와 GPU 메모리의 텐서를 곱하려 하면 런타임 오류가 남

답변달기