라마 모델, 코드와 그림으로 이해하기 파트 3 - RoPE를 중심으로

안녕하세요, 수달이입니다.

 

벌써 시리즈의 세 번째 글이네요. 오늘은 그동안 배운 것들을 종합하여 LlamaModel 클래스를 완성해 보겠습니다. 우선 구성요소부터 살펴볼까요? ت

LlamaModel 구성요소 살펴보기

def __init__(self, config: LlamaConfig):
    super().__init__(config)
    self.padding_idx = config.pad_token_id
    self.vocab_size = config.vocab_size

    self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
    self.layers = nn.ModuleList(
        [LlamaDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
    )
    self.norm = LlamaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
    self.rotary_emb = LlamaRotaryEmbedding(config=config)
    self.gradient_checkpointing = False

    # Initialize weights and apply final processing
    self.post_init()

 

위 코드는 LlamaModel의 __init__ 함수입니다. 네 가지 주요 모듈을 확인할 수 있습니다.

 

1. nn.Embedding: 입력 문장의 각 단어들을 컴퓨터가 이해할 수 있는 벡터로 변환합니다.

2. nn.ModuleList [LlamaDecoderLayer]: 지난 포스팅에서 다룬, 여러 개의 디코더 모듈이 연결된 리스트입니다.

3. LlamaRMSNorm: 역시 지난 시간에 배운 모듈이죠. RMS (root mean square) 를 이용한 정규화 모듈입니다.

4. LlamaRotaryEmbedding: 각 단어 간의 "상대적인" 위치 정보를 표현할 수 있는 위치 임베딩 (position embedding) 입니다. 

 

오늘은 이미 다룬 모듈들은 제외하고 나머지 두 모듈, Embedding 및 LlamaRotaryEmbedding에 대해 자세히 공부해 보겠습니다.

Embedding 모듈

Embedding 모듈은 쉽게 말해 단어 사전입니다. 단어 벡터들을 쭉 나열해 놓고 입력 단어에 해당하는 벡터들을 쏙쏙 뽑아 쓰는 것이죠. 일반적인 사전에서 단어의 뜻을 찾는 것처럼, 임베딩 모듈에서는 단어에 대한 벡터를 찾을 수 있습니다. 아래 그림처럼요.

 

 

 

하지만 위 그림에는 중요한 점 한 가지가 생략되어 있습니다. 실제로 컴퓨터에 저장된 임베딩 모듈은 텍스트 단어 자체를 색인 (index) 으로 사용하지 않습니다. 모든 인덱스는 컴퓨터가 이해할 수 있는 '숫자'로 구성되어 있죠. 따라서 입력 문장은 먼저 토크나이저 (tokenizer) 를 통해 토큰 단위로 분해된 후, 각 토큰이 숫자로 인코딩 (encoding) 되는 과정을 거칩니다. 

 

 

 

그림의 오른쪽에 위치한 임베딩 사전의 크기는 두 가지 요소로 정의됩니다. 바로 '단어 수 (vocab_size)' 와 '벡터의 차원 (hidden_size)' 입니다. 위 그림에 빗대면 테이블의 행과 열의 크기에 해당합니다. 참고로, 라마 3.1 8B 모델의 경우 128,256개의 단어를 가지며, 각 단어 벡터는 4,096차원으로 구성됩니다.

 

정리하자면, 임베딩 모듈은 숫자로 인코딩 된 토큰들을 입력값으로 받아, 각 토큰을 고유한 벡터로 변환하는 역할을 합니다. 입력 문장의 길이를 S라고 한다면 (즉, 토큰의 개수가 S개라면), 임베딩 모듈은 최종적으로 S x hidden_size 크기의 행렬 (matrix) 을 출력합니다.

LlamaRotaryEmbedding

다음으로 배울 모듈은 로터리 임베딩, RoPE 모듈입니다. 이 모듈은 라마 이해하기 파트 1 포스팅에서 어텐션 모듈의 인풋으로 잠시 언급헸었죠. RoPE 논문 원본에는 수식이 가득하지만, (저를 포함하여) 수학이 두려운 분들을 위해 최대한 수식을 적게 사용해서 설명해보겠습니다. 😊

 

RoPE 개념설명

 

RoPE의 핵심은 벡터의 회전입니다. 그래서 이름도 "회전하는" (rotary) 임베딩이죠. 토큰의 위치를 나타내야 하는 임베딩인데 굳이 왜 벡터를 회전시킬까요? 그 이유는 바로, 회전된 두 벡터의 내적 (dot product) 을 통해 토큰 사이의 상대적인 거리를 나타낼 수 있기 때문입니다. 

2차원 벡터 (x_1, y_1) 를 θ만큼 회전 (원점 중심, 반시계 방향으로) 시킨다면, 회전된 벡터 (x_2, y_2) 의 좌표는 다음과 같이 나타낼 수 있습니다.

x_2 = x·cosθ - y·sinθ

y_2 = x·sinθ + y·cosθ

 

위치에 따른 회전

 

RoPE는 각 토큰 벡터를 해당 토큰의 절대적인 위치 m에 비례하여 서로 다른 각도로 회전시킵니다. 즉, 첫 번째 토큰은 m*θ = 1*θ = θ, 두 번째 토큰은 2θ, 세 번째 토큰은 3θ만큼 회전시키는 식이죠.

 

자, 그럼 앞서 예로 든 벡터 (x_1, y_1)가 첫 번째 토큰이고, 여기서 두 칸 떨어진 위치에 벡터 (a_1, b_1)이 있다고 가정해 봅시다. 이때, 벡터 (a_1, b_1)은 세 번째 토큰이므로 RoPE를 적용하면 3θ만큼 회전합니다. 그리고 위와 동일한 공식을 이용하면 다음과 같이 표현됩니다.

 

a_2 =  a·cos3θ - b·sin3θ

b_2 = a·sin3θ + b·cos3θ

 

이제 이 두 벡터의 내적을 구하면 (수식 과정은 생략할게요!) 최종값은 다음과 같습니다.

 

(ax + by) cos(3θ-θ) + (ay - bx) sin(3θ-θ) = (ax + by) cos(2θ) + (ay - bx) sin(2θ)

 

만약 벡터 (a_1, b_1)가 다섯 번째 토큰이었다면, 내적값은 다음과 같습니다.

 

(ax + by) cos(5θ-θ) + (ay - bx) sin(5θ-θ) = (ax + by) cos(4θ) + (ay - bx) sin(4θ)

 

자, 이제 패턴이 보이시나요? 각 토큰 벡터를 절대 위치에 따라 회전시키면, 그 내적값은 두 토큰의 상대적 거리에 의해 결정됩니다. RoPE는 이 간단하지만 똑똑한 로직을 활용하여 모델에 토큰의 상대적 위치 정보를 주입하는 것입니다. 

 

차원(dimension)에 따른 회전

 

지금까지 2차원 벡터의 회전과 그 특징에 대하여 배웠습니다. 그러나, 기억하시겠지만, LLM이 다루는 토큰 벡터는 보통 2차원이 아닙니다. 4096, 8192 등 훨씬 높은 차원의 벡터를 사용하죠. 4096차원의 벡터를 회전시키는 건, 컴퓨터에게도 굉장히 버거운 작업일 것입니다. 그래서 RoPE는 여기서 똑똑한 트릭을 씁니다. 바로 전체 차원을 여러 개의 2차원의 벡터로 쪼개는 것입니다. 예를 들어, 4096차원의 벡터는 총 2048개의 2차원 벡터 쌍이 되는 거죠. 그리고 각 쌍을 다른 각도로 회전시킵니다. 최종적으로, 2차원 벡터 쌍의 회전 각도는 토큰의 절대 위치 m뿐만 아니라 벡터 쌍의 인덱스 j에 의해서도 결정됩니다. 여기서 base는 고정된 값으로 모델 config.json의 rope_theta에 해당합니다. 

 

 

요약하자면, RoPE는 벡터의 회전을 통하여 토큰 간의 상대적인 거리를 나타내며, 회전 각도는 토큰의 절대 위치와 차원에 따라 달라집니다. 이제 코드를 통해 복습해 봅시다.

 

RoPE 코드 이해 - __init__

 

먼저 RoPE 모듈의 __init__ 함수입니다.

 

def __init__(self, config: LlamaConfig, device=None):
    super().__init__()
    # BC: "rope_type" was originally "type"
    if hasattr(config, "rope_scaling") and config.rope_scaling is not None:
        self.rope_type = config.rope_scaling.get("rope_type", config.rope_scaling.get("type"))
    else:
        self.rope_type = "default"
    self.max_seq_len_cached = config.max_position_embeddings
    self.original_max_seq_len = config.max_position_embeddings

    self.config = config
    self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]

    inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device)
    self.register_buffer("inv_freq", inv_freq, persistent=False)
    self.original_inv_freq = self.inv_freq

 

