[Paper Review] TinyLlama: An Open-Source Small Language Model 리뷰

TinyLlama: An Open-Source Small Language Model

TinyLlama

• 기존 연구들
  • Chinchilla’s scaling laws
    • for compute-optimal training, the model size and the number of training tokens should be scaled equally
    • Chinchilla: Gopher (280B)와 똑같은 compute budget을 사용. but 70B param & 4배 더 많은 data
    • 그동안의 LLM 모델들은 underfitting이 되어있었다. model size가 크다고 장땡이 아니라 학습하는 token도 비례하게 커져야한다. Chinchilla는 70B로 gopher보다 모델의 크기가 작지만, 더 많은 token으로 학습을 해서 성능이 더 좋았음
  • LLaMA-1
    • training budget이 아닌 inference budget에 주목
    • 작은 모델이라도 (7B ~ 65B) 오래 학습 시키면 큰 모델보다 성능이 좋다.
  • Chinchilla’s Death
    • Chinchilla’s scaling laws는 모델의 사이즈가 클수록 성능이 좋을 것 이라는 강력한 가정 하에, 모델에 적합한 최적의 학습 토큰 수를 찾으려 했음.
    • 그러나 모델 사이즈보다 어떻게 학습할 것인가 (lr 세팅)와 얼마나 학습할 것인가 가 더 중요하다
    • 모델이 작다고 학습도 조금만 시키는게 아니라 작은 모델도 “오래” “잘” 학습시키면 큰 모델만큼의 성능이 나올 수 있다고 주장
• TinyLlama
  • 작은 모델을 많은 토큰수로 학습시켜보자 !
  • 1.1B의 Transformer decoder-only model를 3T tokens으로 학습
    • pretrained on 1T tokens, 3 epochs
    • architecture & tokenizer: Llama2 와 동일

2. Pretraining

2.1 Pre-training data

1. SlimPajama
   • natural language data
   • SlimPajama – the largest extensively deduplicated, multi-corpora, open-source dataset for training large language models. SlimPajama was created by cleaning and deduplicating the 1.21T token RedPajama dataset from Together. By filtering out low quality data and duplicates, we were able to remove 49.6% of bytes, slimming down the dataset from 1210B to 627B tokens
     • RedPajama: Llama의 pretraining data를 reproduce한 open source data

1. Starcoderdata
   • code data
   • It comprises approximately 250 billion tokens across 86 programming languages. In addition to code, it also includes GitHub issues and text-code pairs that involve natural languages. To avoid data duplication, we remove the GitHub subset of the SlimPajama and only sample code data from the Starcoderdata.

TinyLlaMA

• 위 두 데이터 corpora 를 사용(총 950B tokens)
  • natural data: code data을 7:3으로 샘플링해서 학습
• 3 epoch 학습
  • Scaling Data-Constrained Language Models 에 의하면, 같은 데이터를 4 epoch 까지 repeat해서 학습하는건 새로운 unique data를 추가하는 것과 동일한 효과를 내지만, 그 이상의 epoch을 학습하면 compute resource만 쓰고 별로 효과가 없다고 주장

2.2 Architecture

Llama2와 model archi는 비슷

1. Pre-normalization (GPT3)
   • 학습 안정성을 위해 output 대신 transformer의 sub-layer의 input들을 normalization
   • RMSNorm을 사용
2. SwiGLU activation function (PaLM)
   • 성능 향상을 위해 ReLU 대신 SwiGLU를 사용
   • PaLM 처럼 4d 대신 2/3 4d dimension을 사용했다함
3. Rotary Embeddings (GPTNeo)
   • positional embedding을 제거하고 RoPE (Rotary positional Embedding) 도입
4. Grouped-query Attention
   • To reduce memory bandwidth overhead and speed up inference, we use grouped-query attention (Ainslie et al., 2023) in our model. We have 32 heads for query attention and use 4 groups of key-value heads. With this technique, the model can share key and value representations across multiple heads without sacrificing much performance.

2.3 Speed Optimizations

Llama1도 xFormer 쓰는데.. 380 token/sec/GPU 였었던걸로 알고 있는데.. 엄청 빠르다

2.4 Training

We build our framework based on lit-gpt.3 In adhering to Llama 2 (Touvron et al., 2023b), we employ an autoregressive language modeling objective during the pretraining phase. Consistent with Llama 2’s settings, we utilize the AdamW optimizer (Loshchilov and Hutter, 2019), setting β1 at 0.9 and β2 at 0.95. Additionally, we use a cosine learning rate schedule with maximum learning rate as 4.0 × 10−4 and minimum learning rate as 4.0 × 10−5 . We use 2,000 warmup steps to facilitate optimized learning.4 We set the batch size as 2M tokens. We assign weight decay as 0.1, and use a gradient clipping threshold of 1.0 to regulate the gradient value. We pretrain TinyLlama with 16 A100-40G GPUs in our project.

3. Results