trtllm-build#

usage: trtllm-build [-h] [--checkpoint_dir CHECKPOINT_DIR] [--model_config MODEL_CONFIG]
                    [--build_config BUILD_CONFIG] [--model_cls_file MODEL_CLS_FILE]
                    [--model_cls_name MODEL_CLS_NAME] [--output_dir OUTPUT_DIR]
                    [--max_batch_size MAX_BATCH_SIZE] [--max_input_len MAX_INPUT_LEN]
                    [--max_seq_len MAX_SEQ_LEN] [--max_beam_width MAX_BEAM_WIDTH]
                    [--max_num_tokens MAX_NUM_TOKENS] [--opt_num_tokens OPT_NUM_TOKENS]
                    [--max_encoder_input_len MAX_ENCODER_INPUT_LEN]
                    [--max_prompt_embedding_table_size MAX_PROMPT_EMBEDDING_TABLE_SIZE]
                    [--kv_cache_type KV_CACHE_TYPE] [--paged_kv_cache PAGED_KV_CACHE]
                    [--input_timing_cache INPUT_TIMING_CACHE]
                    [--output_timing_cache OUTPUT_TIMING_CACHE]
                    [--profiling_verbosity {layer_names_only,detailed,none}] [--strip_plan]
                    [--weight_sparsity] [--weight_streaming] [--fast_build]
                    [--workers WORKERS]
                    [--log_level {internal_error,error,warning,info,verbose,debug}]
                    [--enable_debug_output] [--visualize_network VISUALIZE_NETWORK]
                    [--dry_run] [--monitor_memory] [--logits_dtype {float16,float32}]
                    [--gather_context_logits] [--gather_generation_logits]
                    [--gather_all_token_logits] [--lora_dir LORA_DIR [LORA_DIR ...]]
                    [--lora_ckpt_source {hf,nemo}]
                    [--lora_target_modules {attn_qkv,attn_q,attn_k,attn_v,attn_dense,mlp_h_to_4h,mlp_4h_to_h,mlp_gate,cross_attn_qkv,cross_attn_q,cross_attn_k,cross_attn_v,cross_attn_dense,moe_h_to_4h,moe_4h_to_h,moe_gate,moe_router,mlp_router,mlp_gate_up} [{attn_qkv,attn_q,attn_k,attn_v,attn_dense,mlp_h_to_4h,mlp_4h_to_h,mlp_gate,cross_attn_qkv,cross_attn_q,cross_attn_k,cross_attn_v,cross_attn_dense,moe_h_to_4h,moe_4h_to_h,moe_gate,moe_router,mlp_router,mlp_gate_up} ...]]
                    [--max_lora_rank MAX_LORA_RANK]
                    [--speculative_decoding_mode {draft_tokens_external,lookahead_decoding,medusa,explicit_draft_tokens,eagle}]
                    [--max_draft_len MAX_DRAFT_LEN] [--auto_parallel AUTO_PARALLEL]
                    [--gpus_per_node GPUS_PER_NODE]
                    [--cluster_key {A100-SXM-80GB,A100-SXM-40GB,A100-PCIe-80GB,A100-PCIe-40GB,H100-SXM,H100-PCIe,H20,H200-SXM,H200-NVL,A40,A30,A10,A10G,L40S,L40,L20,L4,L2}]
                    [--bert_attention_plugin {auto,float16,float32,bfloat16,int32,disable}]
                    [--gpt_attention_plugin {auto,float16,float32,bfloat16,int32,disable}]
                    [--gemm_plugin {auto,float16,float32,bfloat16,int32,fp8,nvfp4,disable}]
                    [--gemm_swiglu_plugin {fp8,disable}]
                    [--fp8_rowwise_gemm_plugin {auto,float16,float32,bfloat16,int32,disable}]
                    [--nccl_plugin {auto,float16,float32,bfloat16,int32,disable}]
                    [--lora_plugin {auto,float16,float32,bfloat16,int32,disable}]
                    [--dora_plugin {enable,disable}]
                    [--moe_plugin {auto,float16,float32,bfloat16,int32,disable}]
                    [--mamba_conv1d_plugin {auto,float16,float32,bfloat16,int32,disable}]
                    [--low_latency_gemm_plugin {fp8,disable}]
                    [--low_latency_gemm_swiglu_plugin {fp8,disable}]
                    [--gemm_allreduce_plugin {float16,bfloat16,disable}]
                    [--context_fmha {enable,disable}]
                    [--bert_context_fmha_fp32_acc {enable,disable}]
                    [--remove_input_padding {enable,disable}]
                    [--norm_quant_fusion {enable,disable}] [--reduce_fusion {enable,disable}]
                    [--user_buffer {enable,disable}] [--tokens_per_block TOKENS_PER_BLOCK]
                    [--use_paged_context_fmha {enable,disable}]
                    [--use_fp8_context_fmha {enable,disable}]
                    [--fuse_fp4_quant {enable,disable}]
                    [--multiple_profiles {enable,disable}] [--paged_state {enable,disable}]
                    [--streamingllm {enable,disable}] [--use_fused_mlp {enable,disable}]
                    [--pp_reduce_scatter {enable,disable}]

