# LLaMA

This document shows how to build and run a LLaMA model in XTRT-LLM on both single XPU and single node multi-XPU.

## Overview

The XTRT-LLM LLaMA example code is located in [`examples/llama`](./). There are several main files in that folder:

 * [`build.py`](./build.py) to build the engine(s) needed to run the LLaMA model,
 * [`run.py`](./run.py) to run the inference on an input text,

## Support Matrix
  * FP16
  * INT8 & INT4 Weight-Only
  * Tensor Parallel

## Usage

The XTRT-LLM LLaMA example code locates at [examples/llama](./). It takes HF weights as input, and builds the corresponding XTRT engines. The number of XTRT engines depends on the number of XPUs used to run inference.

### Build XTRT engine(s)

Need to prepare the HF LLaMA checkpoint first by following the guides here https://huggingface.co/docs/transformers/main/en/model_doc/llama.

XTRT-LLM LLaMA builds XTRT engine(s) from HF checkpoint. If no checkpoint directory is specified, XTRT-LLM will build engine(s) with dummy weights.

Normally `build.py` only requires single XPU, but if you've already got all the XPUs needed while inferencing, you could enable parallelly building to make the engine building process faster by adding `--parallel_build` argument. Please note that currently `parallel_build` feature only supports single node.

Here're some examples:

```bash
# Build a single-XPU float16 engine from HF weights.
# use_gpt_attention_plugin is necessary in LLaMA.
# It is recommend to use --use_gpt_attention_plugin for better performance

# Build the LLaMA 7B model using a single XPU and FP16.
python build.py --model_dir ./downloads/llama-7b-hf/ \
                --dtype float16 \
                --use_gpt_attention_plugin float16 \
                --output_dir ./downloads/llama-7b-hf/trt_engines/fp16/1-XPU/


# Build the LLaMA 7B model using a single XPU and apply INT8 weight-only quantization.
python build.py --model_dir ./downloads/llama-7b-hf/ \
                --dtype float16 \
                --use_gpt_attention_plugin float16 \
                --use_weight_only \
                --output_dir ./downloads/llama-7b-hf/trt_engines/weight_only/1-XPU/

# Build LLaMA 7B using 2-way tensor parallelism.
python build.py --model_dir ./downloads/llama-7b-hf/ \
                --dtype float16 \
                --use_gpt_attention_plugin float16 \
                --output_dir ./downloads/llama-7b-hf/trt_engines/fp16/2-XPU/ \
                --world_size 2 \
                --tp_size 2 \
                --parallel_build


# Build LLaMA 13B using 2-way tensor parallelism.
python build.py --model_dir ./downloads/llama13b/ \
                --dtype float16 \
                --use_gpt_attention_plugin float16 \
                --output_dir ./downloads/llama13b/trt_engines/fp16/2-XPU/ \
                --world_size 2 \
                --tp_size 2 \
                --parallel_build
```

#### LLaMA v2 Updates
The LLaMA v2 models with 7B and 13B are compatible with the LLaMA v1 implementation. The above
commands still work.


For LLaMA v2 70B, there is a restriction on tensor parallelism that the number of KV heads
must be **divisible by the number of XPUs**. For example, since the 70B model has 8 KV heads, you can run it with
2, 4 or 8 XPUs


```bash
# Build LLaMA 70B using 8-way tensor parallelism.
python build.py --model_dir ./downloads/llama2-70b/ \
                --dtype float16 \
                --use_gpt_attention_plugin float16 \
                --output_dir ./downloads/llama2-70b/trt_engines/fp16/8-XPU/ \
                --world_size 8 \
                --tp_size 8 \
                --parallel_build
```

Same instructions can be applied to fine-tuned versions of the LLaMA v2 models (e.g. 7Bf or llama-2-7b-chat).


Test with `summarize.py`: `pip install nltk rouge_score`

