This example implements a NanoGPT model using Tripy:

1. model.py defines the model as an nvtripy.Module.
2. weight_loader.py loads weights from a HuggingFace checkpoint.
3. example.py runs inference in float16 on input text and displays the output.

1. Install prerequisites:

   python3 -m pip install -r requirements.txt

2. Run the example:

   python3 example.py --input-text "What is the answer to life, the universe, and everything?"

3. [Optional] Use a fixed seed for predictable outputs:

   python3 example.py --input-text "What is the answer to life, the universe, and everything?" --seed=0

quantization.py, uses NVIDIA TensorRT Model Optimizer to quantize the pytorch model.

load_quant_weights_from_hf in weight_loader.py converts the quantization parameters to scales and loads them into the Tripy model.

Use --quant-mode in example.py to enable quantization. Supported modes:

• Weight-only int8 quantization:

  python3 example.py --input-text "What is the answer to life, the universe, and everything?" --seed=0 --quant-mode int8-weight-only

• Weight-only int4 quantization:

  python3 example.py --input-text "What is the answer to life, the universe, and everything?" --seed=0 --quant-mode int4-weight-only