tensorrt_llm_utils

用于TensorRT-LLM检查点导出的工具。

此文件中的一些逻辑是经验性的，如果出现异常，需要不断更新。

函数

convert_to_tensorrt_llm_config	转换为TensorRT-LLM检查点配置。
is_tensorrt_llm_0_8_or_9	如果 tensorrt_llm 版本是 0.8 或 0.9，则返回 true。
prepare_enc_dec_decoder_layer	为编码器-解码器模型的每个解码器层准备配置。
prepare_enc_dec_export_dir	为编码器-解码器模型准备导出目录。

convert_to_tensorrt_llm_config(model_config, weight_keys=['lm_head'])

转换为TensorRT-LLM检查点配置。

Parameters:

• model_config (ModelConfig) – 要转换的model_config。

• weight_keys (Iterable[str]) – 导出到tensorrt_llm检查点的权重字符串的可迭代对象。

is_tensorrt_llm_0_8_or_9()

如果 tensorrt_llm 版本是 0.8 或 0.9，则返回 true。

prepare_enc_dec_decoder_layer(layer_config, model_config, enc_dec, layers)

为编码器-解码器模型的每个解码器层准备配置。

Parameters:

• layer_config (DecoderLayerConfig) –

• model_config (T5Config) –

• enc_dec (str) –

• layers (List[DecoderLayerConfig]) –

prepare_enc_dec_export_dir(tensorrt_llm_config, export_root)

为编码器-解码器模型准备导出目录。

Parameters:

• tensorrt_llm_config (Dict[str, Any]) –

• export_root (Path) –