Padding-Free 및 Packing: 빠르고 효율적으로 LLM 파인튜닝 하기

Padding Token은 입력 시퀸스의 길이를 맞추기 위해 추가되는 특수 토큰이다. 트랜스포머 모델은 병렬적으로 모든 토큰을 처리하기 때문에, 입력 시퀸스의 길이가 동일해야한다. 이에  가장 많이 사용되지만, 메모리와 계산량 측면에서 상당한 낭비를 초래한다. 이에 최근에는 Packing으로 시퀸스의 활용도를 극대화하고, 계산 비용을 최소화할 수 있는 Packing이 제안되였다. 하지만, 종종 Packing을 기반으로 파인튜닝하는 게시물 혹은 영상물을 여럿 찾아볼 수 있지만, 이에 대한 성능의 효율성을 보여주는 것이 대부분이다. 구체적으로 패킹이 파인튜닝하는 과정에서 내부적으로 어떻게 작동하는지에 대해서는 설명하지 않는다. 따라서, 본 게시물에서는 Padding과 Packing를 SFT 기반으로 파인튜닝하는 것에 대해 다루고, 차이점을 구체적으로 살펴보고자 한다. 

---

Padding이 뭘까? 

Padding Token의 간단한 예시 (이미지 출처: Medium)

 

파인튜닝 시 패딩 토큰은 가장 기본적인 배치 처리 방식이다. 일반적으로 GPU와 같은 하드웨어는 병렬 연산 효율을 극대화하기 위해 동일한 크기의 행렬(입력)을 선호하지만, 학습 데이터 셋의 문장(시퀸스) 길이는 제각각 다르다. 이러한 간극을 메우기 위해, 하나의 배치에 포함된 모든 시퀸스의 길이를 동일하게 맞춰야 한다는 제약이 생긴다. 

 

일반적으로 파인튜닝을 시작하기에 앞서, 아래 코드와 같이 "max_seq_length"을 정의한다. 이는 모델이 한 번에 처리할 텍스트의 최대 길이를 1,024개의 토큰으로 제한함을 의미한다. 

#HuggingFace TRL 라이브러리 SFTConfig 예시 
training_arguments = SFTConfig(
          output_dir=output_dir,
          # ... (기타 설정 생략) ... 
          dataset_text_field="text",
          ****max_seq_length=1024****,
          report_to='none'
  )

 

 

만약 입력되는 시퀸스의 길이가 max_seq_length를 초과하게 된다면, 해당 토큰들이 짤리게(truncated) 된다. 반면에 시퀸스의 길이가 max_seq_length 보다 짧다면, 필요한 길이에 맞춰 패딩 토큰 (<PAD>)을 추가하여 길이를 맞춰야 한다. 이러한 패딩 토큰들은 모델 학습에 영향을 주지 않도록 훈련 과정에서 마스킹(masked)으로 처리된다. 

 

예를 들어 max_seq_length =8 설정하고, 토크나이저가 다음과 같이 문장을 분리한다고 가정을 해보겠습니다. 

• 문장 A: "I LIKE TO STUDY" (4 토큰) → I LIKE TO STUDY <PAD> <PAD> <PAD> <PAD>
• 문장 B: "HELLO" (1 토큰) → HELLO <PAD> <PAD> <PAD> <PAD> <PAD> <PAD> <PAD>
• 문장 C: "TODAY IS PRETTY SHINY AND LOVELY AT THE SAME TIME" (10 토큰) → TODAY IS PRETTY SHINY AND LOVELY AT THE

위의 예시처럼 문장 A와 B는 길이가 8보다 짧아 남는 공간이 <PAD> 토큰으로 채워져 있다. 반면 문장 C는 길이가 10이므로, 최대 길이에 맞춰 뒷부분인 "SAME TIME"이 잘려나갔다. 이처럼 패딩은 구현이 상대적으로 단순하지만, 예시에서 볼 수 있듯이 <PAD> 토큰이 전체 데이터의 상당 부분을 차지하게 된다. 모델은 이 <PAD> 토큰에 대해서도 어텐션 계산을 수행하므로(물론 결과는 마스킹 처리되어 무시됨), GPU의 연산 능력과 메모리를 불필요하게 낭비하게 된다. 특히 파인튜닝 데이터셋의 문장 길이가 제각각이고 짧은 문장이 많을수록 비효율적이다.

---

Packing이 뭘까? 

Padding과 Packing의 차이점 (이미지 출처: HuggingFace)

 