命名参数#

--checkpoint_dir

包含 TensorRT-LLM 检查点的目录路径。

--model_config

保存 TensorRT-LLM 检查点配置的文件路径。

--build_config

保存 TensorRT-LLM 构建配置的文件路径。

--model_cls_file

定义自定义 TensorRT-LLM 模型的文件路径。

--model_cls_name

自定义 TensorRT-LLM 模型类名。

--output_dir

保存序列化引擎文件和引擎配置文件的目录路径。

默认值：'engine_outputs'

--max_batch_size

引擎可以调度的最大请求数。

默认值：2048

--max_input_len

一个请求的最大输入长度。

默认值：1024

--max_seq_len, --max_decoder_seq_len

一个请求的最大总长度，包括提示和输出。如果未指定，则该值从模型配置中推导得出。

--max_beam_width

用于束搜索解码的最大束数。

默认值：1

--max_num_tokens

在每个批次中移除填充后，批处理输入 tokens 的最大数量。目前，默认情况下会移除输入填充；您可以通过指定 --remove_input_padding disable 显式禁用它。

默认值：8192

--opt_num_tokens

在每个批次中移除填充后，批处理输入 tokens 的最佳数量。它默认为 max_batch_size * max_beam_width，将此值设置为尽可能接近您工作负载上的实际 tokens 数量。请注意，此参数将来可能会被删除。

--max_encoder_input_len

用于 enc-dec 模型的最长编码器输入长度。设置 max_input_len 为 1 以从长度为 1 的 decoder_start_token_id 开始生成。

默认值：1024

--max_prompt_embedding_table_size, --max_multimodal_len

用于提示调整的最大提示嵌入表大小，或用于多模态模型的多模态最大输入大小。设置大于 0 的值会启用提示调整或多模态输入。

默认值：0

--kv_cache_type

设置 KV 缓存类型（连续、分页或禁用）。对于禁用情况，KV 缓存被禁用，并且仅允许上下文阶段。

--paged_kv_cache

已弃用。对于基于 Transformer 的模型，启用此选项等同于 --kv_cache_type paged。

--input_timing_cache

读取计时缓存的文件路径。如果该文件不存在，则忽略此选项。

--output_timing_cache

写入计时缓存的文件路径。

默认值：'model.cache'

--profiling_verbosity

可能的值：layer_names_only、detailed、none

生成的 TensorRT 引擎的性能分析详细程度。设置为 detailed 允许检查策略选择和内核参数。

默认值：'layer_names_only'

--strip_plan

启用从最终 TensorRT 引擎中剥离权重，假设重拟合权重与构建时提供的权重相同。

默认值：False

--weight_sparsity

启用权重稀疏性。

默认值：False

--weight_streaming

启用将权重卸载到 CPU 并在运行时流式加载。

默认值：False

--fast_build

启用用于更快引擎构建的功能。这可能会导致一些性能下降，并且当前与没有插件的 int8/int4 量化不兼容。

默认值：False

--workers

用于并行构建的工作人员数量。

默认值：1

--log_level

可能的值：internal_error、error、warning、info、verbose、debug

日志记录级别。

