MCPcopy
hub / github.com/NVIDIA/TensorRT-LLM / load_from_ckpt

Method load_from_ckpt

tensorrt_llm/lora_manager.py:738–767  ·  view source on GitHub ↗

Returns the adapter UIDs that were loaded by this call. Note that when an adapter was already loaded before this call, it would not be included in the returned list of UIDs.

(
        self,
        model_dirs_or_files: List[str],
        model_config: Union["ModelConfig", LoraModelConfig],
        uids: Optional[List[str]] = None,
        ckpt_source: str = "hf",
    )

Source from the content-addressed store, hash-verified

736 return LoraManager.get_missing_qkv_modules(self.lora_target_modules)
737
738 def load_from_ckpt(
739 self,
740 model_dirs_or_files: List[str],
741 model_config: Union["ModelConfig", LoraModelConfig],
742 uids: Optional[List[str]] = None,
743 ckpt_source: str = "hf",
744 ) -> List[str]:
745 """Returns the adapter UIDs that were loaded by this call.
746
747 Note that when an adapter was already loaded before this call, it would not be
748 included in the returned list of UIDs.
749 """
750 if ckpt_source == "hf":
751 return self.load_from_hf(
752 model_dirs=model_dirs_or_files,
753 model_config=model_config,
754 uids=uids,
755 )
756 elif ckpt_source == "nemo":
757 # Find all .nemo files from directories or files
758 nemo_files = find_nemo_files(model_dirs_or_files)
759
760 # Pass the actual .nemo files to the loader
761 return self.load_from_nemo(
762 model_files=nemo_files,
763 model_config=model_config,
764 uids=uids,
765 )
766 else:
767 assert False, f"{self.__class__.__name__} does not support source {ckpt_source}"
768
769 def load_from_nemo(
770 self,

Callers 6

from_engineMethod · 0.95
from_dirMethod · 0.95
from_dirMethod · 0.95
add_request_peftMethod · 0.45
_load_lora_adapterMethod · 0.45
_load_prompt_adapterMethod · 0.45

Calls 3

load_from_hfMethod · 0.95
load_from_nemoMethod · 0.95
find_nemo_filesFunction · 0.85

Tested by

no test coverage detected