```bash
python summarize.py --test_trt_llm \
                    --hf_model_location ./downloads/llama-7b-hf \
                    --data_type fp16 \
                    --engine_dir ./downloads/llama-7b-hf/trt_engines/fp16/1-XPU
```

#### SmoothQuant

The smoothquant supports both LLaMA v1 and LLaMA v2. Unlike the FP16 build where the HF weights are processed and loaded into the XTRT-LLM directly, the SmoothQuant needs to load INT8 weights which should be pre-processed before building an engine.

Example:
```bash
python3 hf_llama_convert.py -i ./downloads/llama-7b-hf -o ./downloads/smooth_llama_7B/sq0.8/ -sq 0.8 --tensor-parallelism 1 --storage-type fp16
```

Note `hf_llama_convert.py` run with pytorch, and
1. `torch-cpu` has better accuracy than XPyTorch generally.
2. XPyTorch often use more than 32GB GM, thus more XPU are necessary to finish it.
3. add `-p=1` if run with XPyTorch.

We offer converted data [here](https://fsh.bcebos.com/v1/klx-llm/pretrained_models/quantization/smooth_llama_7B.tar.gz) for LLaMa-7b with sq of 0.6.

[`build.py`](./build.py) add new options for the support of INT8 inference of SmoothQuant models.

`--use_smooth_quant` is the starting point of INT8 inference. By default, it
will run the model in the _per-tensor_ mode.

`--per-token` and `--per-channel` are not supported yet.

Examples of build invocations:

```bash
# Build model for SmoothQuant in the _per_tensor_ mode.
python3 build.py --ft_model_dir=./downloads/smooth_llama_7B/sq0.8/1-XPU/ \
                 --use_smooth_quant \
                 --output_dir ./downloads/smooth_llama_7B/sq0.8/trt_engines/fp16/1-XPU/
```

Note we use `--ft_model_dir` instead of `--model_dir` and `--meta_ckpt_dir` since SmoothQuant model needs INT8 weights and various scales from the binary files.

### Run

Before running the examples, make sure set the environment variables:
```
export PYTORCH_NO_XPU_MEMORY_CACHING=0 # disable XPytorch cache XPU memory.
export XMLIR_D_XPU_L3_SIZE=0           # disable XPytorch use L3.
```
If you are runing with multiple XPUs and no L3 space, you can set `BKCL_CCIX_BUFFER_GM=1` to disable L3.


To run a XTRT-LLM LLaMA model using the engines generated by `build.py`

```bash
# With fp16 inference
python3 run.py --max_output_len=50 \
               --tokenizer_dir ./downloads/llama-7b-hf/ \
               --engine_dir=./downloads/llama-7b-hf/trt_engines/fp16/1-XPU/

# With fp16 inference, SmoothQuant
python3 run.py --max_output_len=50 \
               --tokenizer_dir ./downloads/llama-7b-hf/ \
               --engine_dir=./downloads/smooth_llama_7B/sq0.8/trt_engines/fp16/1-XPU/

```

### Summarization using the LLaMA model

```bash
# Run summarization using the LLaMA 7B model in FP16.
python summarize.py --test_trt_llm \
                    --hf_model_location ./downloads/llama-7b-hf/ \
                    --data_type fp16 \
                    --engine_dir ./downloads/llama-7b-hf/trt_engines/fp16/1-XPU/

# Run summarization using the LLaMA 7B model quantized to INT8.
python summarize.py --test_trt_llm \
                    --hf_model_location ./downloads/llama-7b-hf/ \
                    --data_type fp16 \
                    --engine_dir ./downloads/llama-7b-hf/trt_engines/weight_only/1-XPU/

# Run summarization using the LLaMA 7B model in FP16 using two XPUs.
mpirun -n 2 --allow-run-as-root \
    python summarize.py --test_trt_llm \
                        --hf_model_location ./downloads/llama-7b-hf/ \
                        --data_type fp16 \
                        --engine_dir ./downloads/llama-7b-hf/trt_engines/fp16/2-XPU/
```