[논문 리뷰] TinyLlama: An Open-Source Small Language Model

Peiyuan Zhang, Guangtao Zeng외 2명|arXiv (Cornell University)|2024. 01. 04.

Natural Language Processing Techniques인용 수 57

한 줄 요약

TinyLlama는 ~3조 토큰(≈3 에폭)으로 미리 학습된 컴팩트한 1.1B 디코더-전용 언어 모델로, Llama 2 아키텍처와 오픈 소스 속도 향상을 사용하여 비슷한 크기의 오픈 소스 모델들 중에서 강한 성능을 달성합니다.

ABSTRACT

We present TinyLlama, a compact 1.1B language model pretrained on around 1 trillion tokens for approximately 3 epochs. Building on the architecture and tokenizer of Llama 2, TinyLlama leverages various advances contributed by the open-source community (e.g., FlashAttention and Lit-GPT), achieving better computational efficiency. Despite its relatively small size, TinyLlama demonstrates remarkable performance in a series of downstream tasks. It significantly outperforms existing open-source language models with comparable sizes. Our model checkpoints and code are publicly available on GitHub at https://github.com/jzhang38/TinyLlama.

연구 동기 및 목표

• 매우 대규모 프리트레이닝 데이터로 작은 1.1B 모델이 강한 성능을 달성할 수 있는지 평가한다.
• 오픈 소스 효율성을 활용하여 학습 속도와 메모리 사용량을 개선한다.
• 유사 크기의 기존 오픈 소스 모델과 상식 및 추론 과제에서 TinyLlama를 비교한다.
• 데이터, 코드, 체크포인트를 공개하여 개방성과 재현성을 보여준다.

제안 방법

• Llama 2 스타일의 디코더 아키텍처 및 토크나이저를 사용한다.
• SlimPajama 자연어 데이터와 Starcoderdata 코드 데이터의 혼합물로 프리트레이닝한다 (~950B 토큰).
• 총 약 3조 토큰으로 약 3 에폭 동안 학습한다.
• 속도/효율 최적화를 적용: Fully Sharded Data Parallel (FSDP), FlashAttention, xFormers 조정, 그리고 그룹화된 쿼리 어텐션.
• RoPE 위치 임베딩과 SwiGLU 활성화를 갖는 RMSNorm를 사용한다.
• 프리트레이닝은 AdamW, 코사인 lr 스케줄, 워밍업, 그리고 2,000 워밍업 스텝이 있는 자기회귀 LM 목적에 따르는 것이다.

실험 결과

연구 질문

• RQ1 비정상적으로 큰 데이터셋(~3T 토큰)으로 학습했을 때 1.1B 매개변수 모델이 경쟁력 있는 성능을 달성할 수 있는가?
• RQ2 오픈 소스 효율성 개선이 성능 저하 없이 더 빠른 학습과 더 적은 메모리 사용을 가능하게 하는가?
• RQ3 상식 추론 및 문제 해결 벤치마크에서 TinyLlama가 다른 1B 규모의 오픈 소스 모델과 어떻게 비교되는가?

주요 결과

• TinyLlama는 제로샷 평가에서 여러 다운스트림 태스크에서 OPT-1.3B 및 Pythia-1.4B를 크게 능가한다.
• 유사 규모의 오픈 소스 LM들 사이에서 상식 추론 벤치마크(HellaSwag, OpenBookQA, WinoGrande, ARC, BoolQ, PIQA 등)에서 경쟁력 있는 성능을 달성한다.
• 약 3T 토큰으로 학습하고 효율 최적화를 적용하면 높은 처리량(≈24,000 토큰/초 per A100-40G)을 얻고, 비교 가능한 모델보다 더 적은 GPU-시간이 필요하다.
• TinyLlama는 InstructEval 과제들(MMLU, BBH, HumanEval, DROP)에서 확인된 baselines보다 더 강한 문제 해결 능력을 보여준다.
• 모델은 여전히 오픈 소스이며, 프리트레이닝 코드, 중간 체크포인트 및 데이터 처리 세부 정보가 공개된다.