Run the buildconfig file
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()Last updated
Was this helpful?

