TDD를 통해 ModernBERT 밑바닥부터 이해하기

import tensorflow as tf import numpy as np import pytest from ModernGTETFModel import create_local_sliding_window_mask @pytest.mark.parametrize( "batch_size, seq_len, window_size", [ (1, 10, 4), # batch_size 1, sequence_length 10, window_size 4 (2, 15, 6), # batch_size 2, sequence_length 15, window_size 6 (3, 20, 8), # batch_size 3, sequence_length 20, window_size 8 ], ) def test_create_local_sliding_window_mask( batch_size: int, seq_len: int, window_size: int ): ATTENTION_MASK_DIMENSION: int = 1 HALF_WINDOW_SIZE: int = window_size // 2 # Given: 모든 위치가 유효하다고 전제 하므로 모든 원소를 0.0 으로 초기화 한다. global_mask: tf.zeros = tf.zeros( (batch_size, ATTENTION_MASK_DIMENSION, seq_len, seq_len), dtype=tf.float32 ) # When: 함수를 호출하여 슬라이딩 윈도우 마스크를 생성합니다. result: tf.Tensor = create_local_sliding_window_mask(global_mask, window_size) # Then: tensor 의 shape 가 (batch_size, 1, seq_len, seq_len) 인지를 검증해봅니다. assert result.shape == (batch_size, ATTENTION_MASK_DIMENSION, seq_len, seq_len) # Then: 각 (i, j) 위치에 대하여 윈도우의 내부이면 0.0 이요, 그 외부이면 매우 작은 값이 나온다는 것을 검증 합니다. result_np: np.ndarray = result.numpy() # 인덱스 배열을 생성한 후에, 각 위치에 저장된 값의 절댓값 을 계산 합니다. # 절댓값 행렬 (seq_len, seq_len) 은 두 1D 배열 간의 외적 차이를 계산하여 각 토큰 위치 간의 거리를 나타냅니다. indices: np.ndarray = np.arange(seq_len) diff_matrix: np.ndarray = np.abs(np.subtract.outer(indices, indices)) # 절대값의 차이가 half_w 이하인 윈도우 내부 위치는 True이고, 외부는 False인 마스크를 만들어봅니다. valid_mask: np.ndarray = diff_matrix <= HALF_WINDOW_SIZE # 결과 행렬의 첫 번째 배치의 값을 추출해봅니다. shape 는 (seq_len, seq_len) 이 됩니다. # 추출한 이후에는 윈도우 내부의 값과 윈도우 외부의 값을 가져와봅니다. result_matrix: np.ndarray = result_np[0, 0] window_inner_matrix: np.ndarray = result_matrix[valid_mask] window_outer_matrix: np.ndarray = result_matrix[~valid_mask] # 내부 값들은 0.0 이어야 합니다. np.testing.assert_allclose( actual=window_inner_matrix, desired=0.0, atol=1e-6, err_msg="윈도우 내부 matrix에서 0.0이 아닌 값이 존재하네요.", ) # 외부 값들은 매우 작은 값 이어야 합니다. 여기서 매우 작은 값은 -1e8 미만을 의미 합니다. LOW_THRESHOLD: float = -1e8 assert np.all( outside_values < LOW_THRESHOLD for outside_values in window_outer_matrix ) if __name__ == "__main__": pytest.main([__file__])

def create_local_sliding_window_mask(global_mask_4d, window_size): """ ModernBERT의 local attention을 위한 "양방향 슬라이딩 윈도우" 마스크를 만든 뒤, 원본 global_mask(패딩 토큰 마스킹)와 결합하여 최종 4D float 마스크를 반환합니다. - PyTorch 코드를 예시로: distance = |i - j| window_mask = (distance <= window_size//2) sliding_window_mask = global_attention_mask.masked_fill(~window_mask, -∞) 에 대응. Args: global_mask_4d: shape (batch_size, 1, seq_len, seq_len) - 이미 패딩 토큰 부분은 -∞(float('-inf')) 또는 0.0 으로 구성된 float 마스크. - 보통 BERT식 마스크는 "유효 위치=0.0, 무효 위치=-∞" 형태임 window_size (int): 로컬 윈도우 크기 (예: 128) Returns: final_local_4d (tf.Tensor): shape (batch_size, 1, seq_len, seq_len) - 윈도우 내부이면서 유효 토큰이면 0.0 - 윈도우 밖이거나 패딩이면 -∞ """ # global_mask_4d: [B, 1, S, S] batch_size = tf.shape(global_mask_4d)[0] seq_len = tf.shape(global_mask_4d)[-1] # (S, S)에서의 distance 계산 rows = tf.range(seq_len)[:, None] # shape (S,1) cols = tf.range(seq_len)[None, :] # shape (1,S) distance = tf.abs(rows - cols) # shape (S,S) # distance가 window_size//2 이내면 True half_w = window_size // 2 window_bool_2d = tf.less_equal(distance, half_w) # (S,S), True/False # True면 0.0, False면 -∞인 2D float mask inside_0 = tf.zeros([seq_len, seq_len], dtype=TFDTYPE) outside_inf = tf.fill([seq_len, seq_len], TFDTYPE.min) # -1e9) local_mask_2d = tf.where(window_bool_2d, inside_0, outside_inf) # (S,S), float # shape 확장: [1,1,S,S] local_mask_4d = local_mask_2d[None, None, :, :] # batch 차원만큼 복제: [B,1,S,S] local_mask_4d = tf.tile(local_mask_4d, [batch_size, 1, 1, 1]) # PyTorch의 sliding_window_mask = global_mask.masked_fill(~window_bool_2d, -inf)에 해당 # => local_mask_4d가 window 밖을 -∞로 만들어놓았으므로 # 밖일 때 -∞, 안일 때 0.0 # => global_mask_4d도 이미 "패딩 부분 -∞, 정상 부분 0.0" 형태이므로 # 합산하면 "둘 중 하나라도 -∞이면 -∞"라는 효과가 남 final_local_4d = global_mask_4d + local_mask_4d # (B,1,S,S) return final_local_4d

@pytest.mark.parametrize( "batch_size, seq_len, d_model, num_heads, dropout_rate, training, layer_id", [ ( 1, 128, 64, 8, 0.0, False, 0, ), # Case 1: 배치를 먼저 작게하고, 기본 설정으로 만들어봅니다. layer_id가 0인 경우 layer norm을 적용하지 않고 원시 임베딩. (2, 256, 128, 16, 0.0, False, 0), # Case 2: 중간 배치로 올려보겠습니다. ( 3, 512, 64, 4, 0.1, True, 1, ), # Case 3 : 3 개의 배치에 대해 dropout 0.1, training = True 이고 layer_id가 1 인 상태입니다. # layer_id가 1 인 경우 입력을 먼저 layer norm 하고 q, k, v를 계산합니다. 안정적인 학습을 위한 전략입니다. ], ) def test_ROPE가_존재하지_않는_멀티헤드_어텐션을_테스트한다( batch_size, seq_len, d_model, num_heads, dropout_rate, training, layer_id ): """ test_ROPE가_존재하지_않는_멀티헤드_어텐션을_테스트한다: rope_embed 없이 MHA 의 기본 출력을 검증해봅시다. Given: - 임의의 값으로 채워진 dummy input tensor (shape: [batch_size, seq_len, d_model]). - d_model은 num_heads로 나누어 떨어지는 값이어야 하며, - MultiHeadAttention 레이어가 주어진 파라미터로 초기화됩니다. When: - 해당 레이어가 mask와 rope_embeds 없이 dummy input에 대해 forward pass를 수행합니다. Then: - 출력 tensor의 shape는 [batch_size, seq_len, d_model]이어야 합니다. - 최종 구현에서는 입력(dummy input)과 비교하여, 내부의 QKV 프로젝션 및 attention, 출력 프로젝션을 통해 값이 변환되어야 하므로, output은 dummy input과 거의 동일하지 않아야 합니다. """ # Given : d_model 이 num_heads로 나누어 떨어져야만 합니다. if d_model % num_heads != 0: pytest.fail("d_model must be divisible by num_heads") # Given : 입력 텐서를 무작위의 값으로 생성합니다. shape는 [batch_size, seq_len, d_model] 으로 구성되어야 합니다. dummy_input = tf.random.uniform( shape=(batch_size, seq_len, d_model), dtype=tf.float32 ) # Given : MHA 인스턴스를 생성합니다. multi_head_attention: MultiHeadAttention = MultiHeadAttention( d_model=d_model, num_heads=num_heads, dropout_rate=dropout_rate, layer_id=layer_id, ) # When : 입력 텐서를 전달하여 레이어를 호출해봅니다. output: tf.Tensor = multi_head_attention( inputs=dummy_input, mask=None, rope_embeds=None, training=training ) # Then : 출력 텐서의 shape 가 (batch_size, seq_len, d_model) 과 동일해야만 합니다. expected_shape: tuple = (batch_size, seq_len, d_model) assert ( output.shape == expected_shape ), f"Expected output shape {expected_shape}, but got {output.shape}" # Then : 출력값이 수치적으로 안정적인지를 shape에 대하여 확인해본다. np.testing.assert_allclose( output.numpy(), output.numpy(), err_msg="Output values are not consistent." ) # Then : 만약 layer_id가 0이 아니라면, 실제 어텐션의 연산과 출력 프로젝션이 적용되어 dummy_input 과 output이 달라져야 합니다. if np.allclose(dummy_input.numpy(), output.numpy(), atol=1e-6): pytest.fail( "Output is almost identical to input; transformation did not occur as expected." )

@pytest.mark.parametrize( "batch_size, seq_len, d_model, num_heads", [ (1, 10, 64, 8), # Case 1: 간단한 설정 (2, 15, 128, 8), # Case 2: 조금 더 큰 설정 ], ) def test_ROPE가_존재하는_멀티헤드_어텐션을_테스트한다( batch_size, seq_len, d_model, num_heads ): """ Given a dummy input tensor and a MultiHeadAttention layer, and given dummy rope embeddings (cos, sin) with the correct shape, When the layer is invoked with these rope embeddings, Then the output tensor should have the same shape as (batch_size, seq_len, d_model). """ # Given: d_model이 num_heads로 나누어 떨어져야 합니다. if d_model % num_heads != 0: pytest.skip("d_model must be divisible by num_heads") # Given: 입력 텐서를 무작위로 생성 (shape: [batch_size, seq_len, d_model]) dummy_input = tf.random.uniform((batch_size, seq_len, d_model), dtype=tf.float32) # Given: dropout 없이, layer_id=1 (즉, 로타리 임베딩 적용 케이스)로 MultiHeadAttention 인스턴스를 생성합니다. mha = MultiHeadAttention(d_model, num_heads, dropout_rate=0.0, layer_id=1) # Given: 각 헤드의 차원 계산 (d_model // num_heads) head_dim = d_model // num_heads # Given: rope_embeds를 위한 더미 cos, sin 텐서를 생성합니다. # - 기대하는 rope_embeds의 shape는 [seq_len, 2 * head_dim] # Q: dummy_cos 가 모두 1 이라는 cosine tensor 의 의미는 ? # 실제로는 각 토큰의 위치가 cos 값에 따라 달라지겠지만, 더미 값으로 모든 위치를 1을 반환하게 해서 회전 효과가 없다는 것을 테스트한다. # Q: dummy_sin 이 모두 0 인 sin tensor 의 의미는 ? # 실제로는 임베딩이 sine 값 이 위치에 따라 달라지지만 모든 값을 0 으로 반환하게 해서 회전 효과가 없다는 것을 테스트한다. dummy_cos = tf.ones((seq_len, 2 * head_dim), dtype=tf.float32) # 모든 값 1 dummy_sin = tf.zeros((seq_len, 2 * head_dim), dtype=tf.float32) # 모든 값 0 rope_embeds = (dummy_cos, dummy_sin) # When: rope_embeds를 포함하여 MultiHeadAttention 레이어를 호출합니다. output = mha(dummy_input, mask=None, rope_embeds=rope_embeds, training=False) # Then: 출력 텐서의 shape가 (batch_size, seq_len, d_model)와 동일해야 합니다. expected_shape = (batch_size, seq_len, d_model) assert ( output.shape == expected_shape ), f"Output shape mismatch with rope_embeds: expected {expected_shape}, got {output.shape}" # 그리고, 함수형 검증 도구를 활용하여 출력값의 일관성을 확인합니다. np.testing.assert_allclose( output.numpy(), output.numpy(), err_msg="Output values are inconsistent when using rope_embeds.", )

class MultiHeadAttention(keras.layers.Layer): def __init__( self, d_model: int, num_heads: int, dropout_rate: float = 0.1, layer_id: int = 0, **kwargs, ) -> None: ... def apply_rotary_pos_emb( self, q: tf.Tensor, k: tf.Tensor, cos: tf.Tensor, sin: tf.Tensor ) -> tuple[tf.Tensor, tf.Tensor]: """ Applies the rotary positional embedding to query and key tensors. :param q: Query tensor :param k: Key tensor :param cos: Cosine tensor for rotary embedding :param sin: Sine tensor for rotary embedding :return: Tuple containing modified query and key tensors (q_embed, k_embed) """ # 이 함수는 주어진 텐서를 마지막 차원에서 두 부분으로 분할한 후, # 두 번째 절반의 부호를 반전시키고, 순서를 뒤바꿔서 첫 번째 절반과 결합합니다. # 예를 들어, 입력 벡터가 [x1, x2]라면, 출력은 [-x2, x1]가 됩니다. def _rotate_half(x: tf.Tensor) -> tf.Tensor: # 마지막 차원(axis=-1)을 두 개의 동일한 크기의 텐서로 분할합니다. x1, x2 = tf.split(value=x, num_or_size_splits=2, axis=-1) # x2의 부호를 반전한 후, x1과 함께 이어 붙입니다. # Q: 왜 마지막 차원을 분할하냐? # Transformer에서는 각 토큰마다 하나의 임베딩 벡터가 있으며 # MultiHeadAttention에서는 각 헤드마다 그 벡터의 일부(즉, head_dim)를 사용합니다. # 이 마지막 차원이 실제 피처(특징)들이 담긴 부분이기 때문에, 여기서 정보를 반으로 나누어 회전 변환을 적용하는 것이 자연스럽습니다. return tf.concat([-x2, x1], axis=-1) # 여기서 cos, sin은 원래 [seq_len, 2 * head_dim] 형태일 수 있습니다. # reshape을 통해 [1, 1, seq_len, head_dim] 형태로 바꿉니다. # - 첫 번째 1: 배치 차원에 대해 확장 (모든 배치에 동일한 값을 사용) # - 두 번째 1: 헤드 차원에 대해 확장 (모든 헤드에 동일하게 적용) # - -1: seq_len을 그대로 유지 # - self.head_dim: 마지막 차원은 각 헤드의 차원 크기 num_heads, seq_len, head_dim = 1, 1, -1 cos = tf.reshape(cos, [num_heads, seq_len, head_dim, self.head_dim]) sin = tf.reshape(sin, [num_heads, seq_len, head_dim, self.head_dim]) # 위의 reshape 후, cos와 sin은 [1, 1, 2 * seq_len, head_dim]의 shape를 가지게 됩니다. # 하지만 실제로 적용할 때는 query q와 key k의 시퀀스 길이에 맞춰야 하므로, # 필요한 부분만 슬라이싱하여 [1, 1, seq_len, head_dim]의 shape로 맞춥니다. cos = tf.reshape(cos, [1, 1, -1, self.head_dim]) sin = tf.reshape(sin, [1, 1, -1, self.head_dim]) # tf.shape(q)[2]는 q 텐서의 세 번째 차원, 즉 시퀀스 길이 S를 나타냅니다. # 첫 번째와 두 번째 차원은 그대로 유지하고, 세 번째 차원에서 처음 seq_len (즉, tf.shape(q)[2]) 개의 값만 선택합니다. cos = cos[:, :, : tf.shape(q)[2], :] sin = sin[:, :, : tf.shape(q)[2], :] # 이제 rotary positional embedding을 적용합니다. # 각 query 벡터에 대해, cosine 값과 sine 값을 사용하여 회전 변환을 수행합니다. # 공식은 다음과 같습니다: # q_embed = (q * cos) + (rotate_half(q) * sin) # 이는, 각 원소를 두 부분으로 나누고, 이 두 부분에 대해 회전 행렬의 효과를 벡터화한 형태라고 볼 수 있습니다. q_embed = (q * cos) + (_rotate_half(q) * sin) # 동일하게, key 벡터에도 같은 변환을 적용합니다. k_embed = (k * cos) + (_rotate_half(k) * sin) # 최종적으로, 변환된 query와 key 텐서를 반환합니다. return q_embed, k_embed

@pytest.mark.parametrize( "dim, max_pos, base, scaling_factor, test_seq_len", [ ( 32, 50, 160000.0, 1.0, 30, ), # 케이스 1: test_seq_len(30)이 max_pos(50)보다 작음 → 캐시된 값 사용 기대 ( 32, 50, 160000.0, 1.0, 60, ), # 케이스 2: test_seq_len(60)이 max_pos(50)보다 큼 → 새로운 캐시 계산됨 ( 64, 100, 160000.0, 1.0, 80, ), # 케이스 3: test_seq_len(80)이 max_pos(100) 이하 → 캐시된 값 사용 기대 ], ) def test_ntk_scaling_rotary_embedding_shapes( dim, max_pos, base, scaling_factor, test_seq_len ): """ Given a NTKScalingRotaryEmbedding layer instantiated with specific parameters, and given a dummy input tensor with a specified sequence length, When the layer is called with that input and seq_len is provided, Then the returned cos and sin tensors should have shape (seq_len, dim), and if seq_len <= max_position_embeddings, they should match the cached values. """ # Given: NTKScalingRotaryEmbedding 인스턴스를 생성합니다. rotary_layer = NTKScalingRotaryEmbedding( dim=dim, # 임베딩 차원 max_position_embeddings=max_pos, # 최대 position embeddings base=base, # base 값 (예: 160000.0) scaling_factor=scaling_factor, # 스케일링 팩터 (예: 1.0) ) # Given: 더미 입력 텐서를 생성합니다. 실제 값은 중요하지 않고, 단지 seq_len 정보를 전달하기 위한 용도입니다. dummy_input = tf.random.uniform((1, test_seq_len, dim), dtype=tf.float32) # When: NTKScalingRotaryEmbedding의 call 메서드를 호출하여, cos와 sin 값을 계산합니다. cos, sin = rotary_layer(dummy_input, seq_len=test_seq_len) # Then: 반환된 cos 텐서의 shape가 (test_seq_len, dim)인지 검증합니다. assert cos.shape == ( test_seq_len, dim, ), f"Expected cos shape ({test_seq_len}, {dim}), got {cos.shape}" # Then: 반환된 sin 텐서의 shape가 (test_seq_len, dim)인지 검증합니다. assert sin.shape == ( test_seq_len, dim, ), f"Expected sin shape ({test_seq_len}, {dim}), got {sin.shape}" # Then: test_seq_len이 max_pos 이하인 경우, 캐시된 cos와 sin의 슬라이스와 동일해야 합니다. if test_seq_len <= max_pos: # 캐시된 cos 값을 test_seq_len 만큼 슬라이스합니다. cached_cos = rotary_layer.cos_cached[:test_seq_len] # 캐시된 sin 값을 test_seq_len 만큼 슬라이스합니다. cached_sin = rotary_layer.sin_cached[:test_seq_len] # np.testing.assert_allclose를 사용하여 계산된 cos와 캐시된 cos가 거의 동일한지 확인합니다. np.testing.assert_allclose( cos.numpy(), cached_cos.numpy(), atol=1e-6, err_msg="cos values do not match cached values when seq_len <= max_pos", ) # np.testing.assert_allclose를 사용하여 계산된 sin와 캐시된 sin가 거의 동일한지 확인합니다. np.testing.assert_allclose( sin.numpy(), cached_sin.numpy(), atol=1e-6, err_msg="sin values do not match cached values when seq_len <= max_pos", ) else: # test_seq_len이 max_pos보다 클 경우, 새로운 캐시가 생성되므로 캐시와 비교할 수 없습니다. # 대신, 계산된 cos와 sin 값이 모두 유한한(finite) 값인지 확인합니다. assert np.all( np.isfinite(cos.numpy()) ), "cos contains non-finite values for seq_len > max_pos" assert np.all( np.isfinite(sin.numpy()) ), "sin contains non-finite values for seq_len > max_pos"

class NTKScalingRotaryEmbedding(keras.layers.Layer): def __init__( self, dim, max_position_embeddings=8192, base=160000.0, scaling_factor=1.0, mixed_b=None, **kwargs, ): super().__init__(**kwargs) self.dim = dim self.max_position_embeddings = int(max_position_embeddings * scaling_factor) self.base = base self.scaling_factor = scaling_factor self.mixed_b = mixed_b # 짝수 인덱스에 대한 범위를 생성합니다. 예를 들어, dim=32이면 0,2,4,...,30 # 로타리 임베딩은 보통 임베딩 벡터의 두 개의 연속된 요소(예: [cos, sin])를 한 쌍으로 사용합니다. # 그래서 전체 임베딩 차원 중 절반만큼의 주파수를 계산하면 됩니다. # 여기서는 tf.range를 사용하여 0부터 self.dim까지 2씩 증가하는 값들을 생성하여, 짝수 인덱스만 선택합니다. indices: tf.Tensor = tf.range(0, self.dim, 2, dtype=tf.float32) # 기본 base_inv_freq 계산: 1 / (base^(i/dim)) # 각 짝수 인덱스에 대해, base 값에 따른 역수(inv_freq)를 계산합니다. # indices / self.dim: # 각 인덱스를 임베딩 차원(self.dim)으로 나누어, 0에서 1 사이의 비율 값을 얻습니다. # 예를 들어, 인덱스 0은 0/self.dim = 0, 인덱스 2는 2/self.dim 등으로 계산됩니다. # tf.pow(self.base, (indices / self.dim)): # self.base(예: 160000.0) 값을 위에서 계산한 비율만큼 거듭제곱합니다. # 이는 각 인덱스마다 서로 다른 주파수를 결정하는 역할을 합니다. # 1.0 / (...): # 이렇게 계산된 값의 역수를 취합니다. # 즉, inv_freq = 1 / (self.base^(index/self.dim))가 됩니다. # 이 역수 값들은 이후에 각 위치(position)와 곱해져, cos와 sin 값을 계산하는 데 사용됩니다. self.base_inv_freq = 1.0 / tf.pow( self.base, (indices / tf.cast(self.dim, tf.float32)) ) if self.mixed_b is None: scaled_base = self.base * self.scaling_factor self.scaled_inv_freq = 1.0 / tf.pow( scaled_base, (indices / tf.cast(self.dim, tf.float32)) ) self.scaled_inv_freq = self.scaled_inv_freq / tf.pow( self.scaling_factor, 2.0 / tf.cast(self.dim, tf.float32) ) else: pass self._build_initial_cache() def _build_initial_cache(self): """ 모델이 최대 시퀀스 길이(max_position_embeddings)까지 사용할 수 있도록, 각 위치에 대한 cosine과 sine 값을 미리 계산하여 저장해둡니다. 이렇게 미리 계산해두면, 모델이 실행될 때마다 동일한 값을 반복해서 계산할 필요 없이 캐시된 값을 재사용할 수 있으므로 효율적입니다. """ # 1. 0부터 max_position_embeddings - 1까지의 정수를 생성합니다. # 예를 들어 max_position_embeddings가 8192라면, 0, 1, 2, ..., 8191을 담은 텐서를 만듭니다. t = tf.range(self.max_position_embeddings, dtype=tf.float32) # 2. tf.einsum("i,j->ij", t, self.scaled_inv_freq)를 사용하여 # 각 위치 t와 scaled_inv_freq 벡터 간의 외적을 곱하여, 각 위치에 고유한 "각도"를 생성하기 위해서입니다. # - t의 shape는 (max_position_embeddings,)입니다. # - self.scaled_inv_freq의 shape는 (dim/2,) (예를 들어, 만약 dim이 768이면, (384,)가 될 것입니다.) # - 결과적으로 freqs의 shape는 (max_position_embeddings, dim/2)가 됩니다. freqs = tf.einsum("i,j->ij", t, self.scaled_inv_freq) # 3. 이제 freqs 텐서를 두 번 반복하여(concatenation) shape를 (max_position_embeddings, dim)으로 만듭니다. # 왜 두 번 반복하냐면, 로타리 임베딩에서는 각 위치마다 cosine과 sine 값을 적용하는데, # 보통 임베딩 차원(dim)이 2로 나누어 떨어지므로, 각 절반(dim/2)씩 계산된 값을 두 번 이어 붙여서 전체 임베딩 차원(dim)을 채웁니다. emb = tf.concat([freqs, freqs], axis=-1) # 4. emb 텐서의 각 요소에 대해 cosine과 sine 값을 계산합니다. # 이때, emb의 각 원소는 (t * scaled_inv_freq)로 계산되었으므로, tf.cos(emb)와 tf.sin(emb)는 # 각 위치에 대한 고정된 cosine, sine 값이 됩니다. # 이렇게 계산된 값들을 self.cos_cached와 self.sin_cached에 저장해 두면, # 나중에 입력 시퀀스의 길이에 따라 슬라이스해서 사용할 수 있습니다. self.cos_cached = tf.cos(emb) self.sin_cached = tf.sin(emb) def _compute_new_cache(self, seq_len): # 1. 0부터 seq_len - 1까지의 정수를 생성합니다. # 예를 들어, seq_len이 9000이면, 0부터 8999까지의 텐서를 생성합니다. t = tf.range(seq_len, dtype=tf.float32) # 2. tf.einsum("i,j->ij", t, self.scaled_inv_freq)를 통해 # 각 위치에 대해 scaled_inv_freq 벡터와의 외적을 계산합니다. # 결과 freqs의 shape는 (seq_len, dim/2)가 됩니다. freqs = tf.einsum("i,j->ij", t, self.scaled_inv_freq) # 3. freqs를 두 번 이어 붙여 (concatenate) shape를 (seq_len, dim)으로 만듭니다. emb = tf.concat([freqs, freqs], axis=-1) # 4. emb에 대해 cosine과 sine 값을 계산하여 반환합니다. return tf.cos(emb), tf.sin(emb) def call(self, x, seq_len=None): # 1. 입력 x에서 시퀀스 길이를 결정합니다. # 만약 seq_len이 명시적으로 주어지지 않았다면, x의 두 번째 차원(시퀀스 길이)을 사용합니다. if seq_len is None: seq_len = tf.shape(x)[1] # 2. Eager 실행 모드인지 확인합니다. if tf.executing_eagerly(): # Eager 모드에서는 조건문(if)을 사용하여, 현재 시퀀스 길이가 캐시된 최대값보다 큰지 확인합니다. if seq_len > self.max_position_embeddings: # 만약 시퀀스 길이가 더 길다면, _compute_new_cache 메서드를 호출하여 새로운 cos, sin 값을 계산합니다. cos, sin = self._compute_new_cache(seq_len) else: # 그렇지 않으면, 미리 계산된 캐시(self.cos_cached, self.sin_cached)를 시퀀스 길이에 맞게 슬라이스합니다. cos = self.cos_cached[:seq_len] sin = self.sin_cached[:seq_len] else: # 3. 그래프 모드(비 Eager 모드)에서는 tf.cond를 사용하여 조건 분기를 수행합니다. # - tf.cond는 조건에 따라 두 개의 함수를 실행합니다. def use_new_cache(): return self._compute_new_cache(seq_len) def use_cached(): return self.cos_cached[:seq_len], self.sin_cached[:seq_len] # tf.cond를 사용하여 seq_len이 max_position_embeddings보다 큰지 확인한 후, 해당하는 함수를 실행합니다. cos, sin = tf.cond( tf.greater(seq_len, self.max_position_embeddings), use_new_cache, use_cached, ) # 4. 입력 텐서 x와 데이터 타입을 맞추기 위해, cos와 sin 텐서를 x의 데이터 타입으로 캐스팅합니다. cos = tf.cast(cos, x.dtype) sin = tf.cast(sin, x.dtype) # 5. 최종적으로, 계산된 cos와 sin 텐서를 반환합니다. return cos, sin

class DummyConfig: def __init__(self, d_model, num_heads, intermediate_size): self.hidden_size = d_model self.num_attention_heads = num_heads self.intermediate_size = intermediate_size self.global_rope_theta = 10000.0 # 글로벌 로타리 임베딩 base 값 self.max_position_embeddings = 512 # 전체 최대 시퀀스 길이 self.local_rope_theta = None # 로컬 어텐션 시 별도 값 없으면 None self.global_attn_every_n_layers = 2 # 예: 2번째 레이어마다 글로벌 어텐션 self.local_attention = 128 # 로컬 어텐션 윈도우 크기 @pytest.mark.parametrize( "batch_size, seq_len, d_model, num_heads, intermediate_size, dropout_rate, layer_id", [ ( 1, 10, 64, 8, 128, 0.0, 0, ), # Case 1: 작은 배치, dropout 없이, layer_id=0 (글로벌 어텐션 적용) ( 2, 15, 128, 8, 256, 0.0, 1, ), # Case 2: 중간 배치, dropout 없이, layer_id=1 (로컬 어텐션 적용) ( 3, 20, 64, 4, 128, 0.1, 1, ), # Case 3: dropout 적용, training=False로 실행하여도 dropout layer는 build됨 ], ) def test_transformer_block_output_shape( batch_size, seq_len, d_model, num_heads, intermediate_size, dropout_rate, layer_id ): """ Given a dummy input tensor of shape (batch_size, seq_len, d_model) and a DummyConfig with the specified parameters, When the TransformerBlock is invoked with a simple attention mask (all ones), Then the output tensor should have the same shape as the input tensor. """ # Given: d_model이 num_heads로 나누어 떨어지는지 확인합니다. if d_model % num_heads != 0: pytest.skip("d_model must be divisible by num_heads") # Given: DummyConfig 인스턴스를 생성합니다. config = DummyConfig(d_model, num_heads, intermediate_size) # Given: TransformerBlock 인스턴스를 생성합니다. # layer_id에 따라 내부에서 글로벌/로컬 설정이 달라집니다. transformer_block: TransformerBlock = TransformerBlock( d_model=d_model, num_heads=num_heads, intermediate_size=intermediate_size, dropout_rate=dropout_rate, layer_id=layer_id, config=config, ) # Given: 모델의 테스트용 더미 입력(dummy input)을 생성합니다. # tf.random.uniform 함수를 사용하여, 지정한 shape인 (batch_size, seq_len, d_model)에 맞게 균등 분포의 무작위 값들을 가진 텐서를 만듭니다. # 여기서 batch_size는 한 번에 처리할 데이터 개수, seq_len은 문장(또는 토큰 시퀀스)의 길이, d_model은 각 토큰의 임베딩 차원을 의미합니다. dummy_input = tf.random.uniform((batch_size, seq_len, d_model), dtype=tf.float32) # Given: 단순 attention mask를 생성합니다. # 이 부분에서는 모든 토큰이 유효하다는 가정 하에 attention mask를 만듭니다. # tf.ones를 사용하여, shape이 (batch_size, seq_len)인 2차원 텐서를 생성합니다. # 각 요소의 값은 1로 채워지는데, 이는 해당 위치의 토큰이 “유효(valid)”하다는 의미입니다. # 즉, 이 mask는 패딩 토큰이 없어서 모든 토큰에 대해 어텐션을 수행한다는 전제입니다. attn_mask_2d = tf.ones((batch_size, seq_len), dtype=tf.float32) # 이를 [batch_size, 1, seq_len]로 확장한 뒤, [batch_size, 1, seq_len, seq_len]로 타일링합니다. # 타일링할 때는 새로 추가한 차원을 기준으로, 텐서를 반복(tile)합니다. # [1, seq_len, 1] 인자는 첫 번째 차원은 그대로, 두 번째 차원을 seq_len 만큼 복제하여 (batch_size, seq_len, seq_len)이 되도록 합니다. shaped_mask = tf.tile(tf.expand_dims(attn_mask_2d, axis=1), [1, seq_len, 1]) # 다시 한 번 차원을 확장하여 최종적으로 shape을 (batch_size, 1, seq_len, seq_len)로 만듭니다. shaped_mask = tf.expand_dims(shaped_mask, axis=1) # mask 값: 유효 위치는 0.0, 패딩은 -∞; 여기선 모두 유효하므로 0.0 # shaped_mask와 동일한 shape의 텐서를 생성하되, 모든 값을 0.0으로 채웁니다. 이는 유효한 토큰 위치에 사용됩니다. zeros_mask = tf.zeros_like(shaped_mask, dtype=tf.float32) # tf.fill을 이용하여, shaped_mask와 동일한 shape의 텐서를 생성하고 모든 값을 tf.float32.min으로 채웁니다. # tf.float32.min은 float32 타입에서 표현 가능한 매우 작은 값(실질적으로 -∞에 가까운 값)을 의미합니다. # 이는 패딩 토큰 위치에 사용되어, softmax 계산 시 해당 위치의 attention score를 0으로 만들도록 합니다. neg_inf_mask = tf.fill(tf.shape(shaped_mask), tf.float32.min) # tf.where 함수를 이용하여, shaped_mask의 각 위치가 1.0이면 zeros_mask의 해당 위치 값을 선택하고, 그렇지 않으면 neg_inf_mask의 값을 선택합니다. # 여기서는 shaped_mask가 모두 1로 채워져 있기 때문에, 최종적으로 모든 위치가 0.0이 됩니다. # 결과적으로, 이 simple_mask는 모든 토큰이 유효하여 패딩이 없음을 나타내며, 어텐션 연산 시 모든 위치에 대해 아무런 마스킹 효과가 없음을 의미합니다. simple_mask = tf.where(tf.equal(shaped_mask, 1.0), zeros_mask, neg_inf_mask) # When: TransformerBlock의 call 메서드를 호출합니다. # training=False로 호출하면 드롭아웃 등이 적용되지 않습니다. output = transformer_block(dummy_input, attention_mask=simple_mask, training=False) # Then: 출력 텐서의 shape가 입력 텐서와 동일해야 합니다. expected_shape = (batch_size, seq_len, d_model) assert ( output.shape == expected_shape ), f"Expected output shape {expected_shape}, but got {output.shape}" # Then: 함수형 검증 도구를 사용해 출력값이 모두 유한한지 확인합니다. # (Residual connection과 MLP 계산 후에도 수치 불안정성이 없어야 함) np.testing.assert_allclose( output.numpy(), output.numpy(), atol=1e-6, err_msg="Output values are not consistent or contain non-finite numbers.", )

class TransformerBlock(keras.layers.Layer): def __init__( self, d_model, num_heads, intermediate_size, dropout_rate=0.1, layer_id=0, config=None, **kwargs, ): super().__init__(**kwargs) self.config = config self.layer_id = layer_id # 1. MultiHeadAttention 레이어 생성 # d_model, num_heads, dropout_rate, layer_id를 전달하여 어텐션 연산에 필요한 레이어를 구성합니다. self.attention = MultiHeadAttention(d_model, num_heads, dropout_rate, layer_id) # 2. 로타리 임베딩 파라미터 설정 (글로벌/로컬 구분) # 기본적으로 글로벌 어텐션의 rope_theta와 max_position_embeddings를 사용 self.rope_theta = self.config.global_rope_theta self.max_position_embeddings = self.config.max_position_embeddings """ 기본적으로 config에 정의된 글로벌 로타리 임베딩 값(global_rope_theta)와 최대 포지션 임베딩 길이를 사용합니다. 만약 현재 레이어가 글로벌 어텐션이 아닌 로컬 어텐션을 사용할 경우, 필요한 경우 local_rope_theta를 적용하고 최대 시퀀스 길이를 로컬 어텐션 윈도우 크기로 변경합니다. """ # 만약 현재 layer_id가 글로벌 어텐션이 적용되지 않는(즉, 로컬 어텐션인) 층이면 if (self.layer_id % self.config.global_attn_every_n_layers) != 0: # 만약 로컬 전용 rope_theta가 지정되어 있다면 이를 사용 if self.config.local_rope_theta is not None: self.rope_theta = self.config.local_rope_theta # 로컬 어텐션일 경우 최대 position은 config에서 지정한 local_attention 값으로 사용 self.max_position_embeddings = self.config.local_attention # 3. NTKScalingRotaryEmbedding 레이어 생성 """ 이 레이어는 입력에 대해 cos, sin 값을 미리 계산하여 로타리 임베딩을 제공합니다. 각 헤드마다 차원을 d_model // num_heads로 맞춰주며, 앞서 설정한 최대 포지션과 rope_theta 값을 사용합니다. """ self.rotary_emb = NTKScalingRotaryEmbedding( dim=int(d_model // num_heads), max_position_embeddings=self.max_position_embeddings, base=self.rope_theta, name=f"rotary_embeddings_{layer_id}", ) # 4. 어텐션 드롭아웃 및 MLP 정규화 # 어텐션 출력에 적용할 Dropout 레이어 self.attention_dropout = keras.layers.Dropout(dropout_rate) # MLP에 들어가기 전 residual connection 이후 정규화를 위한 LayerNormalization self.mlp_norm = keras.layers.LayerNormalization(epsilon=1e-5, center=False) # 5. MLP 구성: Wi (GLU를 위해 2배 차원)와 Wo (최종 출력 투영) """ Transformer의 MLP 부분을 Dense 레이어를 사용해 구성합니다. 여기서는 GLU(Gated Linear Unit) 방식을 사용하기 위해, 두 배 크기의 출력 차원을 갖는 Dense 레이어를 먼저 생성한 후 최종 출력 Dense 레이어를 구성합니다. Wi는 입력을 2배 차원으로 변환하여, 이후 GLU에서 두 부분(x_in과 gate)으로 분할할 수 있도록 합니다. Wo는 GLU 처리 후 나온 결과를 원래 모델 차원(d_model)으로 다시 투영합니다. """ # Wi: 중간 차원의 2배 (GLU를 위해 입력과 게이트를 분리하기 위함) self.Wi = keras.layers.Dense( intermediate_size * 2, name="intermediate.dense", use_bias=False ) # Wo: 최종 출력은 d_model 차원으로 투영 self.Wo = keras.layers.Dense(d_model, name="output.dense", use_bias=False) # MLP 출력 후 적용할 Dropout 레이어 self.output_dropout = keras.layers.Dropout(dropout_rate) def gelu_approx(self, x): return tf.nn.gelu(x) def call(self, hidden_states, attention_mask=None, training=False): batch_size = tf.shape(hidden_states)[0] seq_len = tf.shape(hidden_states)[1] # 1. 로타리 임베딩 계산 """ NTKScalingRotaryEmbedding 레이어를 호출하여, 현재 시퀀스 길이에 맞는 cos, sin 텐서를 얻습니다. 이 값들은 이후 MultiHeadAttention에 전달되어, query와 key에 로타리 임베딩을 적용합니다. """ # RoPE: rotary_emb 레이어를 사용해 cos, sin 값을 계산합니다. cos, sin = self.rotary_emb(hidden_states, seq_len=seq_len) # 계산된 cos, sin을 tuple 형태로 저장합니다. rope_embeds = (cos, sin) # 2. 로컬 vs 글로벌 어텐션 판단 # 만약 현재 레이어가 글로벌 어텐션이 적용되지 않는(로컬 어텐션) 경우, window_size를 설정합니다. window_size = None if (self.layer_id % self.config.global_attn_every_n_layers) != 0: window_size = self.config.local_attention # 예: 128 # (2) 어텐션 마스크 처리: # 만약 window_size가 지정되면, 기존의 글로벌 마스크를 로컬 슬라이딩 윈도우 마스크로 변환합니다. # 그렇지 않으면, 글로벌 마스크(모든 토큰 attending)를 그대로 사용합니다. if window_size is not None and window_size > 0: attention_mask = create_local_sliding_window_mask( attention_mask, window_size ) """ MultiHeadAttention 레이어를 호출하여, 어텐션 결과를 얻습니다. 입력(hidden_states)과 어텐션 결과를 더해 residual connection을 구성합니다. """ # 4. MultiHeadAttention 수행 및 Residual 연결 attn_output = self.attention( hidden_states, mask=attention_mask, rope_embeds=rope_embeds, training=training, ) # Residual connection: 입력에 어텐션 출력(attn_output)을 더해줍니다. hidden_states = hidden_states + attn_output # 5. MLP 단계: LayerNormalization → Wi Dense → GLU → dropout → Wo Dense → Residual 연결 normed = self.mlp_norm(hidden_states) # Wi Dense 레이어를 통해 2배 차원의 출력을 생성합니다. mlp_out = self.Wi(normed) # GLU: 출력 mlp_out를 두 부분(x_in, gate)으로 분할합니다. x_in, gate = tf.split(mlp_out, 2, axis=-1) # 활성화 함수 GELU를 x_in에 적용합니다. x_in = tf.nn.gelu(x_in) # x_in과 gate의 element-wise 곱을 통해 GLU 결과를 얻습니다. mlp_out = x_in * gate # training일 경우 dropout을 적용합니다. if training: mlp_out = self.output_dropout(mlp_out, training=training) # Wo Dense 레이어를 통해 최종 출력 차원 d_model으로 투영합니다. mlp_out = self.Wo(mlp_out) # Residual connection: 어텐션 후 결과에 MLP 출력을 더합니다. hidden_states = hidden_states + mlp_out return hidden_states

uv run ModernGTETFWeightConverter.py uv run model_conversion_validator.py Encoder Layer 22: PyTorch shape: torch.Size([2, 842, 768]) dims: [batch_size=2, seq_len=842, hidden_dim=768] TensorFlow shape: (2, 842, 768) dims: [batch_size=2, seq_len=842, hidden_dim=768] -> MSE: 32.437252 -> CLS Token Cosine Similarity: 1.000000 [TensorFlow] Final Embeddings Shape: (2, 768) === 3) PT vs. TF 최종 임베딩 비교 === [[ 0.22076873 -1.3506409 -2.0690584 ... -0.38223675 1.0725555 -0.5355745 ] [ 0.12727773 -1.3734885 -2.333182 ... -0.3719053 1.4427259 -0.78435 ]] [[ 0.2207656 -1.350643 -2.0690598 ... -0.38224083 1.0725522 -0.5355765 ] [ 0.12727615 -1.3734777 -2.3331754 ... -0.37190926 1.442732 -0.7843542 ]] ===== Queries ===== [0] 이 모델은 무엇을 하는 모델인가요? [1] 이 모델은 무엇을 하는 모델인가요?이 모델은 무엇을 ... ===== PyTorch Embeddings (shape) ===== (2, 768) ===== TF Embeddings (shape) ===== (2, 768) ===== Pairwise Cosine Similarity (PT vs TF) ===== Query 0 Cosine Similarity: 1.0000 Query 1 Cosine Similarity: 1.0000 ===== MSE (PT vs TF) ===== MSE: 0.000000 ===== Sample Differences (first query, first 5 dims) ===== [ 3.1292439e-06 2.1457672e-06 1.4305115e-06 -1.4305115e-06 2.6226044e-06]