[None][refactor] Table-drive the lazy-safetensors model-type check in HfWeightLoader

tensorrt_llm/_torch/models/checkpoints/hf/weight_loader.py

@@ -52,6 +52,9 @@
 _RANK_STRIPED_IO_POLICY = "rank_striped_read_ahead"
 _SUPPORTED_IO_POLICIES = (_NATIVE_IO_POLICY, _RANK_STRIPED_IO_POLICY)
 _SUPPORTED_REQUESTED_IO_POLICIES = (_AUTO_IO_POLICY, ) + _SUPPORTED_IO_POLICIES
+# Model families whose checkpoints are too large to materialize in host RAM;
+# their models stream rank-local slices out of the lazy mmapped handles.
+_LAZY_SAFETENSORS_MODEL_TYPES = ("kimi_k3", "kimi_linear")
 # Default to a single cached checkpoint: each entry pins a full copy of the
 # raw weights in CPU RAM, so callers wanting cross-model caching must opt in
 # via TRTLLM_HF_WEIGHT_CACHE_MAX_ENTRIES.
@@ -352,18 +355,25 @@ def _with_weight_cache(self,
         return weights
 
     @staticmethod
-    def _is_kimi_k3_checkpoint(checkpoint_dir: str) -> bool:
-        """Kimi K3 checkpoints (~1.5 TB) must not be materialized in host RAM."""
+    def _requires_lazy_safetensors(checkpoint_dir: str) -> bool:
+        """Whether this checkpoint must stay mmapped instead of being read
+        into host RAM.
+
+        The listed model families ship checkpoints too large to materialize
+        in host RAM (Kimi K3 is about 1.5 TB), and their models stream
+        rank-local slices (expert-parallel expert ranges) out of the lazy
+        handles during `load_weights`.
+        """
         config_path = os.path.join(checkpoint_dir, "config.json")
         if not os.path.isfile(config_path):
             return False
         # Do not swallow read/parse failures: every rank must take the same
-        # branch here (the non-Kimi path enqueues collectives), so a
-        # rank-local transient error routing one rank differently would
-        # deadlock the job. Propagating fails fast on all ranks instead.
+        # branch here (the eager path enqueues collectives), so a rank-local
+        # transient error routing one rank differently would deadlock the job.
+        # Propagating fails fast on all ranks instead.
         with open(config_path) as f:
             model_type = json.load(f).get("model_type")
-        return model_type in ("kimi_k3", "kimi_linear")
+        return model_type in _LAZY_SAFETENSORS_MODEL_TYPES
 
     def _load_lazy_safetensors(
             self,
@@ -687,6 +697,7 @@ def _start_rank_striped_read_ahead(
         active_communicator,
         **kwargs,
     ) -> tuple[dict[str, Any], RankStripedReadAheadSession | None]:
+        """Load weights via rank-striped read-ahead, or fall back to native."""
         node_communicator = None
         split_error = None
         try:
@@ -715,7 +726,7 @@ def _start_rank_striped_read_ahead(
         eligibility_reason = None
         preflight_error = None
         try:
-            if self._is_kimi_k3_checkpoint(checkpoint_dir):
+            if self._requires_lazy_safetensors(checkpoint_dir):
                 eligibility_reason = (
                     "the checkpoint requires model-specific lazy SafeTensors loading"
                 )
@@ -877,7 +888,8 @@ def _load_weights_native(self,
                              _local_communicator=None,
                              _allow_prefetch: bool = True,
                              **kwargs) -> dict[str, Any]:
-        if self._is_kimi_k3_checkpoint(checkpoint_dir):
+        """Load weights with the native (no read-ahead) I/O policy."""
+        if self._requires_lazy_safetensors(checkpoint_dir):
             return self._load_lazy_safetensors(checkpoint_dir, use_consolidated)
         weight_files = glob.glob(f"{checkpoint_dir}/*.safetensors")
         # Some model checkpoint directories contain not only the sharded safetensors, but one

tests/unittest/_torch/models/checkpoints/hf/test_weight_loader.py

@@ -13,15 +13,18 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import json
 import mmap
 import os
 import threading
+from pathlib import Path
 from unittest import mock
 
 import pytest
 
 from tensorrt_llm._torch.models.checkpoints import HfWeightLoader
 from tensorrt_llm._torch.models.checkpoints.base_weight_loader import ConsumableWeightsDict
+from tensorrt_llm._torch.models.checkpoints.hf.weight_loader import _LAZY_SAFETENSORS_MODEL_TYPES
 from tensorrt_llm.mapping import Mapping
 
 pytestmark = pytest.mark.cpu_only
@@ -502,7 +505,6 @@ def test_kimi_k3_lazy_load_records_the_checkpoint_dir(tmp_path):
     without this the model silently fell back to the shared mapping and the
     step was OOM-killed.
     """
-    import json
 
     import safetensors.torch
     import torch
@@ -519,3 +521,26 @@ def test_kimi_k3_lazy_load_records_the_checkpoint_dir(tmp_path):
         assert weights.checkpoint_dir == str(tmp_path)
     finally:
         loader.cleanup()
+
+
+@pytest.mark.parametrize("model_type", _LAZY_SAFETENSORS_MODEL_TYPES)
+def test_requires_lazy_safetensors_for_every_listed_model_type(
+    tmp_path: Path, model_type: str
+) -> None:
+    """Each model type in the table routes to the lazy (mmapped) load path."""
+    (tmp_path / "config.json").write_text(json.dumps({"model_type": model_type}))
+    assert HfWeightLoader._requires_lazy_safetensors(str(tmp_path)) is True
+
+
+@pytest.mark.parametrize("config", [{"model_type": "llama"}, {}])
+def test_requires_lazy_safetensors_is_false_for_other_checkpoints(
+    tmp_path: Path, config: dict[str, str]
+) -> None:
+    """Any other model type, or no model type at all, takes the eager path."""
+    (tmp_path / "config.json").write_text(json.dumps(config))
+    assert HfWeightLoader._requires_lazy_safetensors(str(tmp_path)) is False
+
+
+def test_requires_lazy_safetensors_is_false_without_a_config(tmp_path: Path) -> None:
+    """A directory without config.json cannot opt in to lazy loading."""
+    assert HfWeightLoader._requires_lazy_safetensors(str(tmp_path)) is False