使用GptManager / cpp运行时运行gpt-2b + LoRA

首先构建一个启用了LoRA和inflight-batching的模型。

git-lfs clone https://huggingface.co/qychen/luotuo-lora-7b-0.1
git-lfs clone https://huggingface.co/kunishou/Japanese-Alpaca-LoRA-7b-v0
BASE_MODEL=llama-7b-hf

python examples/llama/convert_checkpoint.py --model_dir ${BASE_MODEL} \
    --output_dir /tmp/llama_7b/trt_ckpt/fp16/1-gpu/ \
    --dtype float16

trtllm-build --checkpoint_dir /tmp/llama_7b/trt_ckpt/fp16/1-gpu/ \
    --output_dir /tmp/llama_7b_with_lora_qkv/trt_engines/fp16/1-gpu/ \
    --remove_input_padding enable \
    --gpt_attention_plugin float16 \
    --context_fmha enable \
    --paged_kv_cache enable \
    --gemm_plugin float16 \
    --lora_plugin float16 \
    --max_batch_size 128 \
    --max_input_len 512 \
    --max_seq_len 562 \
    --lora_dir Japanese-Alpaca-LoRA-7b-v0 \
    --max_lora_rank 8 \
    --lora_target_modules "attn_q" "attn_k" "attn_v"

要将LoRAs传递到cpp运行时，它们必须转换为以下格式。下面的脚本将把Hugging Face的LoRA模型转换为正确的NumPy张量。

python3 tensorrt_llm/examples/hf_lora_convert.py -i Japanese-Alpaca-LoRA-7b-v0 -o Japanese-Alpaca-LoRA-7b-v0-weights --storage-type float16
python3 tensorrt_llm/examples/hf_lora_convert.py -i luotuo-lora-7b-0.1 -o luotuo-lora-7b-0.1-weights --storage-type float16

请参考tensorrtllm_backend文档以获取使用Triton的多LoRA示例。

LoRA 张量格式详情

要使用GptManager运行LoraWeights的推理，InferenceRequests必须包含LoraWeights（lora_weights）和LoraConfig（lora_config）参数。

LoraTaskId 给定LoRA的唯一任务ID。

首次使用特定LoRA进行推理时，必须提供lora_task_id、lora_weights和lora_config。LoRA将被缓存，因此后续对同一任务的请求仅需要lora_task_id。如果缓存已满，最旧的LoRA将被移除以为新的LoRA腾出空间。如果lora_task_id未被缓存，将返回错误。

LoraWeights 包含所有LoRA的权重。目前，这应该包括所有TP和PP等级的权重。权重张量的形状为 [num_lora_modules_layers, D x Hi + Ho x D ]。最后一个维度保存了相关模块（例如，attn_qkv）和模型层的输入/输出适配器权重。

每个输入/输出张量首先被展平，然后按照上述格式连接在一起。第一个维度（大小为num_lora_module_layers）为每个模块层都有一个条目（即，attn_q layer1有一个条目，attn_k layer1也有一个条目）。

D=adapter_size (即 R 值), Hi=hidden_size_in, Ho=hidden_size_out.

LoraConfig 是一个配置张量，用于标识 LoraWeights 每个元素的模块ID、层ID和适配器大小。它的形状为 [num_lora_modules_layers, 3]。最后一个维度包含 [module_id, layer_idx, adapter_size D (即 R 值)]。

此功能支持如 https://arxiv.org/pdf/2106.09685.pdf 中描述的LoRAs。

示例LoRA张量

这里是一个模型的LoraWeights和LoraConfig张量的示例，该模型的tp=1，pp=1，4层，隐藏大小为4。以下张量适用于具有q和k适配器的LoRA。

# loraConfig
[
  [1, 0, 2]
  [2, 0, 4]
  [1, 1, 2]
  [2, 1, 4]
  [1, 2, 2]  # Note that the final 2 layers only adapt `q`
  [1, 3, 8]
]
# Note: The loraConfig tensor configures the loraWeights tensor.
#       The contents of each row of loraWeights is specified be the corresponding row in loraConfig