이 코드 블록에서 가장 중요한 부분은 rope_init_fn, 즉 RoPE 초기화 함수입니다. RoPE 초기화 함수는 역주파수 (inv_freq) 라는 RoPE의 핵심 파라미터를 계산하는데요. 이것이 바로 이전 섹션에서 언급한 회전 각도 θ_j 에 해당합니다. 가장 기본적인 형태의 RoPE 초기화 함수를 살펴보겠습니다. (라마 3.1 모델은 여기서 살짝 변형된 형태의 함수를 사용합니다.)

* 가독성을 위해 일부 부분은 생략했습니다. 전체 코드는 여기에서 보실 수 있습니다. 

 

def _compute_default_rope_parameters(config):
    base = config.rope_theta
    partial_rotary_factor = config.partial_rotary_factor if hasattr(config, "partial_rotary_factor") else 1.0
    head_dim = getattr(config, "head_dim", None) or config.hidden_size // config.num_attention_heads
    dim = int(head_dim * partial_rotary_factor)

    attention_factor = 1.0  # Unused in this type of RoPE

    # Compute the inverse frequencies
    inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2) / dim))
    
    return inv_freq, attention_factor

 

코드의 마지막에서 두 번째 줄을 보시면 역주파수의 계산 공식을 보실 수 있습니다. 이 공식을 그래프로 그려볼까요? 

 

