Examples of running the convert_checkpoint.py script
This process will help you create the arguments that go with the checkpoint.py script
Examples of running the convert_checkpoint.py
script
checkpoint.py
scriptThe script provides three examples demonstrating different use cases for running the
convert_heckpoint.py
script.Each example shows the specific command to run
convert_checkpoint.py
with the corresponding arguments.The examples cover scenarios such as converting a checkpoint using a single GPU with FP16 or BF16 precision, and converting a checkpoint with specific model hyperparameters.
1. Convert a checkpoint using a single GPU and FP16:
python convert_checkpoint.py --model_dir ./path/to/model/directory \
--output_dir ./tllm_checkpoint_1gpu_fp16 \
--dtype float16
2. Convert a checkpoint using a single GPU and BF16:
python convert_checkpoint.py --model_dir ./path/to/model/directory \
--output_dir ./tllm_checkpoint_1gpu_bf16 \
--dtype bfloat16
3. Convert a checkpoint with specific model hyperparameters:
python convert_checkpoint.py --model_dir ./path/to/model/directory \
--output_dir ./path/to/output/directory \
--dtype float16 \
--n_layer 32 \
--n_head 32 \
--n_embd 4096
README: Converting Checkpoints with TensorRT-LLM
This guide explains how to use the run_convert_checkpoint.py
script to run the convert_checkpoint.py
script with transparent configurations using a YAML file.
First download the scripts from Github
git clone https://github.com/Continuum-Labs-HQ/tensorrt-continuum.git
You will need to manually move the files into the right directory for now
Prerequisites
Python 3.x
TensorRT-LLM
PyYAML (install with
pip install pyyaml
)
Configuration
Populate the
config.yaml
file in the same directory as therun_convert_checkpoint.py
script. This file will contain the configurations for theconvert_checkpoint.py
script.Open the
config.yaml
file and specify the desired configurations. The file should have the following structure:
model:
model_dir: ./path/to/model
output_dir: ./path/to/output
dtype: float16
checkpoint:
tp_size: 1
pp_size: 1
vocab_size: 32000
n_positions: 2048
n_layer: 32
n_head: 32
n_embd: 4096
inter_size: 11008
# Additional checkpoint arguments
# ...
Adjust the values in the
config.yaml
file according to your requirements. You can refer to the comments in the file for suggestions and available choices.
Running the Script
Open a terminal and navigate to the directory containing the
run_convert_checkpoint.py
script and theconfig.yaml
file.Run the following command:
python3 run_convert_checkpoint.py
The script will read the configurations from the
config.yaml
file and construct the corresponding command-line arguments for theconvert_checkpoint.py
script.The
convert_checkpoint.py
script will be executed with the specified configurations.
How It Works
The run_convert_checkpoint.py
script does the following:
It uses the
argparse
module to parse the command-line argument--config
, which specifies the path to the YAML configuration file (default isconfig.yaml
).It loads the configurations from the specified YAML file using the
yaml.safe_load()
function.It extracts the relevant configuration values from the loaded YAML data, such as
model_dir
,output_dir
, anddtype
.It constructs the command-line arguments for the
convert_checkpoint.py
script based on the extracted configuration values.It iterates over the checkpoint arguments in the YAML file and adds them to the command-line arguments list. Boolean values are handled separately, adding only the argument flag if the value is
True
.Finally, it uses the
subprocess.run()
function to execute theconvert_checkpoint.py
script with the constructed command-line arguments.
By using the run_convert_checkpoint.py
script and the config.yaml
file, you can easily configure and run the convert_checkpoint.py
script without explicitly specifying all the command-line arguments.
This approach provides a more transparent and manageable way to set the configurations for the checkpoint conversion process.
Last updated
Was this helpful?