DOCSEARCHbyAlgolia
Submit Docs

TinyLlama Documentation

repository·main·Indexed 27 days ago

https://github.com/jzhang38/tinyllama

TinyLlama is a project for pretraining a highly efficient 1.1B parameter Llama-architecture model on 3 trillion tokens. Designed for high-performance training and low-footprint inference, it is ideal for edge devices and speculative decoding. The repository provides instructions for pretraining on multiple nodes, full-parameter and QLora finetuning, and running a Gradio-based chatbot demo. It includes technical specifications for its 22-layer architecture and supports optimizations like Flash Attention 2 and FSDP.

Tokens
4K
Snippets
11
Records
23
Agent score
93%

What's inside TinyLlama

  1. Overview of TinyLlama-1.1B

    main
    TinyLlama is a project aimed at pre-training a 1.1B parameter Llama model on 3 trillion tokens. It uses the same architecture and tokenizer as Llama 2, making it compatible with many open-source Llama-based projects. Due to its small size, it is suitable for applications with limited compute and memory, such as edge devices (a 4-bit quantized version requires only ~550MB RAM), speculative decoding for larger models, and real-time dialogue in games.

  2. Run the Tinyllama Gradio Chatbot locally

    main

    Launch the chatbot interface by running the app.py script. Once the server starts, the terminal will display a local URL. Open this URL in your web browser to interact with the chatbot.

    If you are running this on a remote server, use SSH local port forwarding to access the interface: ssh -L [local port]:localhost:[remote port] [username]@[server address]

    python TinyLlama/chat_gradio/app.py

  3. Fine-tuning TinyLlama

    main

    Fine-tuning code and inference scripts are located in the sft directory.

    • Chat Models: The project has released chat models (e.g., TinyLlama-1.1B-Chat-v0.1) fine-tuned on the OpenAssistant dataset.
    • Low-Memory Fine-tuning: For GPUs with less than 4GB of RAM, it is recommended to use Qlora and bitsandbytes.

  4. Evaluate TinyLlama using Instruct-Eval Benchmarks

    main

    To evaluate TinyLlama's problem-solving capabilities (MMLU, BBH, HumanEval, and DROP) using the instruct-eval suite, run the corresponding main.py commands for each task. You can specify different GPUs using CUDA_VISIBLE_DEVICES to run evaluations in parallel.

    CUDA_VISIBLE_DEVICES=0 python main.py mmlu --model_name llama --model_path PY007/TinyLlama-1.1B-intermediate-step-480K-1T
CUDA_VISIBLE_DEVICES=1 python main.py bbh --model_name llama --model_path PY007/TinyLlama-1.1B-intermediate-step-480K-1T
CUDA_VISIBLE_DEVICES=2 python main.py drop --model_name llama --model_path PY007/TinyLlama-1.1B-intermediate-step-480K-1T
CUDA_VISIBLE_DEVICES=3 python main.py humaneval  --model_name llama  --n_sample 1 --model_path PY007/TinyLlama-1.1B-intermediate-step-480K-1T

  5. Use TinyLlama for Speculative Decoding

    main
    TinyLlama can be used to assist in the speculative decoding of larger models. Examples and instructions for using TinyLlama with llama.cpp for speculative decoding can be found in the speculative_decoding/README.md file within the repository.

  8. Finetune TinyLlama

    main

    The repository includes a simple full-parameter finetuning and inference script located in the sft directory. The V0.1 chat model was finetuned using this script with the openassistant-guanaco dataset.

    For environments with limited memory (less than 4GB RAM), it is recommended to use QLora and bitsandbytes instead of the provided full-parameter script.

  9. Use HuggingFace Assisted Generation for Speculative Decoding

    main

    You can implement speculative decoding using HuggingFace's 'Assisted Generation' feature. This involves using a small assistant model (e.g., PY007/TinyLlama-1.1B-Chat-v0.1) to speed up a larger target model (e.g., guanaco-7b to guanaco-33b).

    Note: Due to INT8 quantization and causal masking in assisted generation, the output of greedy decoding may differ from native decoding in rare occasions. Implementation code for testing this setup can be found in instruct_hf_assisted_decoding.py.

  10. Install dependencies for TinyLlama pretraining

    main

    To prepare the environment for pretraining, ensure CUDA 11.8 is installed. You must install PyTorch Nightly, build xformers from source, and install flash-attention along with its fused operators. Finally, install the remaining requirements.

    Note: Building xformers and flash-attention may take 5 minutes or longer and may appear to hang or produce many warnings.

    # Install Pytorch Nightly
pip install --index-url https://download.pytorch.org/whl/nightly/cu118 --pre 'torch>=2.1.0dev'

# Build XFormers from Source
pip uninstall ninja -y && pip install ninja -U
pip install -v -U git+https://github.com/facebookresearch/xformers.git@main#egg=xformers

# Install Flash-Attention 2 and fused operators
git clone https://github.com/Dao-AILab/flash-attention
cd flash-attention
python setup.py install
cd csrc/rotary && pip install .
cd ../layer_norm && pip install .
cd ../xentropy && pip install .
cd ../.. && rm -rf flash-attention

# Install Remaining Dependencies
pip install -r requirements.txt tokenizers sentencepiece

  11. Evaluate TinyLlama using GPT4All Benchmarks

    main

    To evaluate TinyLlama's commonsense reasoning ability following the GPT4All evaluation suite, use the lm-eval-harness tool. The evaluation typically reports acc_norm by default across tasks like HellaSwag, OpenBookQA, WinoGrande, ARC (Easy and Challenge), BoolQ, and PIQA.

    python main.py \
    --model hf-causal \
    --model_args pretrained=PY007/TinyLlama-1.1B-Chat-v0.1,dtype="float" \
    --tasks hellaswag,openbookqa,winogrande,arc_easy,arc_challenge,boolq,piqa\
    --device cuda:0 --batch_size 32

  12. Prepare datasets for pretraining

    main

    TinyLlama pretraining requires the SlimPajama and Starcoderdata datasets.

    1. Download: Use git lfs to clone the datasets from HuggingFace. Note that SlimPajama requires ~893GB and Starcoderdata requires ~290GB.
    2. Tokenize: Use the provided scripts to tokenize the datasets and divide them into chunks. The processed data will require approximately 1.8T of storage.
    # Download Datasets
cd /path/to/dataset
git lfs install
git clone https://huggingface.co/datasets/cerebras/SlimPajama-627B
git clone https://huggingface.co/datasets/bigcode/starcoderdata

# Tokenize data
python scripts/prepare_starcoder.py --source_path /path/to/starcoderdata/ --tokenizer_path data/llama --destination_path data/slim_star_combined --split train --percentage 1.0
python scripts/prepare_slimpajama.py --source_path /path/to/SlimPajama --tokenizer_path data/llama  --destination_path data/slim_star_combined --split validation --percentage 1.0
python scripts/prepare_slimpajama.py --source_path /path/to/SlimPajama --tokenizer_path data/llama  --destination_path data/slim_star_combined --split train --percentage 1.0