RWKV 아키텍처를 Hugging Face Transformers에 통합

TL;DR

RWKV는 트랜스포머 어텐션을 모방한 새로운 RNN‑기반 아키텍처로, 이제 Hugging Face transformers 라이브러리에서 공식적으로 지원됩니다. 이를 통해 개발자는 RNN의 속도와 메모리 효율성을 유지하면서 매우 긴 컨텍스트를 처리할 수 있는 오픈‑소스 언어 모델에 접근할 수 있습니다.

RWKV 프로젝트 개요

RWKV 프로젝트는 Bo Peng(GitHub: BlinkDL)이 이끌며 활발한 Discord 커뮤니티가 유지하고 있습니다. Stability AI가 학습에 사용된 GPU를 기부했습니다. 프로젝트 로드맵에는 성능 개선(RWKV.cpp, 양자화 등), 확장성 향상(데이터셋 처리) 및 챗‑파인튜닝·멀티모달 파인튜닝과 같은 연구 확장이 포함됩니다. 커뮤니티 구성원은 공식 Discord 채널에 참여해 기여할 수 있습니다.

RWKV가 RNN과 트랜스포머를 연결하는 방식

RNN의 한계와 트랜스포머의 장점

  • 전통적인 RNN은 각 시간 단계마다 동일한 가중치를 재사용하므로, 기울기 소실 문제와 장거리 메모리 부족을 초래합니다. LSTM과 GRU가 이를 부분적으로 완화하지만 여전히 매우 긴 시퀀스에서는 어려움을 겪습니다.
  • 트랜스포머는 자체 어텐션을 통해 모든 토큰을 병렬로 처리하며, query, key, value 투영을 사용해 어텐션 점수를 계산합니다. 이 설계는 장거리 의존성 문제를 해결하고 기존 RNN에 비해 학습 속도를 크게 높입니다.

RWKV의 하이브리드 설계

  • RWKV는 Apple의 Attention‑Free Transformer에서 영감을 받아 RNN과 호환 가능한 형태로 단순화되었습니다.
  • 트랜스포머 스타일의 임베딩, layer‑norm, 인과 언어 모델 헤드를 유지하면서, 어텐션 레이어를 재귀 기반 공식으로 교체해 자체 어텐션과 동일한 표현력을 제공합니다.
  • TokenShift와 SmallInitEmb와 같은 추가 트릭(공식 GitHub README에 문서화됨)이 모델이 GPT 수준의 성능을 달성하도록 필요합니다.

RWKV 아키텍처의 기술적 하이라이트

장기 컨텍스트 처리 능력

  • RWKV는 ctx8192(8 192 토큰) 컨텍스트 윈도우를 1 024 토큰 모델과 동일한 추론 속도와 메모리 사용량으로 처리할 수 있습니다.
  • 실험적인 손실 곡선은 컨텍스트 길이가 길어질수록 모델 크기에 관계없이 언어 모델 손실이 개선됨을 보여주며, 효과적인 장거리 기억 능력을 입증합니다.

학습 효율성

  • 기존 RNN과 달리 RWKV는 “선형화된 GPT” 방식으로 학습할 수 있어 배치 간 병렬 처리가 가능하고 전통적인 순환 모델보다 빠르게 수렴합니다.
  • 현재 학습 파이프라인은 14 B 파라미터까지 확장 가능하며, RWKV‑4 시리즈의 수치 안정성을 위한 지속적인 수정이 진행 중입니다.

사용 가능한 모델 체크포인트

순수 언어 모델(RWKV‑4)

  • 모델 크기는 약 170 M부터 14 B 파라미터까지 다양합니다.
  • 모든 모델은 The Pile 데이터셋으로 사전 학습되었으며, 최신 베이스라인과 비교 테스트를 통해 비슷한 성능을 보였습니다.

명령어‑파인튜닝 챗 모델(RWKV‑4 Raven)

  • Raven 시리즈는 ALPACA, CodeAlpaca, Guanaco, GPT‑4All, ShareGPT 등 다양한 명령어 데이터셋으로 RWKV‑4를 파인튜닝합니다.
  • 영어 전용, 영어 + 중국어 + 일본어 등 다양한 언어 조합과 1.5 B, 7 B, 14 B 등 크기별 변형이 존재합니다.
  • 모든 체크포인트는 Hugging Face Hub의 RWKV 조직 아래에 호스팅됩니다.

🤗 Transformers와 함께 RWKV 사용하기

텍스트 생성 예시

from transformers import pipeline
model_id = "RWKV/rwkv-4-169m-pile"
pipe = pipeline("text-generation", model=model_id)
print(pipe("In a shocking finding, scientist discovered a herd of dragons...", max_new_tokens=20))

이 파이프라인은 트랜스포머 기반 생성기와 견줄 만한 일관된 이어쓰기를 반환합니다.

챗 모델(Raven) 예시

from transformers import AutoModelForCausalLM, AutoTokenizer
model_id = "RWKV/rwkv-raven-1b5"
model = AutoModelForCausalLM.from_pretrained(model_id).to(0)
tokenizer = AutoTokenizer.from_pretrained(model_id)
prompt = "### Instruction: Tell me about ravens\n### Response:"
inputs = tokenizer(prompt, return_tensors="pt").to(0)
output = model.generate(inputs["input_ids"], max_new_tokens=100)
print(tokenizer.decode(output[0], skip_special_tokens=True))

모델은 Alpaca 스타일의 명령어 형식을 따르며 상세한 응답을 생성합니다.

원본 RWKV 가중치를 Hugging Face 형식으로 변환하기

convert_rwkv_checkpoint_to_hf.py 변환 스크립트가 transformers 레포에 포함되어 있습니다. 사용자는 원본 체크포인트를 Hub 레포에 업로드한 뒤 다음 명령을 실행합니다.

python convert_rwkv_checkpoint_to_hf.py \
  --repo_id RAW_HUB_REPO \
  --checkpoint_file RAW_FILE \
  --output_dir OUTPUT_DIR

--push_to_hub와 --model_name 옵션을 추가하면 변환된 모델을 바로 Hub에 업로드할 수 있습니다.

향후 방향

  • 다국어 RWKV – 다국어 코퍼스와 토크나이저 작업이 진행 중이며, 모델의 언어 커버리지를 확대할 예정입니다.
  • 커뮤니티 연구 – Discord 채널에서는 새로운 학습 레시피, 벤치마크, 아키텍처 개선 프로젝트가 활발히 진행됩니다.
  • 압축 및 가속 – RWKV는 행렬‑벡터 연산만을 사용하므로 양자화(4‑bit/8‑bit), ONNX 내보내기, 광자 가속기와 같은 실험적 하드웨어에 적합합니다. optimum 라이브러리와 rwkv.cpp, rwkv-cpp-cuda와 같은 레포와의 통합은 추론 속도를 더욱 향상시킬 것입니다.

감사의 글

Hugging Face 팀은 Bo Peng, RWKV 커뮤니티 및 Johan Wind(RWKV 블로그 포스트), ArEnSc(초기 Transformers PR), Merve Noyan, Maria Khalusova, Pedro Cuenca 등 리뷰와 통합 지원에 기여한 분들께 감사를 표합니다.

인용

연구에서 RWKV를 사용할 경우, RWKV‑LM 레포지토리의 CITATION.cff 파일에 제공된 정보를 인용해 주세요.

Sources