# loraWeights
# Note: that 'in weights' and 'out weights' are 'A' and 'B' in the LoRA paper.
[
  [ <2 x 4 in weights>, <4 x 2 out weights> <padding> ]  # `q` adapter for layer 0
  [ <4 x 4 in weights>, <4 x 4 out weights> <padding> ]  # `k` adapter for layer 0
  [ <2 x 4 in weights>, <4 x 2 out weights> <padding> ]  # `q` adapter for layer 1
  [ <4 x 4 in weights>, <4 x 4 out weights> <padding> ]  # `k` adapter for layer 1
  [ <2 x 4 in weights>, <4 x 2 out weights> <padding> ]  # `q` adapter for layer 2
  [ <8 x 4 in weights>, <4 x 8 out weights>           ]  # `q` adapter for layer 3. Note the final layer has a adapter size of 8
]

LoRA 模块 ID 映射

模块名称（在`convert_checkpoint.py`脚本中指定）	模块 ID	描述
attn_qkv	0	组合的 qkv 适配器
attn_q	1	q 适配器
attn_k	2	k 适配器
attn_v	3	v 适配器
注意力密集层	4	注意力机制中密集层的适配器
mlp_h_to_4h	5	用于llama2适配器，用于注意力/RMSNorm后的门控mlp层：向上投影
mlp_4h_to_h	6	用于llama2适配器，用于注意力/RMSNorm后的门控mlp层：向下投影
mlp_gate	7	用于在注意力/RMSNorm之后的gated mlp的llama2适配器：gate
cross_attn_qkv	8	用于交叉注意力的组合qkv适配器
cross_attn_q	9	用于交叉注意力的q适配器
cross_attn_k	10	用于交叉注意力的k适配器
cross_attn_v	11	用于交叉注意力的v适配器
cross_attn_dense	12	交叉注意力中密集层的适配器
moe_h_to_4h	13	用于专家MLP层的mixtral适配器：向上投影
moe_4h_to_h	14	用于专家MLP层的mixtral适配器：向下投影
moe_gate	15	用于专家MLP层的mixtral适配器：门
moe_router	16	用于专家路由器层的mixtral适配器
mlp_router	17	用于共享专家门层的qwen2-moe适配器

LoraCache 配置

核心思想是，我们将在TRT-LLM中拥有一个固定大小的两级LoRA缓存。高级缓存位于主机上，低级缓存位于GPU上（与现有的KV缓存不同）。两者的大小均可由用户配置。

CPU缓存被配置为最大大小。GPU缓存被配置为引擎加载后空闲GPU内存的百分比。随着请求的到来，LoRAs被存储在主机缓存中。

当请求被调度执行时，LoRAs会被加载到GPU缓存中。

使用张量并行的LoRA

LoRA的张量并行分区是特殊的。有两种情况：RowLinear 和 ColumnLinear。假设我们有一个线性层，输入特征大小为 K，输出特征大小为 N。那么，权重的形状为 [K, N]。

首先，考虑这个线性层是一个ColumnLinear层。当我们对权重进行分区时，我们使用tp_size按列分割权重。然后，有tp_size个分割的权重，这些权重的形状是[K, N // tp_size]。当我们在这样的ColumnLinear层上应用LoRA适配器时，原始两个权重的形状是[K, lora_rank]和[lora_rank, N]。因此，我们只对第二个权重进行分区，并得到tp_size个分割的权重，形状为[lora_rank, N // tp_size]。对于第一个权重，每个GPU保持相同的完整权重（形状为[K, lora_rank]）。

接下来，考虑这个线性层是一个RowLinear层。当我们对权重进行分区时，我们使用tp_size按行分割权重。然后，有tp_size个分割的权重，这些权重的形状是[K // tp_size, N]。当我们在这样的RowLinear层上应用LoRA适配器时，原始两个权重的形状是[K, lora_rank]和[lora_rank, N]。因此，我们只对第一个权重进行分区，并得到tp_size个分割的权重，其形状为[K // tp_size, lora_rank]。对于第二个权重，每个GPU保持相同的完整权重（形状为[lora_rank, N]）。