The ModelLoader is used to build an end-to-end model for a single-gpu. It accepts model name or a local model dir, and will download the model if necessary.
| 95 | |
| 96 | |
| 97 | class ModelLoader: |
| 98 | ''' The ModelLoader is used to build an end-to-end model for a single-gpu. |
| 99 | It accepts model name or a local model dir, and will download the model if necessary. |
| 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 |
| 140 | if isinstance(self.llm_args.model, Module): |
| 141 | # Build engine from user provided model |
| 142 | self._build_pipeline.append( |
| 143 | ("Build TensorRT LLM engine", |
| 144 | self._build_engine_from_inmemory_model)) |
| 145 | return |
| 146 | |
| 147 | if (self.model_obj.is_hub_model |
| 148 | and self._model_format is not _ModelFormatKind.TLLM_ENGINE): |
| 149 | # Download HF model if necessary |
| 150 | if self.model_obj.model_name is None: |
| 151 | raise ValueError( |
| 152 | "Either model_dir or model should be provided to ModelConfig." |
| 153 | ) |
| 154 | self._build_pipeline.append( |
no outgoing calls