MCPcopy Index your code
hub / github.com/NVIDIA/TensorRT-LLM / __init__

Method __init__

tensorrt_llm/llmapi/llm_utils.py:102–136  ·  view source on GitHub ↗
(self,
                 llm_args: LlmArgs,
                 workspace: Optional[str | tempfile.TemporaryDirectory] = None,
                 llm_build_stats: Optional["LlmBuildStats"] = None)

Source from the content-addressed store, hash-verified

100 '''
101
102 def __init__(self,
103 llm_args: LlmArgs,
104 workspace: Optional[str | tempfile.TemporaryDirectory] = None,
105 llm_build_stats: Optional["LlmBuildStats"] = None):
106 self.llm_args = llm_args
107 self._workspace = workspace or tempfile.TemporaryDirectory()
108 self.llm_build_stats = llm_build_stats or LlmBuildStats()
109
110 self.model_obj = _ModelWrapper(self.llm_args.model)
111 self.speculative_model_obj = _ModelWrapper(
112 self.llm_args.speculative_model
113 ) if self.llm_args.speculative_model is not None else None
114
115 if isinstance(self.llm_args, TrtLlmArgs):
116 self.convert_checkpoint_options = self.llm_args._convert_checkpoint_options
117 self.rank = mpi_rank()
118 self.global_rank = global_mpi_rank()
119 self.mapping = llm_args.parallel_config.to_mapping()
120
121 self._build_pipeline = []
122
123 # For model from hub, the _model_dir is None, and will updated once downloaded
124 self._model_dir: Optional[
125 Path] = self.model_obj.model_dir if self.model_obj.is_local_model else None
126
127 self._speculative_model_dir: Optional[
128 Path] = self.speculative_model_obj.model_dir if self.speculative_model_obj is not None and self.speculative_model_obj.is_local_model else None
129 self._model_info: Optional[_ModelInfo] = None
130 self._model_format = self.llm_args.model_format
131
132 if isinstance(self.llm_args, TrtLlmArgs):
133 assert self.llm_args.build_config
134 self.build_config = self.llm_args.build_config
135
136 self._gather_build_steps()
137
138 def _gather_build_steps(self):
139 # Prepare the model processing pipeline

Callers

nothing calls this directly

Calls 6

_gather_build_stepsMethod · 0.95
LlmBuildStatsClass · 0.85
_ModelWrapperClass · 0.85
global_mpi_rankFunction · 0.85
to_mappingMethod · 0.80
mpi_rankFunction · 0.50

Tested by

no test coverage detected