paged attention vs flash attention

HanJu Han·2024년 10월 29일

LLM 최적화

목록 보기
11/16
  1. 핵심 목적과 접근 방식:
  • 플래시 어텐션:
    • GPU 메모리 대역폭 최적화에 중점
    • 타일 단위로 행렬 곱셈을 수행
    • HBM과 SRAM 사이의 메모리 이동 최소화
  • 페이지 어텐션:
    • GPU 메모리 용량 관리에 중점
    • 페이지 단위로 KV 캐시 관리
    • 메모리 할당과 재사용에 초점
  1. 메모리 접근 방식:
플래시 어텐션:
- 블록 단위로 나누어 SRAM에서 계산
- 중간 결과를 HBM에 저장하지 않음
- IO 복잡도: O(N²/L) (L: SRAM 크기)

페이지 어텐션:
- 페이지 단위로 KV 캐시 관리
- 필요한 페이지만 메모리에 유지
- IO 복잡도: O(N × P) (P: 페이지 크기)
  1. 최적화 대상:
플래시 어텐션:
- 메모리 대역폭 사용량
- GPU SRAM 활용
- 행렬 연산 속도

페이지 어텐션:
- 메모리 용량 사용량
- KV 캐시 관리
- 추론 시 메모리 효율성
  1. 실제 구현 차이:
graph TB
    subgraph "플래시 어텐션"
        F1[입력 시퀀스] --> F2[블록 분할]
        F2 --> F3[SRAM 로드]
        F3 --> F4[블록 단위 계산]
        F4 --> F5[부분 결과 누적]
        F5 --> F6[최종 결과]
    end

    subgraph "페이지 어텐션"
        P1[입력 시퀀스] --> P2[페이지 분할]
        P2 --> P3[KV 캐시 저장]
        P3 --> P4[페이지별 어텐션]
        P4 --> P5[캐시 관리]
        P5 --> P6[최종 결과]
    end
  1. 성능 특성:
플래시 어텐션:
- 훈련 시 더 효율적
- 더 빠른 행렬 연산
- 메모리 대역폭 제한 해결

페이지 어텐션:
- 추론 시 더 효율적
- 더 적은 메모리 사용
- 긴 시퀀스 처리에 유리
  1. 사용 사례별 차이:
플래시 어텐션 유리:
- 모델 학습
- 대규모 배치 처리
- 행렬 연산이 많은 경우

페이지 어텐션 유리:
- 실시간 추론
- 긴 컨텍스트 처리
- 제한된 메모리 환경
  1. 조합 가능성:
  • 두 방식은 상호 배타적이지 않음
  • 함께 사용하면 더 나은 성능 가능:
    • 페이지 어텐션으로 메모리 관리
    • 플래시 어텐션으로 연산 최적화
  1. 실제 적용 예:
플래시 어텐션:
- GPT 모델 학습
- BERT 파인튜닝
- 대규모 변환기 모델

페이지 어텐션:
- 채팅 서비스
- 문서 처리
- 실시간 번역

핵심 차이는 플래시 어텐션이 연산 최적화에 중점을 두는 반면, 페이지 어텐션은 메모리 관리에 초점을 맞춘다는 것입니다. 실제 응용에서는 두 방식의 장점을 조합하여 사용하는 것이 이상적일 수 있습니다.


파인튜닝에서는 일반적으로 KV 캐시를 사용하지 않습니다. KV 캐시는 주로 추론(inference) 단계에서 사용되는 최적화 기법

  1. 파인튜닝 vs 추론 시의 차이:

파인튜닝:

  • 매 배치마다 가중치가 업데이트됨
  • 모든 토큰에 대해 새로운 어텐션 계산 필요
  • 역전파(backpropagation)가 필요
  • 그래디언트 계산이 필요

추론:

  • 가중치가 고정되어 있음
  • 이전 계산 결과를 재사용 가능
  • 역전파가 필요 없음
  • 순전파(forward pass)만 수행
  1. KV 캐시를 파인튜닝에서 사용하지 않는 이유:
a) 가중치 업데이트 문제:
- 파인튜닝: W₁ → W₂ → W₃ (가중치 변화)
- KV 캐시: K₁V₁(W₁) ≠ K₁V₁(W₂) (캐시 무효화)

b) 배치 처리 특성:
파인튜닝: 
    배치 1: [문장1, 문장2, 문장3]
    배치 2: [문장4, 문장5, 문장6]
    → 매 배치마다 새로운 계산 필요

추론:
    입력: "안녕하세요"
    다음 토큰: "반갑" → KV 캐시 활용 가능
  1. 메모리 사용 패턴:
