We present an open-source reproduction of Meta AI's
LLaMa 2. However, with significantly reduced model sizes,
LiteLlama-460M-1T has 460M parameters trained with 1T tokens.
We train our models on part of
RedPajama dataset. We use the
GPT2Tokenizer to tokenize the text.
The model was trained with ~1T tokens (0.98T). num of tokens = stepslengthbatch_size=4996791024192=98240888832≈0.98T.
The training curve is at this
WandB project.
The experimental checkpoints can be directly loaded by
Transformers library. The following code snippet shows how to load the our experimental model and generate text with it.
1import torch
2from transformers import AutoTokenizer, AutoModelForCausalLM
3
4model_path = 'ahxt/LiteLlama-460M-1T'
5
6model = AutoModelForCausalLM.from_pretrained(model_path)
7tokenizer = AutoTokenizer.from_pretrained(model_path)
8model.eval()
9
10prompt = 'Q: What is the largest bird?\nA:'
11input_ids = tokenizer(prompt, return_tensors="pt").input_ids
12tokens = model.generate(input_ids, max_length=20)
13print( tokenizer.decode(tokens[0].tolist(), skip_special_tokens=True) )
14# Q: What is the largest bird?\nA: The largest bird is a black-headed gull.
Detailed results can be found
hereThis model was developed by
Xiaotian Han from Texas A&M University at the DATA Lab under the supervision of Prof.
Xia "Ben" Hu, and the model is released under MIT License.