默认值：'info'

--enable_debug_output

启用调试输出。

默认值：False

--visualize_network

将 TensorRT 网络导出为 ONNX 的目录路径，以便在引擎构建之前进行调试。

--dry_run

运行整个构建过程，除了实际的引擎构建，用于调试。

默认值：False

--monitor_memory

在引擎构建期间启用内存监视器。

默认值：False

Logits 参数#

--logits_dtype

可能的值：float16、float32

logits 的数据类型。

--gather_context_logits

启用收集上下文 logits。

默认值：False

--gather_generation_logits

启用收集生成 logits。

默认值：False

--gather_all_token_logits

同时启用 gather_context_logits 和 gather_generation_logits。

默认值：False

LoRA 参数#

--lora_dir

LoRA 权重的目录。如果提供多个目录，则第一个目录用于配置。

--lora_ckpt_source

可能的值：hf、nemo

LoRA 检查点的来源类型。

默认值：'hf'

--lora_target_modules

可能的值：attn_qkv、attn_q、attn_k、attn_v、attn_dense、mlp_h_to_4h、mlp_4h_to_h、mlp_gate、cross_attn_qkv、cross_attn_q、cross_attn_k、cross_attn_v、cross_attn_dense、moe_h_to_4h、moe_4h_to_h、moe_gate、moe_router、mlp_router、mlp_gate_up

应用 LoRA 的目标模块名称。仅当启用 lora_plugin 时有效。

--max_lora_rank

不同 LoRA 模块的最大 LoRA 秩。它用于计算 LoRA 插件的工作空间大小。

默认值：64

推测性解码参数#

--speculative_decoding_mode

可能的值：draft_tokens_external、lookahead_decoding、medusa、explicit_draft_tokens、eagle

推测性解码模式。

--max_draft_len

用于推测性解码目标模型的最长草稿 tokens 长度。

默认值：0

自动并行参数#

--auto_parallel

用于自动并行的 MPI 世界大小。

默认值：1

--gpus_per_node

在多节点设置中，每个节点拥有的 GPU 数量。这是一个集群规范，可以大于/小于世界大小。此选项仅用于通过 --auto_parallel 指定的自动并行。

默认值：8

--cluster_key

可能的值：A100-SXM-80GB、A100-SXM-40GB、A100-PCIe-80GB、A100-PCIe-40GB、H100-SXM、H100-PCIe、H20、H200-SXM、H200-NVL、A40、A30、A10、A10G、L40S、L40、L20、L4、L2

目标 GPU 类型的唯一名称。如果未指定，则从当前 GPU 类型推断。此选项仅用于通过 --auto_parallel 指定的自动并行。

插件配置参数#

--bert_attention_plugin

可选选项: auto, float16, float32, bfloat16, int32, disable

该插件使用高效内核，并为类 BERT 编码器模型的注意力层启用 KV 缓存的原地更新。

默认值: 'auto'

--gpt_attention_plugin

可选选项: auto, float16, float32, bfloat16, int32, disable

该插件使用高效内核，并为类 GPT 解码器模型的注意力层启用 KV 缓存的原地更新。

默认值: 'auto'

--gemm_plugin

可选选项: auto, float16, float32, bfloat16, int32, fp8, nvfp4, disable

GEMM 插件利用 NVIDIA cuBLASLt 来执行 GEMM 操作。注意：它只对非量化的 gemm 操作有效（FP8 除外）。注意：对于 FP8，它还需要检查点中的相同校准。

--gemm_swiglu_plugin

可选选项: fp8, disable

Gated-MLP 中的 GEMM + SwiGLU 融合将两个 Matmul 操作和一个 SwiGLU 操作组合到一个内核中。目前，这仅支持 Hopper 上的 FP8 精度。

默认值: 'disable'

--fp8_rowwise_gemm_plugin

可选选项: auto, float16, float32, bfloat16, int32, disable

用于 fp8 的量化 GEMM，它对激活使用每个 token 的动态缩放，对权重使用每个通道的静态缩放。注意：它还需要检查点中的相同校准。

默认值: 'disable'

--nccl_plugin

可选选项: auto, float16, float32, bfloat16, int32, disable

