Load tensor from shards This function contains following steps: 1. Translate tllm_key into external key(s). 2. Load tensor/tensors partially according to layer attributes. 3. Call preprocess() if it is not None. 4. Call layer's post processing
(self,
tllm_key: str,
preprocess: Callable[[int], None] = None,
skip_tp: bool = False,
custom_postprocess_kwargs: dict = {})
| 228 | return res |
| 229 | |
| 230 | def load(self, |
| 231 | tllm_key: str, |
| 232 | preprocess: Callable[[int], None] = None, |
| 233 | skip_tp: bool = False, |
| 234 | custom_postprocess_kwargs: dict = {}): |
| 235 | """Load tensor from shards |
| 236 | |
| 237 | This function contains following steps: |
| 238 | 1. Translate tllm_key into external key(s). |
| 239 | 2. Load tensor/tensors partially according to layer attributes. |
| 240 | 3. Call preprocess() if it is not None. |
| 241 | 4. Call layer's post processing function. |
| 242 | 5. Return the dict for updating weight dict. |
| 243 | |
| 244 | Args: |
| 245 | tllm_key (str): TRT-LLM key from model iterators |
| 246 | preprocess (function, Optional): Customized preprocess function for step 3. |
| 247 | skip_tp (bool): Skip TP in case of the derived TP config is inappropriate. |
| 248 | """ |
| 249 | tp_rank = self.model.config.mapping.tp_rank |
| 250 | |
| 251 | sub_module = self.model |
| 252 | for attr in tllm_key.split(".")[:-1]: |
| 253 | sub_module = getattr(sub_module, attr) |
| 254 | param = self.model |
| 255 | for attr in tllm_key.split("."): |
| 256 | param = getattr(param, attr) |
| 257 | if param.is_buffer: |
| 258 | return {} |
| 259 | assert sub_module is not None and param is not None, f"{tllm_key} got Nonetype for parameter or parent module." |
| 260 | |
| 261 | tllm_to_externel_key_dict = getattr(sub_module, |
| 262 | "tllm_to_externel_key_dict", None) |
| 263 | tp_dim = getattr(sub_module, "tp_dim", -1) |
| 264 | require_weight_transpose = ( |
| 265 | isinstance(sub_module, WeightOnlyGroupwiseQuantColumnLinear) |
| 266 | or isinstance(sub_module, WeightOnlyGroupwiseQuantRowLinear)) |
| 267 | if tp_dim >= 0 and require_weight_transpose: |
| 268 | if sub_module.prequant_scaling_factor is not None: |
| 269 | if tllm_key.endswith("prequant_scaling_factor"): |
| 270 | tp_dim = 1 - tp_dim |
| 271 | elif tllm_key.endswith("weights_scaling_factor"): |
| 272 | tp_dim = -1 |
| 273 | elif tllm_key.endswith("weight"): |
| 274 | tp_dim = 1 - tp_dim |
| 275 | tp_size = getattr(sub_module, "tp_size", 1) |
| 276 | # Disable auto TP when num_kv_heads is invalid for split |
| 277 | if getattr(sub_module, "is_qkv", |
| 278 | False) and self.model.config.num_key_value_heads < tp_size: |
| 279 | tp_dim = -1 |
| 280 | tp_size = 1 |
| 281 | if skip_tp: |
| 282 | tp_dim = -1 |
| 283 | tp_size = 1 |
| 284 | if isinstance(sub_module, MOEWeightWrapper): |
| 285 | tp_rank = self.model.config.mapping.moe_tp_rank |
| 286 | external_key = self.translate_to_external_key( |
| 287 | tllm_key, tllm_to_externel_key_dict) |
no test coverage detected