tensorrt_llm_utils

Utils for TensorRT-LLM checkpoint export.

Some of the logics in this file are empirical and needs constant update if exceptions occur.

Functions

convert_to_tensorrt_llm_config	Convert to TensorRT-LLM checkpoint config.
is_tensorrt_llm_0_8_or_9	Returns true if tensorrt_llm version is 0.8 or 0.9.
prepare_enc_dec_decoder_layer	Prepare the config for each decoder layer of encoder-decoder model.
prepare_enc_dec_export_dir	Prepare the export directory for encoder-decoder model.

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

Convert to TensorRT-LLM checkpoint config.

Parameters:

• model_config (ModelConfig) – The model_config to convert.

• weight_keys (Iterable[str]) – The iterable of string of weights exported to the tensorrt_llm checkpoint.

is_tensorrt_llm_0_8_or_9()

Returns true if tensorrt_llm version is 0.8 or 0.9.

prepare_enc_dec_decoder_layer(layer_config, model_config, enc_dec, layers)

Prepare the config for each decoder layer of encoder-decoder model.

Parameters:

• layer_config (DecoderLayerConfig) –

• model_config (T5Config) –

• enc_dec (str) –

• layers (List[DecoderLayerConfig]) –

prepare_enc_dec_export_dir(tensorrt_llm_config, export_root)

Prepare the export directory for encoder-decoder model.

Parameters:

• tensorrt_llm_config (Dict[str, Any]) –

• export_root (Path) –