패킹은 패딩 토큰의 비효율성을 해결하기 위한 접근법이다. 이는 여러 개의 짧은 시퀸스를 하나의 시퀸스로 이어 붙여, 모델의 max_seq_length를 최대한 알차게 채우는 방식이다. 예를 들어, 문장 A("I LIKE TO STUDY", 4 토큰)와 문장 B("HELLO", 1 토큰)를 max_seq_length=8에 패킹하면 다음과 같다.

• 먼저 두 문장을 문장 끝을 나타내는 특수 토큰(예: <EOS>)으로 연결한다.
• 이를 하나의 시퀀스로 만들면 I LIKE TO STUDY <EOS> HELLO가 되며, 총 7개의 토큰을 차지한다.
• 최대 길이에 맞춰 패딩을 하나만 추가하면 최종 시퀀스는 I LIKE TO STUDY <EOS> HELLO <PAD> 가 된다.

이처럼 패딩만 사용했을 때는 두 문장을 처리하기 위해 총 16개의 토큰 공간이 필요했지만, 패킹 같은 경우에는 8개의 공간만으로 두 문장을 모두 처리할 수 있다. 이처럼 낭비되던 <PAD> 토큰이 크게 줄어들어 GPU 자원을 훨씬 효율적으로 사용할 수 있다. 하지만, 패킹을 제대로 활용할려면 짚고 넘어가야할 것이 있다. 단순하게 하나의 시퀸스로 합쳐진 문장들이 서로 영향을 주지 않도록 막는 것이 매우 중요하며, 이에 대한 어텐션 마스크를 직접 구현해야한다. 

 

일반적으로 언어 모델에 사용되는 Causal Attention Mask는 각 토큰이 오직 자신과 자신보다 앞에 위치한 토큰들만 참조하도록 제한하는 역할을 한다. 이는 단순하게 모델이 다음 단어를 예측할 때 미래의 정답을 미리 엿보는 것을 방지하는 역할을 수행하는 것과 같다. 하지만, 여러 문장이 하나로 합쳐진 패킹된 시퀸스에서는 이러한 단순한 어텐션 마스크는 오히려 문제를 일으킨다. 예를 들어, 'I LIKE TO STUDY <EOS> HELLO'라는 패킹된 시퀀스가 있다면, 일반적인 규칙에 따라 'HELLO' 토큰은 자신보다 앞에 있는 'STUDY'나 'LIKE' 같은, 전혀 다른 문맥의 토큰 정보를 그대로 참조하게 된다. 이처럼 서로 다른 문장의 정보가 뒤섞이는 현상을 Cross-Contamination이라고 칭한다.

 

패킹은 이러한 문제점을 방지하기 위해, 마치 문장들 사이에 "벽"을 세우는 것과 같은 block-diagonal attention mask를 구현해서 사용해야 한다. 이는 각 문장을 하나의 블록으로 간주하여, 토큰들이 오직 자신이 속한 블록 내에서만 어텐션하도록 강제하는 방법이다. 이에 대한 예시는 아래와 같다. 

Sequence: [A1 A2 A3 | B1 B2 | C1 C2 C3]
Attention Mask (Block-Diagonal)

[
  [1 1 1, 0 0, 0 0 0 0],  # A블록은 A블록 내에서만 어텐션
  [1 1 1, 0 0, 0 0 0 0],
  [1 1 1, 0 0, 0 0 0 0],

  [0 0 0, 1 1, 0 0 0 0],  # B블록은 B블록 내에서만 어텐션
  [0 0 0, 1 1, 0 0 0 0],

  [0 0 0, 0 0, 1 1 1 1],  # C블록은 C블록 내에서만 어텐션
  [0 0 0, 0 0, 1 1 1 1],
  [0 0 0, 0 0, 1 1 1 1],
  [0 0 0, 0 0, 1 1 1 1]
]

 

이와 같은 방식은 매우 효과적이지만, 기존의 어텐션 마스크 처리 방식을 크게 수정해야 하므로 TRL, Unsloth 같은 유명한 파인튜닝 라이브러리에 구현되어 있지 않다. 이에 따라 padding-free 라는 새로운 대안을 도입하였으며, 이는 패딩의 비효율성을 개선하고 패킹의 구현에 대한 복잡성을 최소화한 것이다. 

---

[cu_seqlens and FlashAttention] Packing-free가 뭘까? 

패딩-프리 배치(Padding-free batching)는 패딩의 필요성을 제거하는 데 초점을 둔다. 이 방식은 시퀸스들을 <PAD> 토큰을 추가하여 고정된 길이로 맞추는 대신, 모든 입력 시퀸스를 하나의 연속적인 텐서로 평탄화(flatten)하게 된다. 그리고 각 개별 시퀸스가 어디에서 시작하고 끝나는지를 나타내기 위해 가장 핵심적인 메타데이터인 cu_seqlens(누적 시퀸스 길이)를 사용하게 된다. 

 

