Run the buildconfig file
To build a script that parses the buildconfig.yaml
arguments and passes them to the build.py
script, you can create a new Python script that reads the YAML configuration file, extracts the relevant settings, and constructs the appropriate command-line arguments for the trtllm-build
command.
Here's an example of how you can create such a script:
import argparse
import subprocess
import yaml
def parse_buildconfig(config_file):
with open(config_file, 'r') as f:
config = yaml.safe_load(f)
args = []
# Model Configuration
if 'model' in config:
model_config = config['model']
if 'model_dir' in model_config:
args.extend(['--model_dir', model_config['model_dir']])
if 'output_dir' in model_config:
args.extend(['--output_dir', model_config['output_dir']])
if 'dtype' in model_config:
args.extend(['--dtype', model_config['dtype']])
# Checkpoint Configuration
if 'checkpoint' in config:
checkpoint_config = config['checkpoint']
if 'checkpoint_dir' in checkpoint_config:
args.extend(['--checkpoint_dir', checkpoint_config['checkpoint_dir']])
if 'tp_size' in checkpoint_config:
args.extend(['--tp_size', str(checkpoint_config['tp_size'])])
if 'pp_size' in checkpoint_config:
args.extend(['--pp_size', str(checkpoint_config['pp_size'])])
# Add more checkpoint configuration options as needed
# Build Configuration
if 'build' in config:
build_config = config['build']
if 'max_input_len' in build_config:
args.extend(['--max_input_len', str(build_config['max_input_len'])])
if 'max_output_len' in build_config:
args.extend(['--max_output_len', str(build_config['max_output_len'])])
if 'max_batch_size' in build_config:
args.extend(['--max_batch_size', str(build_config['max_batch_size'])])
if 'max_beam_width' in build_config:
args.extend(['--max_beam_width', str(build_config['max_beam_width'])])
# Add more build configuration options as needed
return args
def main():
parser = argparse.ArgumentParser(description='Parse buildconfig.yaml and run trtllm-build')
parser.add_argument('--config', type=str, required=True, help='Path to the buildconfig.yaml file')
args = parser.parse_args()
buildconfig_args = parse_buildconfig(args.config)
command = ['trtllm-build'] + buildconfig_args
subprocess.run(command, check=True)
if __name__ == '__main__':
main()
In this script:
We define a function called
parse_buildconfig
that takes the path to thebuildconfig.yaml
file as input. It reads the YAML file, extracts the relevant settings from the model, checkpoint, and build configurations, and constructs a list of command-line arguments based on the settings.We define the
main
function that usesargparse
to parse the command-line arguments. It expects a--config
argument that specifies the path to thebuildconfig.yaml
file.Inside the
main
function, we call theparse_buildconfig
function to parse the YAML file and obtain the command-line arguments.We construct the
trtllm-build
command by concatenating the base command with the parsed command-line arguments.Finally, we use
subprocess.run
to execute thetrtllm-build
command with the provided arguments.
To use this script, save it to a file (e.g., run_trtllm_build.py
) and run it from the command line, providing the path to the buildconfig.yaml
file:
trtllm-build --checkpoint_dir ./tllm_checkpoint_1gpu_fp16 \
--output_dir ./llama/7B/trt_engines/fp16/1-gpu \
--gemm_plugin float16
python3 buildrun.py --config buildconfig.yaml
This script will parse the buildconfig.yaml
file, extract the relevant settings, and pass them as command-line arguments to the trtllm-build
command.
Note: Make sure you have the necessary dependencies installed (yaml
and argparse
) before running the script.
Last updated
Was this helpful?