여기서 x축은 벡터의 차원, y축은 역주파수 (또는 회전각) 를 나타냅니다. 차원이 낮을수록 회전각이 크고, 토큰 위치 m이 조금만 변해도 회전 각도 m * θ_j가 크게 변합니다. 반면, 고차원의 벡터 쌍들은 회전각이 작아, 토큰 위치 m의 변화에도 살짝만 변할 뿐입니다. 이렇게 차원에 따라 벡터 쌍들을 다른 속도로 회전시키면, 각 벡터 쌍들을 통해 서로 다른 정보를 수집할 수 있습니다. 위치에 민감한 저차원의 벡터 쌍으로는 가까운 토큰과의 관계를 파악하고, 상대적으로 둔감한 고차원의 벡터 쌍으로는 멀리 떨어진 토큰과의 관계, 즉 거시적인 정보를 학습할 수 있는 것이죠.

 

RoPE 코드 이해 - forward

 

다시 RoPE 개념 설명 섹션에서 언급한 벡터의 회전 공식을 떠올려보면, 우리에게 남은 것은 회전각에 따른 코사인과 사인 값을 계산하는 일입니다. 이 작업이 바로 forward 함수에서 진행됩니다. 역시 가독성을 위해 일부 라인은 생략했습니다. (전체 코드는 여기)

 

def forward(self, x, position_ids):

    inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1)
    position_ids_expanded = position_ids[:, None, :].float()
    
    with torch.autocast(device_type=device_type, enabled=False):
        freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)
        emb = torch.cat((freqs, freqs), dim=-1)
        cos = emb.cos()
        sin = emb.sin()

    return cos, sin

 

 

첫 두 줄은 토큰 위치 m과 회전 주파수 θ_j을 배치 (batch) 단위로 곱하기 위한 준비 과정입니다. self.inv_freq 에는 앞서 계산한 θ_j 값이 들어있고, 모양 변환 과정을 거쳐 (batch_size, d/2, 1) 차원을 갖습니다 (d/2는 벡터 '쌍'의 개수입니다). position_ids는 토큰의 위치 값 m들이 담긴 텐서이며, 모양은 (batch_size, sequence_length) 입니다. 이 또한 모양 변환을 통해 (batch_size, 1, sequence_length)로 만듭니다.

 

이렇게 준비된 두 텐서를 곱하고(@) 모양을 살짝 바꿔주면, (batch_size, sequence_length, d/2) 모양의 텐서가 나옵니다. 즉, freqs는 m*θ_j 값을 담고 있는 텐서입니다. 다음으로는 freqs를 자기 자신과 이어 붙여서, 차원을 d/2에서 d로 2배 늘립니다. 이제 완성된 회전각 텐서에 cos과 sin 함수를 각각 적용하여, 최종적으로 필요한 코사인 텐서와 사인 텐서가 만들어집니다. 각각 (batch_size, sequence_length, d) 모양의 텐서입니다. 

 

RoPE의 적용 - 벡터는 언제 회전할까

 

RoPE forward 함수를 통해 cos, sin 값을 계산해두었는데요. 실제로 이 값들은 어텐션 모듈에서 사용됩니다. RoPE는 어텐션 모듈 내에서 쿼리 및 키 벡터에만 적용됩니다. 즉, 어텐션 모듈은 아래의 apply_rotary_pos_emb 함수를 호출해 쿼리, 키 벡터를 변형시킵니다. 

 

def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):
    cos = cos.unsqueeze(unsqueeze_dim)
    sin = sin.unsqueeze(unsqueeze_dim)
    q_embed = (q * cos) + (rotate_half(q) * sin)
    k_embed = (k * cos) + (rotate_half(k) * sin)
    return q_embed, k_embed

 

첫 두 줄의 unsqueeze 함수는 쿼리 및 키 벡터와의 행렬 곱셈이 가능하도록 cos, sin의 차원을 확장시킵니다. 그다음 줄에서는 rotate_half라는 함수가 눈에 띄는데요.

 

