DOCSEARCHbyAlgolia
Submit Docs

AirLLM

repository·main·Indexed 12 days ago

https://github.com/lyogavin/airllm

A library that reduces VRAM requirements for large language model inference by decomposing models into layers and streaming them. It enables massive models, such as Llama 3.1 405B or DeepSeek-V3, to run on consumer-grade GPUs (as low as 4GB) without requiring traditional quantization or pruning. Supports popular model families including Llama, Qwen, DeepSeek, Mistral, and others.

Tokens
11.2K
Snippets
41
Records
50
Agent score
98%

What's inside AirLLM

  1. Supported Models in AirLLM

    main

    AirLLM supports virtually every popular open LLM by passing its Hugging Face ID to AutoModel.from_pretrained(...).

    Supported model families include:

    • Llama (2 / 3 / 3.1 / 3.3 / 4)
    • Qwen (1 / 2 / 2.5 / 3, including MoE and FP8)
    • DeepSeek (V2 / V3 / R1)
    • Mistral & Mixtral
    • Phi
    • Gemma
    • ChatGLM
    • Baichuan
    • InternLM
    • Yi

  2. How AirLLM manages VRAM for large models

    main

    AirLLM enables running massive models on low-end hardware by only keeping one layer on the GPU at a time. Consequently, the required VRAM depends on the model's layer size rather than the total model size.

    Typical VRAM requirements:

    • ~1–2 GB: Qwen3 / Mistral / Phi (≈8B)
    • ~1–3 GB: Qwen3-30B / Mixtral (MoE)
    • ~3 GB: Qwen3-235B (MoE)
    • ~4 GB: Llama 3.x 70B (full precision)
    • ~8 GB: Llama 3.1 405B
    • ~12 GB: DeepSeek-V3 (671B)

  3. Concept: DPO vs PPO for RLHF

    main

    Direct Preference Optimization (DPO) is presented as a highly efficient alternative to the traditional PPO (Proximal Policy Optimization) method for RLHF.

    Key differences and advantages of DPO:

    • No Reward Model Required: Unlike PPO, which requires training an additional reward model and an SFT model (consuming significant GPU memory), DPO optimizes the policy language model directly using a mathematical transformation. This stabilizes training and reduces hardware requirements.
    • Stability: By removing the reward model, DPO avoids the instability and convergence issues often associated with reinforcement learning in PPO.
    • Efficiency: DPO significantly reduces GPU memory consumption and increases training speed because it only requires training the SFT model (which is then optimized via DPO loss).

  4. Run inference with Anima 33B

    main

    To run inference, you need to load the base model (e.g., timdettmers/guanaco-33b-merged) and then apply the Anima PEFT adapter (e.g., lyogavin/Anima33B) using the peft library.

    Ensure you have installed the required dependencies first.

    # imports
from peft import PeftModel
from transformers import GenerationConfig, LlamaForCausalLM, LlamaTokenizer
import torch

# create tokenizer
base_model = "timdettmers/guanaco-33b-merged"
tokenizer = LlamaTokenizer.from_pretrained(base_model)
	
# base model
model = LlamaForCausalLM.from_pretrained(
        base_model,
        torch_dtype=torch.float16,
        device_map="auto",
    )
    
# LORA PEFT adapters
adapter_model = "lyogavin/Anima33B"

model = PeftModel.from_pretrained(
        model,
        adapter_model,
        #torch_dtype=torch.float16,
    )
model.eval()
	
# prompt
prompt = "中国的首都是哪里？"
inputs = tokenizer(prompt, return_tensors="pt")
	
# Generate
generate_ids = model.generate(**inputs, max_new_tokens=30)
print(tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0])

  5. Reproduce Anima 33B model training

    main

    To reproduce the training of the Anima 33B model (tested on a single 80GB H100 or dual 40GB A100), follow these steps:

    1. Install the required dependencies.
    2. Navigate to the training directory and run the training script.

    Note: The training uses the guanaco-33b backbone and the guanaco_belle_merge_v1.0 dataset via QLoRA.

    # 1. install dependencies
pip install -r requirements.txt
# 2. 
cd training
./run_Amina_training.sh

  6. Install AirLLM via pip

    main

    Install the airllm package using pip to begin using the library for low-memory LLM inference.

    pip install airllm

  8. Fine-tune other models based on Anima

    main

    You can use Anima as a base to fine-tune other models. You must modify the --dataset and --dataset_format arguments in the execution script to point to your specific dataset.

    Steps:

    1. Install dependencies.
    2. Run the fine-tuning script located in the training directory.
    # 1. install dependencies
pip install -r requirements.txt
# 2. 
cd training
./run_finetune_raining_based_on_Anima.sh

  9. Train with Anima QLoRA DPO

    main

    Anima provides a low-cost implementation of Direct Preference Optimization (DPO) using the QLoRA framework. This allows for RLHF (Reinforcement Learning from Human Feedback) training on large models (e.g., 33B) using a single GPU by eliminating the need for a separate reward model.

    Prerequisites

    1. Data Preparation: Use a dataset format similar to hh-rlhf. Each entry must contain two keys: chosen (the preferred output) and rejected (the non-preferred output) for a given prompt.
    2. SFT Model: You must first train a Supervised Fine-Tuning (SFT) model. This model serves as the starting point and reference for the DPO training process to prevent excessive deviation.

    Training Steps

    Install the required dependencies and execute the training script from the rlhf directory.

    # 1. install dependencies
pip install -r requirements.txt

# 2. run DPO training
cd rlhf
./run_dpo_training.sh

  10. Enable model compression for 3x speed up

    main

    You can achieve up to 3x inference speedup using block-wise quantization-based model compression with minimal accuracy loss. This method quantizes weights to reduce disk loading bottlenecks.

    Requirements:

    1. Install bitsandbytes: pip install -U bitsandbytes
    2. Ensure airllm version is > 2.0.0: pip install -U airllm

    Usage: Pass the compression argument to from_pretrained() with either '4bit' or '8bit'.

    model = AutoModel.from_pretrained("garage-bAInd/Platypus2-70B-instruct",
                     compression='4bit' # or '8bit'
                    )

  11. Install Anima 100K dependencies

    main

    To use Anima 100K for training or inference, you must install several specific dependencies, including flash-attn and specialized subdirectories from the flash-attention repository. Ensure you set your CUDA_HOME environment variable correctly before installation.

    Note: You must update the CUDA_HOME path to match your local environment.

    # Please update the path of `CUDA_HOME`
export CUDA_HOME=/usr/local/cuda-11.8
pip install transformers==4.31.0
pip install sentencepiece
pip install ninja
pip install flash-attn --no-build-isolation
pip install git+https://github.com/HazyResearch/flash-attention.git#subdirectory=csrc/rotary
pip install git+https://github.com/HazyResearch/flash-attention.git#subdirectory=csrc/xentropy
pip install accelerate
pip install bitsandbytes
pip install evaluate
pip install git+https://github.com/huggingface/peft.git@v0.4.0
pip install wandb

  12. Reproduce Anima 33B training

    main

    To reproduce the training of the Anima 33B model, ensure you have a single 80GB H100 GPU or a multi-GPU environment with 2x A100 40GB. Follow these steps:

    1. Install the required dependencies.
    2. Navigate to the training directory and execute the training script.

    This process is tested and verified for these hardware configurations.

    # 1. 依存関係をインストール
pip install -r requirements.txt
# 2. 
cd training
./run_Amina_training.sh