ruGPT3XL 컨텍스트를 8k 토큰으로 확장: 희소 어텐션과 파인튜닝
개발자들은 Megatron-LM의 ruGPT3XL을 llama.cpp의 GGUF 지원을 통해 HuggingFace 형식으로 성공적으로 변환했습니다. 테스트 결과, gazeta 데이터셋에서 기대했던 12.05 대신 50.1로 퍼플렉서티(PPL)가 악화되었습니다. 원인은 희소 어텐션이 GPT-2의 조밀한 nn.MultiheadAttention으로 대체된 것이었습니다.
Megatron-LM의 원래 메커니즘을 복원하자 PPL이 11.68로 감소했습니다. 원본 모델의 주장된 지표와의 상관관계는 R=0.93에 도달했습니다.
gazeta에서의 PPL 비교
| 모델 | PPL (gazeta) |
|----------------|--------------|
| ruGPT3Small | — |
| ruGPT3Medium | 25.2 |
| ruGPT3Large | 16.8 |
| ruGPT3XL (조밀)| 50.1 |
| ruGPT3XL (희소)| 11.68 |
ruGPT3XL의 희소 어텐션 메커니즘
희소 어텐션은 논문 "Generating Long Sequences with Sparse Transformers"(arXiv:1904.10509)의 교대 패턴을 사용합니다:
- 짝수 레이어(0,2,4…): 블록 희소 어텐션. 각 토큰은 128 토큰의 로컬 윈도우와 정기적인 간격의 글로벌 블록을 참조합니다. 다른 헤드는 서로 다른 글로벌 위치를 사용합니다.
- 홀수 레이어(1,3,5…): 표준 인과적 조밀 어텐션.
이를 통해 2차 대신 거의 선형적인 메모리 증가를 보장합니다. 컨텍스트를 4배 증가시키면 KV 캐시와 활성화에 대한 메모리가 3-4배 증가합니다.
llama.cpp에서는 LLM_ARCH_RUGPT3XL 아키텍처(PR #21161)가 완전한 희소 어텐션 지원과 함께 추가되었습니다. 배치 크기가 1보다 큰 훈련 중 마스크 오류가 수정되었습니다.
Triton과 SDPA를 통한 최적화
matmul + softmax + matmul(이거 모드)로 구현하면 RTX 4090에서 약 6280 토큰/초를 달성합니다.
속도 향상:
F.scaled_dot_product_attention(SDPA): +40%(약 8800 토큰/초).- 배치 크기 증가(5×2048)와 함께 SDPA: 추가 +25%.
- Inductor와 함께
torch.compile: 기준 대비 ×1.85(약 11600 토큰/초), Triton 커널을 생성합니다.
컨텍스트를 8k로 확장하는 전략
아키텍처적 제한:
- 학습된 절대 위치 임베딩(nn.Embedding(2048,2048))은 외삽되지 않습니다.
- 희소 어텐션 그리드는
max_position_embeddings // sparse_block_size에 의존합니다.
파인튜닝 원칙
- 위치 임베딩 타일링: 새로운 위치에 대해 0-2047을 순환적으로 복제합니다. 짧은 컨텍스트에서 안정성을 유지합니다.
- 혼합 데이터셋: 60% 긴 예제(EOS를 통한 gazeta), 40% 짧은 청크.
- 단계적 확장: RTX 4090에서 2k→4k(2.6시간), 그 후 4k→8k(3.9시간).
매개변수: gazeta 훈련, 3 에폭, lr=5e-6 코사인 감쇠, bfloat16, 그래디언트 체크포인팅. PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True(38.5 GB 피크)로 OOM 해결.
모델: evilfreelancer/ruGPT3XL-8k.
gazeta 테스트에서의 PPL 결과
| 모델 | 2048 | 4096 | 8192 |
|----------------|------|------|------|
| ruGPT3XL 베이스 | 11.68| — | — |
| ruGPT3XL-4k | 11.75| 12.04| — |
| ruGPT3XL-8k | 11.77| 11.99| 13.00 |
2k에서 최소한의 퇴보, 모델은 품질을 유지합니다.
핵심 요약
- 희소 어텐션은 중요합니다: 조밀 구현은 11.68 대신 PPL 50.1을 초래합니다.
- APE 타일링은 효과적입니다: 기본 컨텍스트를 저하시키지 않고 파인튜닝을 가능하게 합니다.
- 8k는 RTX 4090에서 가능합니다: expandable_segments로 38.5 GB.
- 단계적 접근은 과적합을 최소화합니다: 2k→4k→8k를 6.5시간 내에.
- Triton/SDPA는 1.85배 속도 향상: torch.compile이 최적화된 커널을 생성합니다.
— Editorial Team
아직 댓글이 없습니다.