def rotate_half(x):
    """Rotates half the hidden dims of the input."""
    x1 = x[..., : x.shape[-1] // 2]
    x2 = x[..., x.shape[-1] // 2 :]
    return torch.cat((-x2, x1), dim=-1)

 

이 함수는 벡터 회전 공식을 똑똑하게 계산하기 위한 트릭입니다. 입력 벡터 x를 반으로 쪼갠 후, 뒷부분을 앞으로 가져와 부호를 바꾸고, 앞부분은 뒤로 보내는 역할을 합니다. 이렇게 하면 여러 벡터 쌍을 한 번에, 훨씬 효율적으로 처리할 수 있습니다. 2차원 벡터로 예를 들어보겠습니다.

 

 

 

 

마지막 단계의 수식이 익숙하지 않으신가요? 포스팅 앞부분에서 언급했던 벡터 회전 공식입니다. 이 rotate_half 함수를 사용하면 여러 벡터 쌍을 동시에 처리할 수 있고 (위 그림에서 벡터들이 옆으로 길어진다고 상상해보세요), for-loop을 사용하는 것보다 훨씬 효율적으로 계산할 수 있습니다. 

 

드디어! 길고 길었던 RoPE 모듈 설명이 끝났습니다. RoPE가 개념적으로는 간단하지만, 그 원리를 자세히 이해하는 데는 꽤 시간이 드는 것 같습니다. 이제 정말 쉬운 부분만 남았으니 마지막 섹션으로 넘어가 보죠. :)

LlamaModel Forward

임베딩과 RoPE 모듈까지 배우면서 LlamaModel의 주요 구성요소에 대한 공부는 모두 끝났는데요. 이제 이 모듈들이 어떻게 조합되어 최종적으로 문맥 정보가 풍부한 벡터를 만들어내는지, forward 함수를 통해 알아보겠습니다. 이 LlamaModel이 바로 다음 포스팅에서 다룰 최종 언어 모델의 핵심 엔진이 됩니다. (전체 코드는 여기를 참고해 주세요.)

 

 

 

LlamaModel 클래스의 인풋은 Input ID입니다. 입력 문장이 토크나이저를 통해 토큰 단위로 분해된 후, 숫자로 변환된 것이죠. 위 그림에서는 단순화를 위해 토큰 ID 하나만 표시했지만, 실제로는 여러 개의 토큰 ID를 인풋으로 넘길 수 있습니다. 각 토큰은 임베딩 모듈을 거쳐 더 큰 차원 (hidden_size) 으로 확장되고, 여러 개의 디코더 모듈을 순차적으로 통과합니다. 라마 3.1 8B 모델의 경우, 총 32개의 디코더 레이어를 통과하게 됩니다. 참고로 각 디코더는 RoPE 모듈이 미리 계산해 둔 cos, sin 값을 공유하여 사용합니다. 디코더 시리즈를 통과한 벡터는 마지막으로 정규화 모듈을 거칩니다. 각 토큰은 임베딩 모듈을 통과한 순간부터 동일한 차원 (hidden_size) 을 유지하므로, 최종 아웃풋은 (batch_size, sequence_length, hidden_size)의 모양을 가집니다. 정리하자면 LlamaModel의 순서는 다음과 같습니다.

 

토큰 ID -> 임베딩으로 차원 확장 -> 디코더 시리즈 -> 정규화

마무리

생각보다 길고 길었던 세 번째 포스트가 마무리됐네요. 이번 포스트의 복병은 RoPE 모듈이었는데요. 논문으로 이해하기 어려웠던 RoPE의 개념을 좀 더 쉽게 이해하는 계기가 되셨기를 바랍니다. 이번 포스팅에서 꼭 기억해야 할 점은 다음과 같습니다.

 

1. RoPE의 핵심 개념:

     a. 두 토큰의 상대적 거리는 각 토큰 벡터를 회전시킨 뒤, 그 내적으로 나타낼 수 있습니다.

     b. 벡터는 cosθ, sinθ 함수를 통해 각도 θ만큼 회전합니다.

2. RoPE의 동작 방식:

    a. 쿼리 및 키 벡터를 각각 2차원의 벡터 쌍으로 나누어, (1) 토큰의 위치 (2) 차원 인덱스에 따라 다른 각도로 회전시킵니다.

    b. 어텐션 연산에 필요한 cos, sin 값을 미리 계산해두어 디코더 내 어텐션 모듈에 제공합니다. 

3. LlamaModel 클래스는 토큰 ID를 인풋으로 받아 임베딩, 디코더 시리즈를 통해 고차원의 벡터를 출력합니다.

 

자, 이제 정말로 어려운 파트는 모두 통과하셨습니다! 다음 포스팅에서는 최종 보스인 LlamaModelForCausalLM에 대해 알아보겠습니다. 오늘 배운 LlamaModel과는 무엇이 다른지 알아볼 테니 다음 시간까지 기다려주세요.

 

그럼 모두 Happy LLM!