trtllm build configuration file
To create a script that takes the buildconfig.yaml file and parses its arguments into the trtllm build command line, you can use the argparse and yaml modules in Python. Here's an example script that accomplishes this:
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(['--checkpoint_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(['--logits_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'])])
if 'vocab_size' in checkpoint_config:
args.extend(['--vocab_size', str(checkpoint_config['vocab_size'])])
if 'n_positions' in checkpoint_config:
args.extend(['--n_positions', str(checkpoint_config['n_positions'])])
if 'n_layer' in checkpoint_config:
args.extend(['--n_layer', str(checkpoint_config['n_layer'])])
if 'n_head' in checkpoint_config:
args.extend(['--n_head', str(checkpoint_config['n_head'])])
if 'n_embd' in checkpoint_config:
args.extend(['--n_embd', str(checkpoint_config['n_embd'])])
if 'inter_size' in checkpoint_config:
args.extend(['--inter_size', str(checkpoint_config['inter_size'])])
if 'n_kv_head' in checkpoint_config:
args.extend(['--n_kv_head', str(checkpoint_config['n_kv_head'])])
if 'rms_norm_eps' in checkpoint_config:
args.extend(['--rms_norm_eps', str(checkpoint_config['rms_norm_eps'])])
if 'bos_token_id' in checkpoint_config:
args.extend(['--bos_token_id', str(checkpoint_config['bos_token_id'])])
if 'eos_token_id' in checkpoint_config:
args.extend(['--eos_token_id', str(checkpoint_config['eos_token_id'])])
if 'tie_word_embeddings' in checkpoint_config:
args.extend(['--tie_word_embeddings', str(checkpoint_config['tie_word_embeddings'])])
if 'use_cache' in checkpoint_config:
args.extend(['--use_cache', str(checkpoint_config['use_cache'])])
if 'torch_dtype' in checkpoint_config:
args.extend(['--torch_dtype', checkpoint_config['torch_dtype']])
if 'hidden_act' in checkpoint_config:
args.extend(['--hidden_act', checkpoint_config['hidden_act']])
# 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'])])
if 'max_prompt_embedding_table_size' in build_config:
args.extend(['--max_prompt_embedding_table_size', str(build_config['max_prompt_embedding_table_size'])])
if 'gather_context_logits' in build_config:
args.extend(['--gather_context_logits', str(build_config['gather_context_logits'])])
if 'gather_generation_logits' in build_config:
args.extend(['--gather_generation_logits', str(build_config['gather_generation_logits'])])
if 'strongly_typed' in build_config:
args.extend(['--strongly_typed', str(build_config['strongly_typed'])])
if 'profiling_verbosity' in build_config:
args.extend(['--profiling_verbosity', build_config['profiling_verbosity']])
if 'enable_debug_output' in build_config:
args.extend(['--enable_debug_output', str(build_config['enable_debug_output'])])
if 'max_draft_len' in build_config:
args.extend(['--max_draft_len', str(build_config['max_draft_len'])])
if 'use_refit' in build_config:
args.extend(['--use_refit', str(build_config['use_refit'])])
if 'weight_sparsity' in build_config:
args.extend(['--weight_sparsity', str(build_config['weight_sparsity'])])
if 'max_encoder_input_len' in build_config:
args.extend(['--max_encoder_input_len', str(build_config['max_encoder_input_len'])])
if 'use_fused_mlp' in build_config:
args.extend(['--use_fused_mlp', str(build_config['use_fused_mlp'])])
if 'dry_run' in build_config:
args.extend(['--dry_run', str(build_config['dry_run'])])
if 'visualize_network' in build_config:
args.extend(['--visualize_network', str(build_config['visualize_network'])])
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()# TensorRT-LLM Build Configuration File
model:
model_dir: ./llama-2-7b-chat-hf # Path to the pretrained model directory
output_dir: ./llama-2-7b-chat-engine # Path to save the built engine
dtype: float16 # Data type for the model (choices: float32, float16, bfloat16)
checkpoint:
checkpoint_dir: ../llama-2-7b-chat-hf-output # Path to the TensorRT-LLM checkpoint directory
tp_size: 1 # Tensor parallelism size, increase for multi-GPU tensor parallelism
pp_size: 1 # Pipeline parallelism size, increase for multi-GPU pipeline parallelism
vocab_size: 32000 # Vocabulary size of the model
n_positions: 2048 # Maximum number of positions (sequence length)
n_layer: 32 # Number of layers in the model
n_head: 32 # Number of attention heads
n_embd: 4096 # Hidden size of the model
inter_size: 11008 # Intermediate size of the model's feed-forward layers
#meta_ckpt_dir: # Path to the meta checkpoint directory
#n_kv_head: # Number of key-value heads (defaults to n_head if not specified)
#rms_norm_eps: 1e-6 # Epsilon value for RMS normalization
#use_weight_only: false # Enable weight-only quantization
#weight_only_precision: int8 # Precision for weight-only quantization (choices: int8, int4)
#smoothquant: 0.5 # Smoothquant parameter for quantization
#per_channel: false # Enable per-channel quantization
#per_token: false # Enable per-token quantization
#int8_kv_cache: false # Enable int8 quantization for key-value cache
#ammo_quant_ckpt_path: # Path to the quantized checkpoint file in .npz format
#per_group: false # Enable per-group quantization for GPTQ/AWQ quantization
#load_by_shard: false # Load the pretrained model shard-by-shard
#hidden_act: silu # Activation function used in the model (default: silu)
#rotary_base: 10000.0 # Base value for rotary positional embeddings
#group_size: 128 # Group size used in GPTQ quantization
#dataset_cache_dir: # Path to the dataset cache directory
#load_model_on_cpu: false # Load the model on CPU
#use_parallel_embedding: false # Enable embedding parallelism
#embedding_sharding_dim: 0 # Dimension for embedding sharding (choices: 0, 1)
#use_embedding_sharing: false # Enable embedding sharing to reduce engine size
#workers: 1 # Number of workers for parallel checkpoint conversion
#moe_num_experts: 0 # Number of experts for Mixture of Experts (MoE) layers
#moe_top_k: 0 # Top-k value for MoE layers (defaults to 1 if moe_num_experts is set)
#moe_tp_mode: 0 # Parallelism mode for distributing MoE experts in tensor parallelism
#moe_renorm_mode: 1 # Renormalization mode for MoE gate logits
#save_config_only: false # Only save the model configuration without building the engine
#disable_weight_only_quant_plugin: false # Disable the weight-only quantization plugin
build:
max_input_len: 256 # Maximum input sequence length
max_output_len: 256 # Maximum output sequence length
max_batch_size: 8 # Maximum batch size
max_beam_width: 1 # Maximum beam width for beam search
#max_num_tokens: # Maximum number of tokens to generate
#opt_num_tokens: # Optimal number of tokens to generate
max_prompt_embedding_table_size: 0 # Maximum size of the prompt embedding table
gather_context_logits: false # Gather context logits during generation
gather_generation_logits: false # Gather generation logits during generation
strongly_typed: false # Enable strongly typed network definition
#builder_opt: # Builder optimization level
profiling_verbosity: layer_names_only # Profiling verbosity level (choices: layer_names_only, detailed, none)
enable_debug_output: false # Enable debug output
max_draft_len: 0 # Maximum draft length for Medusa-style generation
use_refit: false # Enable engine refitting
#input_timing_cache: # Path to the input timing cache file
#output_timing_cache: # Path to save the output timing cache file
lora_config: # Configuration for LoRA (Low-Rank Adaptation)
#lora_dir: # Path to the LoRA checkpoint directory
#lora_target_modules: # Target modules for LoRA adaptation
#lora_ckpt_source: hf # Source of LoRA checkpoints (choices: hf, nemo)
#max_lora_rank: 4 # Maximum rank for LoRA adaptation
auto_parallel_config: # Configuration for automatic parallelization
#enabled: false # Enable automatic parallelization
#tp_size: 1 # Tensor parallelism size for automatic parallelization
#pp_size: 1 # Pipeline parallelism size for automatic parallelization
#max_memory_MB: 80000 # Maximum memory in MB for automatic parallelization
#max_dram_memory_MB: 30000 # Maximum DRAM memory in MB for automatic parallelization
#compile_max_memory_MB: 17000 # Maximum memory in MB for compilation during automatic parallelization
#compile_max_dram_memory_MB: 8000 # Maximum DRAM memory in MB for compilation during automatic parallelization
#debug_mode: false # Enable debug mode for automatic parallelization
weight_sparsity: false # Enable weight sparsity
plugin_config: # Configuration for plugins
#use_custom_all_reduce: false # Use custom all-reduce plugin
#use_fp8_all_reduce: false # Use FP8 all-reduce plugin
#use_fp8_cast_plugin: false # Use FP8 cast plugin
#use_async_malloc: false # Use asynchronous memory allocation plugin
#use_paged_context_fmha: false # Use paged context fused multi-head attention plugin
#use_fp8_context_fmha: false # Use FP8 context fused multi-head attention plugin
#lora_plugin: # Configuration for LoRA plugin
#type: # Type of LoRA plugin
max_encoder_input_len: 1024 # Maximum encoder input sequence length for encoder-decoder models
use_fused_mlp: false # Use fused MLP layers
dry_run: false # Perform a dry run without building the engine
visualize_network: false # Visualize the network graphThe updated script includes all the relevant configurations from the buildconfig.yaml file.
It parses the model, checkpoint, and build configurations and constructs the corresponding command-line arguments for the trtllm-build command.
python3 buildrun.py --config buildconfig.yamlpython build_trtllm.py --config buildconfig.yamlThis configuration file is divided into three main sections:
Model Configuration
Specifies the paths to the pretrained model directory and the output directory where the built engine will be saved. It also allows you to set the data type for the model (float32, float16, or bfloat16).
Checkpoint Configuration
Defines the settings related to the TensorRT-LLM checkpoint, such as the checkpoint directory, tensor parallelism size, pipeline parallelism size, and various model-specific parameters like vocabulary size, number of layers, attention heads, hidden size, etc.
Many of these settings are optional and can be uncommented and adjusted based on the specific model requirements.
Build Configuration
Contains the parameters for building the TensorRT engine, including maximum input and output sequence lengths, maximum batch size, beam width, prompt embedding table size, and various optimization and debugging options.
It also allows you to configure LoRA (Low-Rank Adaptation), automatic parallelization, weight sparsity, and plugin-specific settings.
Last updated
Was this helpful?