파인튜닝 시:
- 그래디언트 저장 공간 필요
- 중간 활성화값 저장 필요
- 배치 단위 병렬 처리

추론 시:
- 순차적 토큰 생성
- 이전 계산 결과 재사용
- KV 캐시로 메모리 최적화
  1. 연산 흐름 비교:
sequenceDiagram
    participant F as 파인튜닝
    participant I as 추론
    
    rect rgb(200, 230, 200)
    Note over F: 배치 입력
    F->>F: 어텐션 계산
    F->>F: 순전파
    F->>F: 손실 계산
    F->>F: 역전파
    F->>F: 가중치 업데이트
    end
    
    rect rgb(230, 200, 200)
    Note over I: 토큰 입력
    I->>I: KV 캐시 확인
    I->>I: 필요한 경우만 계산
    I->>I: 다음 토큰 생성
    I->>I: KV 캐시 업데이트
    end
  1. 실제 예시:
# 파인튜닝 시:
for batch in training_data:
    outputs = model(batch)
    loss = criterion(outputs, targets)
    loss.backward()
    optimizer.step()
    # KV 캐시 사용 불가

# 추론 시:
kv_cache = None
for i in range(max_length):
    output, kv_cache = model(input_ids, kv_cache)
    # KV 캐시 활용 가능
  1. 대안적 최적화 기법(파인튜닝 시):
1) 그래디언트 체크포인팅
2) 혼합 정밀도 훈련
3) 배치 크기 최적화
4) 메모리 효율적인 옵티마이저
  1. KV 캐시가 필요한 상황:
적합한 경우:
- 챗봇 응답 생성
- 텍스트 자동 완성
- 코드 생성

부적합한 경우:
- 모델 훈련
- 파인튜닝
- 배치 처리

결론적으로, KV 캐시는 추론 시에만 사용되는 최적화 기법이며, 파인튜닝 과정에서는 가중치 업데이트와 역전파 과정 때문에 사용할 수 없습니다. 대신 파인튜닝에서는 다른 최적화 기법들을 활용하여 효율성을 높입니다.


페이지드 어텐션과 플래시 어텐션의 학습과 추론 시의 활용

페이지드 어텐션

페이지드 어텐션은 원래 추론 최적화를 위해 설계되었습니다. 추론 과정에서는 이전 토큰들의 KV 값을 재사용할 수 있다는 점이 큰 장점이 됩니다. 이는 마치 교과서를 읽을 때 중요한 내용을 형광펜으로 표시해두고 나중에 다시 참조하는 것과 비슷합니다.

학습 시에는 상황이 조금 다릅니다. 학습 과정에서는 역전파(backpropagation)를 위해 모든 중간 계산 값들을 저장해야 하기 때문에 KV 캐시의 이점이 크게 줄어듭니다. 마치 시험을 볼 때는 모든 풀이 과정을 자세히 적어야 하는 것과 비슷합니다.

플래시 어텐션

플래시 어텐션은 학습과 추론 모두에서 효과적으로 사용할 수 있습니다. 이는 플래시 어텐션이 메모리 접근 자체를 최적화하는 방식으로 설계되었기 때문입니다.

학습 시:

1. 순전파(forward pass) 동안:
   - 블록 단위로 계산을 나누어 처리
   - 중간 결과를 효율적으로 저장

2. 역전파(backward pass) 동안:
   - 저장된 중간 결과를 사용해 기울기 계산
   - 블록 단위로 병렬 처리 가능

추론 시:

1. 입력 시퀀스를 블록으로 나눔
2. 각 블록을 병렬로 처리
3. 결과를 효율적으로 통합

실제 사용 예시

긴 문서를 처리하는 AI 모델을 예로 들어보겠습니다:

학습 단계에서:

  • 플래시 어텐션: 전체 학습 과정에서 안정적인 성능 발휘
  • 페이지드 어텐션: 메모리 사용량은 줄일 수 있지만, 학습 속도 향상은 제한적

추론 단계에서:

  • 플래시 어텐션: 병렬 처리로 빠른 응답 가능
  • 페이지드 어텐션: KV 캐시를 활용해 효율적인 순차 처리 가능

이러한 특성 때문에 실제 현장에서는:
1. 학습 시에는 주로 플래시 어텐션을 사용하고
2. 추론 시에는 두 기술을 상황에 맞게 선택하거나 함께 사용하는 경우가 많습니다

이는 마치 학교에서 공부할 때는 모든 내용을 자세히 기록하고(플래시 어텐션), 실제 시험을 볼 때는 효율적인 방법을 선택적으로 사용하는 것(페이지드 어텐션 또는 플래시 어텐션)과 비슷합니다.

profile
시리즈를 기반으로 작성하였습니다.

0개의 댓글