(self,
llm_args: LlmArgs,
workspace: Optional[str | tempfile.TemporaryDirectory] = None,
llm_build_stats: Optional["LlmBuildStats"] = None)
| 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 |
nothing calls this directly
no test coverage detected