Classification with LoRA and LLAMA/Mistral example

[bjayakumar]

System Info

TensorRT-LLM 0.8
Nvidia A100
Cuda 12.2

Who can help?

I can run LLAMA with LoRA weights for text generative tasks like completion and summarization. However, I have trained a model with classification LoRA. When I tried to run with run.py, I can able to do generative tasks. However, I wanted to do classifcation.

Can you post an example where LLM like LLAMA or MISTRAL is fine tuned with classification head using LoRA and convert those checkpoints and do classification by running run.py?

It appears model_runner does only decode?

TensorRT-LLM/tensorrt_llm/runtime/model_runner.py

Line 788 in 89ba1b1

outputs = self.session.decode(

Information

• The official example scripts
• My own modified scripts

Tasks

• An officially supported task in the examples folder (such as GLUE/SQuAD, ...)
• My own task or dataset (give details below)

Reproduction

.

Expected behavior

A working example for classification with LoRa and LLM

actual behavior

Unavailable

additional notes

.

[byshiue]

The requirement is not very clear for me. Do you want to get the context logits of model instead of generating tokens? If so, you could get the logits of context by adding --gather_context_logits during building engine.

[bjayakumar]

Sorry I was not clear.

I took HF LLM like Mistral and used it for classification task using LoRA.
I have saved the model.

I can build an engine with Foundation model and LoRA which does classification.

When I run, run.py with that engine, I can only get generation output not classification

python ../run.py --engine_dir "$ENGINE_DIR" \
              --max_output_len 10 \
              --tokenizer_dir ${BASE_LLAMA_MODEL} \
              --input_text "i want to pay for my book" \
              --lora_task_uids 0 \
              --use_py_session --top_p 0.5 --top_k 0

I would like to know how to get classification output.

[byshiue]

TensorRT-LLM only provided the examples about generation. For classification, you might need to change the example script.

[badrinath89]

Can you pleaser share pointer or direction to solve this?

What changes should I make in the code?

[bjayakumar]

TensorRT-LLM only provided the examples about generation. For classification, you might need to change the example script.

How do you change the example script? Can you suggest a psuedo code?