diff --git a/tensorrt_llm/_torch/models/checkpoints/hf/weight_loader.py b/tensorrt_llm/_torch/models/checkpoints/hf/weight_loader.py index fa6eef1a06d9..1c7d59b07dd0 100644 --- a/tensorrt_llm/_torch/models/checkpoints/hf/weight_loader.py +++ b/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 diff --git a/tests/unittest/_torch/models/checkpoints/hf/test_weight_loader.py b/tests/unittest/_torch/models/checkpoints/hf/test_weight_loader.py index 27defe3720be..ccf1973bb7f2 100644 --- a/tests/unittest/_torch/models/checkpoints/hf/test_weight_loader.py +++ b/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