이와 같은 방법은 HuggingFace의 TRL(Transformer Reinforcement Learning)과 같은 프레임워크를 FlashAttention과 결합되어 대규모 모델을 학습할 때 가장 큰 이점을 제공한다. FlashAttention는 이러한 메타데이터를 사용하여 최적화된 분할 방식으로 어텐션(attention)을 계산할 수 있다.  

 

일반적인 패딩 기반 배치 방식에서는 위에서 다루었듯이 다음과 같다. 

# 일반적인 패딩 기반 배치 방식
[
  [A1, A2, A3, PAD, PAD],
  [B1, B2, B3, B4, B5],
]

 

첫 번째 시퀀스는 두 번째 시퀀스의 길이에 맞추기 위해 패딩 토큰이 추가되었다. 하지만, 이는 패딩 토큰에 대한 불필요한 연산을 초래한다. 이러한 패딩 토큰들은 어텐션 메커니즘이나 손실(loss) 계산에 영향을 주지 않도록 학습 중에 명시적으로 마스킹(masking) 처리가 되어야 한다. 반면, 패딩-프리 배치 방식은 모든 시퀀스를 하나의 연속적인 토큰 배열로 평탄화하여 패딩 토큰의 필요성을 없애버린다.  

# 패딩-프리 기반 배치 방식
[ A1, A2, A3, B1, B2, B3, B4, B5 ]

 

예시와 같이 평탄화된 시퀀스는 개별 시퀸스의 시작 및 끝 지점을 명시적으로 표현할 수 있는 메타데이터  다음과 같이 정의한다. 

• 시퀀스 A: 인덱스 0에서 시작, 3 앞에서 끝 (길이 3) 
• 시퀀스 B: 인덱스 3에서 시작, 8 앞에서 끝 (길이 5) 

이러한 정보를 바탕으로 모델은 각 개별 시퀀스 내에서만 어텐션을 적용하게 된다. 각 시퀸스는 완벽하게 분리되어 명시적으로 구현된 어텐션 마스크 없이도 시퀸스 간의 정보유출을 방지할 수 있다. 즉, 블록-대각선(block-diagonal) 형태의 어텐션 구조가 암묵적으로 강제되는 것이다. 

 

이와 같은 블록-대각선 구조의 효율적인 계산을 가능하게 하는 핵심적인 기술은 FlashAttention이다. 일반적으로 사용되는 어텐션 메커니즘에서 발생하는 가장 큰 문제점은 GPU의 SRAM 및 HBM 간의 데이터 전송에서 발생하게 된다. 이말은 즉슨, GPU의 메모리가 연산하는 것에 사용되는 것이 아닌, 데이터를 옮기는 과정에서 많이 낭비가 된다는 의미이다. 이를 더 구체젝으로 설명하면 다음과 같다.

• 기존 어텐션 방식에서의 계산 과정에서 시퀸스 길이(N)에 비례하는 어텐션 스코어 행렬(N X N)이 생성 된다. 
• 어텐션 스코어 행렬의 크기가 너무 크기 때문에, GPU 코어 바로 옆에 붙어 있는 SRAM에 전부 다 올라가지 못하게 된다.
• 결국 이러한 어텐션 스코어 행렬을 처리하기 위해, 상대적으로 느린 HBM에 저장했다가, Softmax 계산을 위해 다시 불러와야 하는 비효율적으로 동작하게 되는 것이다. 

결국, 위와 같이 GPU를 통해 연산을 수행하는 것 보다, 데이터를 HBM으로 보내고 다시 가져오는 것을 기다리는 데 더 많은 GPU의 메모리와 시간을 소모하게 되는 것과 같다. 이에 FlashAttention에서는 이러한 명목 현상을 없애기 위해 Tiling과 Kernel Fusion 이라는 기법을 사용한다. 이때 이러한 기법들은 위에서 언급되었던 cu_seqlens 메타데이터와 좋은 시너지를 보인다. 즉, 기존 방식처럼 거대한 전체 어텐션 행렬을 계산한 뒤 마스크를 통해 0으로 만드는 것이 아니라, 필요한 블록(각 시퀀스 내부)에 대해서만 연산을 집중하는 방식이다. 

 

따라서, 패딩-프리 배치는 데이터를 효율적으로 구성하는 방법론을 제시하고, FlashAttention은 가장 효율적으로 처리하는 커널을 제공한다. 이 둘의 조합은 특히 SFT(Supervised Fine-Tuning)처럼 문장의 길이가 천차만별인 데이터셋을 다룰 때, 훈련 속도와 효율성을 극대화하여 대체적으로 많은 파인튜닝 기법에서 활용되고 있는 추세이다.