Method from_checkpoint

tensorrt_llm/models/modeling_utils.py:726–762

(
        cls,
        ckpt_dir: str,
        rank: Optional[int] = None,
        config: Optional[PretrainedConfig] = None,
        *,
        preprocess_weights_hook: Optional[Callable[[Dict[str, Tensor]],
                                                   Dict[str, Tensor]]] = None)

    @classmethod
    def from_checkpoint(
        cls,
        ckpt_dir: str,
        rank: Optional[int] = None,
        config: Optional[PretrainedConfig] = None,
        *,
        preprocess_weights_hook: Optional[Callable[[Dict[str, Tensor]],
                                                   Dict[str, Tensor]]] = None):
        if config is None:
            config = PretrainedConfig.from_json_file(
                os.path.join(ckpt_dir, 'config.json'))

        if rank is not None:
            config.set_rank(rank)

        rank = config.mapping.rank
        if config.mapping.cp_size > 1:
            # cp_tp_pp rank -> tp_pp rank: because different cp ranks share the same ckpt.
            cp_size = config.mapping.cp_size
            # rank = pp_rank × tp_size × cp_size + tp_rank × cp_size + cp_rank.
            # rank // cp_size is equivalent to pp_rank × tp_size + tp_rank.
            rank = rank // cp_size
        weights_path = os.path.join(ckpt_dir, f'rank{rank}.safetensors')

        assert os.path.isfile(weights_path)
        weights = safetensors.torch.load_file(weights_path)
        is_checkpoint_pruned = getattr(config, 'is_pruned', False)

        if preprocess_weights_hook is not None:
            weights = preprocess_weights_hook(weights)

        weights = preprocess_weights(weights,
                                     config,
                                     from_pruned=is_checkpoint_pruned)
        model = cls(config)
        model.load(weights, from_pruned=is_checkpoint_pruned)
        return model

    def load(self, weights, from_pruned=False):
        required_names = set()