We demonstrate how to run inference (next token prediction) with the LLaMA base model in the generate.py script:
python generate.py --prompt "Hello, my name is"Output:
Hello my name is TJ. I have a passion for the outdoors, love hiking and exploring. I also enjoy traveling and learning new things. I especially enjoy long walks, good conversation and a friendly smile.
The script assumes you have downloaded and converted the weights and saved them in the ./checkpoints folder as described here.
Note All scripts support argument customization
With the default settings, this will run the 7B model and require ~26 GB of GPU memory (A100 GPU).
On GPUs with bfloat16 support, the generate.py script will automatically convert the weights and consume about ~14 GB.
For GPUs with less memory, or ones that don't support bfloat16, enable quantization (--quantize llm.int8):
python generate.py --quantize llm.int8 --prompt "Hello, my name is"This will consume about ~10 GB of GPU memory or ~8 GB if also using bfloat16.
See python generate.py --help for more options.
You can also use GPTQ-style int4 quantization, but this needs conversions of the weights first:
python quantize/gptq.py --output_path checkpoints/lit-llama/7B/llama-gptq.4bit.pth --dtype bfloat16 --quantize gptq.int4GPTQ-style int4 quantization brings GPU usage down to about ~5GB. As only the weights of the Linear layers are quantized, it is useful to also use --dtype bfloat16 even with the quantization enabled.
With the generated quantized checkpoint generation quantization then works as usual with --quantize gptq.int4 and the newly generated checkpoint file:
python generate.py --quantize gptq.int4 --checkpoint_path checkpoints/lit-llama/7B/llama-gptq.4bit.pth