NCCL 插件包装 NCCL 运算符以支持多 GPU 甚至多节点。

默认值: 'auto'

--lora_plugin

可选选项: auto, float16, float32, bfloat16, int32, disable

启用 LoRA。

默认值: 'disable'

--dora_plugin

可选选项: enable, disable

启用 DoRA。

默认值: 'disable'

--moe_plugin

可选选项: auto, float16, float32, bfloat16, int32, disable

启用一些自定义内核以加速 MoE 模型的 MoE 层。

默认值: 'auto'

--mamba_conv1d_plugin

可选选项: auto, float16, float32, bfloat16, int32, disable

启用自定义内核以加速 Mamba 的 conv1d 运算符。

默认值: 'auto'

--low_latency_gemm_plugin

可选选项: fp8, disable

专为低延迟场景优化的 GEMM 插件。

默认值: 'disable'

--low_latency_gemm_swiglu_plugin

可选选项: fp8, disable

专为低延迟场景优化的 GEMM + SwiGLU 融合插件。

默认值: 'disable'

--gemm_allreduce_plugin

可选选项: float16, bfloat16, disable

GEMM + AllReduce 内核融合插件。

默认值: 'disable'

--context_fmha

可选选项: enable, disable

在上下文阶段启用融合多头注意力，将触发使用单个内核执行 MHA/MQA/GQA 块的内核。

默认值: 'enable'

--bert_context_fmha_fp32_acc

可选选项: enable, disable

为 bert_attention_plugin 中的上下文 FMHA 启用 FP32 累加器。如果禁用，则使用 FP16，预计性能更好，但准确性可能更差。

默认值: 'disable'

--remove_input_padding

可选选项: enable, disable

将不同的 token 打包在一起，从而减少计算量和内存消耗。

默认值: 'enable'

--norm_quant_fusion

可选选项: enable, disable

将 LayerNorm 和量化内核融合到单个内核中，从而提高端到端性能。

默认值: 'disable'

--reduce_fusion

可选选项: enable, disable

将 AllReduce 后的 ResidualAdd 和 LayerNorm 内核融合到单个内核中，从而提高端到端性能。

默认值: 'disable'

--user_buffer

可选选项: enable, disable

消除通信内核中从本地缓冲区到共享缓冲区的额外复制，从而提高端到端性能。此功能必须与 –reduce_fusion enable 一起启用，并且目前仅支持 FP8 LLAMA 模型。

默认值: 'disable'

--tokens_per_block

定义每个分页 kv 缓存块中包含多少个 token。

默认值: 32

--use_paged_context_fmha

可选选项: enable, disable

允许高级功能，如 KV 缓存重用和分块上下文。

默认值: 'enable'

--use_fp8_context_fmha

可选选项: enable, disable

当 FP8 量化被激活时，可以通过启用 FP8 Context FMHA 来进一步加速注意力机制

默认值: 'enable'

--fuse_fp4_quant

可选选项: enable, disable

是否将 FP4 量化融合到注意力内核中。

默认值: 'disable'

--multiple_profiles

可选选项: enable, disable

在构建的引擎中启用多个 TensorRT 优化配置文件，当禁用 GEMM 插件时，这将有利于性能，因为更多的优化配置文件有助于 TensorRT 有更多机会选择更好的内核。注意：此功能会增加引擎构建时间，但预计不会有其他不利影响。

默认值: 'disable'

--paged_state

可选选项: enable, disable

启用分页状态，这有助于更有效地管理 RNN 状态的内存。

默认值: 'enable'

--streamingllm

可选选项: enable, disable

启用 [StreamingLLM](https://arxiv.org/abs/2309.17453)，它使用窗口注意力机制在长文本上执行高效且稳定的 LLM。

默认值: 'disable'

--use_fused_mlp

可选选项: enable, disable

在 Gated-MLP 中启用水平融合，该融合将两个 Matmul 操作组合成一个操作，然后是一个单独的 SwiGLU 内核。

默认值: 'enable'

--pp_reduce_scatter

可选选项: enable, disable

启用具有 ReduceScatter + AllGather 的流水线并行优化，针对大型 MoE 模型。

默认值: 'disable'