# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.

#

# Licensed under the Apache License, Version 2.0 (the "License");

# you may not use this file except in compliance with the License.

# You may obtain a copy of the License at

#

# http://www.apache.org/licenses/LICENSE-2.0

#

# Unless required by applicable law or agreed to in writing, software

# distributed under the License is distributed on an "AS IS" BASIS,

# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

# See the License for the specific language governing permissions and

# limitations under the License.

import argparse

import logging

import sys

from nemo.export import TensorRTLLM

LOGGER = logging.getLogger("NeMo")

def get_args(argv):

parser = argparse.ArgumentParser(

formatter_class=argparse.ArgumentDefaultsHelpFormatter,

description=f"Exports nemo models stored in nemo checkpoints to TensorRT-LLM",

)

parser.add_argument("-nc", "--nemo_checkpoint", required=True, type=str, help="Source .nemo file")

parser.add_argument(

"-mt",

"--model_type",

type=str,

required=True,

choices=["gptnext", "gpt", "llama", "falcon", "starcoder", "gemma"],

help="Type of the model. gptnext, gpt, llama, falcon, and starcoder are only supported."

" gptnext and gpt are the same and keeping it for backward compatibility",

)

parser.add_argument(

"-mr", "--model_repository", required=True, default=None, type=str, help="Folder for the trt-llm model files"

)

parser.add_argument("-ng", "--num_gpus", default=1, type=int, help="Number of GPUs for the deployment")

parser.add_argument("-tps", "--tensor_parallelism_size", type=int, help="Tensor parallelism size")

parser.add_argument("-pps", "--pipeline_parallelism_size", type=int, help="Pipeline parallelism size")

parser.add_argument(

"-dt",

"--dtype",

choices=["bfloat16", "float16", "fp8", "int8"],

default="bfloat16",

type=str,

help="dtype of the model on TensorRT-LLM",

)

parser.add_argument("-mil", "--max_input_len", default=256, type=int, help="Max input length of the model")

parser.add_argument("-mol", "--max_output_len", default=256, type=int, help="Max output length of the model")

parser.add_argument("-mbs", "--max_batch_size", default=8, type=int, help="Max batch size of the model")

parser.add_argument("-mnt", "--max_num_tokens", default=None, type=int, help="Max number of tokens")

parser.add_argument("-ont", "--opt_num_tokens", default=None, type=int, help="Optimum number of tokens")

parser.add_argument(

"-mpet", "--max_prompt_embedding_table_size", default=None, type=int, help="Max prompt embedding table size"

)

parser.add_argument(

"-npkc", "--no_paged_kv_cache", default=False, action='store_true', help="Enable paged kv cache."

)

parser.add_argument(

"-drip",

"--disable_remove_input_padding",

default=False,

action='store_true',

help="Disables the remove input padding option.",

)

parser.add_argument(

"-mbm",

'--multi_block_mode',

default=False,

action='store_true',

help='Split long kv sequence into multiple blocks (applied to generation MHA kernels). \

It is beneifical when batchxnum_heads cannot fully utilize GPU.',

)

parser.add_argument(

'--use_lora_plugin',

nargs='?',

const=None,

choices=['float16', 'float32', 'bfloat16'],

help="Activates the lora plugin which enables embedding sharing.",

)

parser.add_argument(

'--lora_target_modules',

nargs='+',

default=None,

choices=[

"attn_qkv",

"attn_q",

"attn_k",

"attn_v",

"attn_dense",

"mlp_h_to_4h",

"mlp_gate",

"mlp_4h_to_h",

],

help="Add lora in which modules. Only be activated when use_lora_plugin is enabled.",

)

parser.add_argument(

'--max_lora_rank',

type=int,

default=64,

help='maximum lora rank for different lora modules. '

'It is used to compute the workspace size of lora plugin.',

)

parser.add_argument("-dm", "--debug_mode", default=False, action='store_true', help="Enable debug mode")

args = parser.parse_args(argv)

return args

def nemo_export_trt_llm(argv):

args = get_args(argv)

if args.debug_mode:

loglevel = logging.DEBUG

else:

loglevel = logging.INFO

LOGGER.setLevel(loglevel)

LOGGER.info("Logging level set to {}".format(loglevel))

LOGGER.info(args)

if args.dtype != "bfloat16":

LOGGER.error(

"Only bf16 is currently supported for the optimized deployment with TensorRT-LLM. "

"Support for the other precisions will be added in the coming releases."

)

return

try:

trt_llm_exporter = TensorRTLLM(model_dir=args.model_repository, load_model=False)

LOGGER.info("Export to TensorRT-LLM function is called.")

trt_llm_exporter.export(

nemo_checkpoint_path=args.nemo_checkpoint,

model_type=args.model_type,

n_gpus=args.num_gpus,

tensor_parallel_size=args.tensor_parallelism_size,

pipeline_parallel_size=args.pipeline_parallelism_size,

max_input_token=args.max_input_len,

max_output_token=args.max_output_len,

max_batch_size=args.max_batch_size,

max_num_tokens=args.max_num_tokens,

opt_num_tokens=args.opt_num_tokens,

max_prompt_embedding_table_size=args.max_prompt_embedding_table_size,

paged_kv_cache=(not args.no_paged_kv_cache),

remove_input_padding=(not args.disable_remove_input_padding),

dtype=args.dtype,

enable_multi_block_mode=args.multi_block_mode,

use_lora_plugin=args.use_lora_plugin,

lora_target_modules=args.lora_target_modules,

max_lora_rank=args.max_lora_rank,

save_nemo_model_config=True,

)

LOGGER.info("Export is successful.")

except Exception as error:

LOGGER.error("Error message: " + str(error))

raise error

if __name__ == '__main__':

nemo_export_trt_llm(sys.argv[1:])