diff --git a/examples/llm_sparsity/weight_sparsity/README.md b/examples/llm_sparsity/weight_sparsity/README.md index 97563aff007..42bac030ce9 100644 --- a/examples/llm_sparsity/weight_sparsity/README.md +++ b/examples/llm_sparsity/weight_sparsity/README.md @@ -135,28 +135,39 @@ ROUGE scores: {'rouge1': 42.174, 'rouge2': 19.2724, 'rougeL': 28.6989, 'rougeLsu Please refer to [link](../llm_eval/README.md#Evaluation-scripts-for-LLM-tasks) for more details of how to evaluate the sparsified models on other benchmarks, such as MMLU and HumanEval. -## Export TensorRT-LLM Checkpoint +## Export and Serve a Hugging Face Checkpoint -To export the PTS pytorch model to a TensorRT-LLM checkpoint, run the following command: +Export the PTS model to a Hugging Face checkpoint: ```sh -python export_trtllm_ckpt.py --model_name_or_path meta-llama/Llama-2-7b-hf \ +python export_hf_ckpt.py --model_name_or_path meta-llama/Llama-2-7b-hf \ --model_max_length 1024 \ --dtype fp16 \ - --modelopt_restore_path saved_models_Llama-2-7b-hf_sparsegpt_tp1_pp1/pts_modelopt_state.pth \ - --output_dir saved_models_Llama-2-7b-hf_sparsegpt_tp1_pp1/trtllm/ckpt_pts + --modelopt_restore_path saved_models_Llama-2-7b-hf_sparsegpt_tp1_pp1/pts/pts_modelopt_state.pth \ + --output_dir saved_models_Llama-2-7b-hf_sparsegpt_tp1_pp1/hf/pts ``` -To export the finetuned pytorch model to a TensorRT-LLM checkpoint, run the following command: +Export the finetuned model in the same way: ```sh -python export_trtllm_ckpt.py --model_name_or_path meta-llama/Llama-2-7b-hf \ +python export_hf_ckpt.py --model_name_or_path meta-llama/Llama-2-7b-hf \ --model_max_length 1024 \ --dtype fp16 \ --modelopt_restore_path saved_models_Llama-2-7b-hf_sparsegpt_tp1_pp1/finetuned/finetuned_modelopt_state.pth \ - --output_dir saved_models_Llama-2-7b-hf_sparsegpt_tp1_pp1/trtllm/ckpt_finetuned + --output_dir saved_models_Llama-2-7b-hf_sparsegpt_tp1_pp1/hf/finetuned ``` -## Build TensorRT-LLM Engine +The output is a standard Hugging Face checkpoint, including the tokenizer and sparse model weights. +TensorRT-LLM loads it directly; no checkpoint conversion or engine build is required. -For guidance on how to build TensorRT-LLM engines, please refer to [link](https://nvidia.github.io/TensorRT-LLM/commands/trtllm-build.html#trtllm-build) and use the `--weight_sparsity` flag. +Serve the exported checkpoint and set the deployment parallelism at startup: + +```sh +trtllm-serve saved_models_Llama-2-7b-hf_sparsegpt_tp1_pp1/hf/pts \ + --tp_size 1 \ + --pp_size 1 \ + --host 0.0.0.0 \ + --port 8000 +``` + +After the server starts, send requests to the OpenAI-compatible `/v1/chat/completions` endpoint. diff --git a/examples/llm_sparsity/weight_sparsity/export_trtllm_ckpt.py b/examples/llm_sparsity/weight_sparsity/export_hf_ckpt.py similarity index 84% rename from examples/llm_sparsity/weight_sparsity/export_trtllm_ckpt.py rename to examples/llm_sparsity/weight_sparsity/export_hf_ckpt.py index e23cdb677b0..7ef22c52eca 100644 --- a/examples/llm_sparsity/weight_sparsity/export_trtllm_ckpt.py +++ b/examples/llm_sparsity/weight_sparsity/export_hf_ckpt.py @@ -23,8 +23,7 @@ import modelopt.torch.opt as mto import modelopt.torch.sparsity as mts -from modelopt.torch.export import get_model_type -from modelopt.torch.export.trtllm import export_tensorrt_llm_checkpoint +from modelopt.torch.export import export_hf_checkpoint DEFAULT_PAD_TOKEN = "[PAD]" @@ -104,8 +103,6 @@ def main(args): model=model, ) - # Export the sparse model to trt-llm checkpoint - model_type = get_model_type(model) if args.modelopt_restore_path: print(f"Loading sparsity state from {args.modelopt_restore_path}") if not os.path.isfile(args.modelopt_restore_path): @@ -113,17 +110,11 @@ def main(args): mto.restore(model, args.modelopt_restore_path) - print(f"Exporting trt-llm checkpoint to {args.output_dir}") + print(f"Exporting Hugging Face checkpoint to {args.output_dir}") with torch.inference_mode(): model = mts.export(model) - export_tensorrt_llm_checkpoint( - model, - model_type, - torch.float16, - export_dir=args.output_dir, - inference_tensor_parallel=args.inference_tensor_parallel, - inference_pipeline_parallel=args.inference_pipeline_parallel, - ) + export_hf_checkpoint(model, export_dir=args.output_dir) + tokenizer.save_pretrained(args.output_dir) if __name__ == "__main__": @@ -141,22 +132,11 @@ def main(args): parser.add_argument("--dtype", help="Model data type.", default="fp16") parser.add_argument( "--model_max_length", + type=int, default=2048, help="Maximum sequence length. Sequences will be right padded (and possibly truncated).", ) parser.add_argument("--output_dir", default="output_dir") - parser.add_argument( - "--inference_tensor_parallel", - help="Number of tensor parallel groups for inference.", - type=int, - default=1, - ) - parser.add_argument( - "--inference_pipeline_parallel", - help="Number of pipeline parallel groups for inference.", - type=int, - default=1, - ) parser.add_argument( "--trust_remote_code", help="Set trust_remote_code for Huggingface models and tokenizers",