🌐 US-Proxy
class="logged-out env-production page-responsive" style="word-wrap: break-word;" >
Skip to content

Commit 6c87b6a

Browse files
moraxumoraxu
committed
[None][refactor] Table-drive the lazy-safetensors model-type check in HfWeightLoader
`HfWeightLoader._is_kimi_k3_checkpoint` decided, from `config.json`'s `model_type`, whether a checkpoint must stay mmapped instead of being materialized in host RAM, and hard-coded the two Kimi model types in the method body. Rename it to `_requires_lazy_safetensors` and move the model types into a module-level `_LAZY_SAFETENSORS_MODEL_TYPES` tuple, so the next model family whose checkpoint is too large for host RAM opts in by adding one entry instead of editing a Kimi-named predicate. Behavior is unchanged: the tuple holds exactly the two model types the method used to name. Add unit tests that exercise the predicate for every listed model type, for a non-lazy model type, for a config without `model_type`, and for a directory without `config.json`. Signed-off-by: Michal Guzek <mguzek@nvidia.com>
1 parent 01dc618 commit 6c87b6a

2 files changed

Lines changed: 40 additions & 8 deletions

File tree

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

Lines changed: 18 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -52,6 +52,9 @@
5252
_RANK_STRIPED_IO_POLICY = "rank_striped_read_ahead"
5353
_SUPPORTED_IO_POLICIES = (_NATIVE_IO_POLICY, _RANK_STRIPED_IO_POLICY)
5454
_SUPPORTED_REQUESTED_IO_POLICIES = (_AUTO_IO_POLICY, ) + _SUPPORTED_IO_POLICIES
55+
# Model families whose checkpoints are too large to materialize in host RAM;
56+
# their models stream rank-local slices out of the lazy mmapped handles.
57+
_LAZY_SAFETENSORS_MODEL_TYPES = ("kimi_k3", "kimi_linear")
5558
# Default to a single cached checkpoint: each entry pins a full copy of the
5659
# raw weights in CPU RAM, so callers wanting cross-model caching must opt in
5760
# via TRTLLM_HF_WEIGHT_CACHE_MAX_ENTRIES.
@@ -352,18 +355,25 @@ def _with_weight_cache(self,
352355
return weights
353356

354357
@staticmethod
355-
def _is_kimi_k3_checkpoint(checkpoint_dir: str) -> bool:
356-
"""Kimi K3 checkpoints (~1.5 TB) must not be materialized in host RAM."""
358+
def _requires_lazy_safetensors(checkpoint_dir: str) -> bool:
359+
"""Whether this checkpoint must stay mmapped instead of being read
360+
into host RAM.
361+
362+
The listed model families ship checkpoints too large to materialize
363+
in host RAM (Kimi K3 is about 1.5 TB), and their models stream
364+
rank-local slices (expert-parallel expert ranges) out of the lazy
365+
handles during ``load_weights``.
366+
"""
357367
config_path = os.path.join(checkpoint_dir, "config.json")
358368
if not os.path.isfile(config_path):
359369
return False
360370
# Do not swallow read/parse failures: every rank must take the same
361-
# branch here (the non-Kimi path enqueues collectives), so a
362-
# rank-local transient error routing one rank differently would
363-
# deadlock the job. Propagating fails fast on all ranks instead.
371+
# branch here (the eager path enqueues collectives), so a rank-local
372+
# transient error routing one rank differently would deadlock the job.
373+
# Propagating fails fast on all ranks instead.
364374
with open(config_path) as f:
365375
model_type = json.load(f).get("model_type")
366-
return model_type in ("kimi_k3", "kimi_linear")
376+
return model_type in _LAZY_SAFETENSORS_MODEL_TYPES
367377

368378
def _load_lazy_safetensors(
369379
self,
@@ -715,7 +725,7 @@ def _start_rank_striped_read_ahead(
715725
eligibility_reason = None
716726
preflight_error = None
717727
try:
718-
if self._is_kimi_k3_checkpoint(checkpoint_dir):
728+
if self._requires_lazy_safetensors(checkpoint_dir):
719729
eligibility_reason = (
720730
"the checkpoint requires model-specific lazy SafeTensors loading"
721731
)
@@ -877,7 +887,7 @@ def _load_weights_native(self,
877887
_local_communicator=None,
878888
_allow_prefetch: bool = True,
879889
**kwargs) -> dict[str, Any]:
880-
if self._is_kimi_k3_checkpoint(checkpoint_dir):
890+
if self._requires_lazy_safetensors(checkpoint_dir):
881891
return self._load_lazy_safetensors(checkpoint_dir, use_consolidated)
882892
weight_files = glob.glob(f"{checkpoint_dir}/*.safetensors")
883893
# Some model checkpoint directories contain not only the sharded safetensors, but one

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

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@
2222

2323
from tensorrt_llm._torch.models.checkpoints import HfWeightLoader
2424
from tensorrt_llm._torch.models.checkpoints.base_weight_loader import ConsumableWeightsDict
25+
from tensorrt_llm._torch.models.checkpoints.hf.weight_loader import _LAZY_SAFETENSORS_MODEL_TYPES
2526
from tensorrt_llm.mapping import Mapping
2627

2728
pytestmark = pytest.mark.cpu_only
@@ -519,3 +520,24 @@ def test_kimi_k3_lazy_load_records_the_checkpoint_dir(tmp_path):
519520
assert weights.checkpoint_dir == str(tmp_path)
520521
finally:
521522
loader.cleanup()
523+
524+
525+
@pytest.mark.parametrize("model_type", _LAZY_SAFETENSORS_MODEL_TYPES)
526+
def test_requires_lazy_safetensors_for_every_listed_model_type(tmp_path, model_type):
527+
"""Each model type in the table routes to the lazy (mmapped) load path."""
528+
import json
529+
530+
(tmp_path / "config.json").write_text(json.dumps({"model_type": model_type}))
531+
assert HfWeightLoader._requires_lazy_safetensors(str(tmp_path)) is True
532+
533+
534+
@pytest.mark.parametrize("config", [{"model_type": "llama"}, {}])
535+
def test_requires_lazy_safetensors_is_false_for_other_checkpoints(tmp_path, config):
536+
import json
537+
538+
(tmp_path / "config.json").write_text(json.dumps(config))
539+
assert HfWeightLoader._requires_lazy_safetensors(str(tmp_path)) is False
540+
541+
542+
def test_requires_lazy_safetensors_is_false_without_a_config(tmp_path):
543+
assert HfWeightLoader._requires_lazy_safetensors(str(tmp_path)) is False

0 commit comments

Comments
 (0)