deprecate trtllm-build in weight_sparsity

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.

examples/llm_sparsity/weight_sparsity/export_trtllm_ckpt.py → 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,26 +103,18 @@ 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):
             raise FileNotFoundError(f"Sparsity state file {args.modelopt_restore_path} not found.")
 
         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",