avoid import all models when import trtllm

Author: chang-l

tensorrt_llm/_torch/__init__.py

@@ -1,5 +1,4 @@
 from .llm import LLM
 from .model_config import MoeLoadBalancerConfig
-from .models.checkpoints.base_checkpoint_loader import BaseCheckpointLoader
 
-__all__ = ["LLM", "MoeLoadBalancerConfig", "BaseCheckpointLoader"]
+__all__ = ["LLM", "MoeLoadBalancerConfig"]

tensorrt_llm/llmapi/llm_args.py

@@ -1908,7 +1908,7 @@ class TorchLlmArgs(BaseLlmArgs):
         default=None,
         description="The checkpoint loader to use for this LLM instance.",
         json_schema_extra={
-            "type": "Optional[tensorrt_llm._torch.BaseCheckpointLoader]"
+            "type": "Optional[tensorrt_llm._torch.models.checkpoints.BaseCheckpointLoader]"
         },
     )
 

tests/unittest/api_stability/references/llm.yaml

@@ -67,7 +67,7 @@ methods:
         annotation: Optional[tensorrt_llm.llmapi.llm_args.CudaGraphConfig]
         default: null
       checkpoint_loader:
-        annotation: Optional[tensorrt_llm._torch.BaseCheckpointLoader]
+        annotation: Optional[tensorrt_llm._torch.models.checkpoints.BaseCheckpointLoader]
         default: null
       checkpoint_format:
         annotation: Optional[str]