diff --git a/docs/source/developer-guide/telemetry.md b/docs/source/developer-guide/telemetry.md index 9a67bfa2c5dd..ad822e040273 100644 --- a/docs/source/developer-guide/telemetry.md +++ b/docs/source/developer-guide/telemetry.md @@ -153,7 +153,7 @@ unset or when the safety sanitizer rejects the runtime value. | `kv_connector_config.connector` | `Optional[str]` | `categorical` | allowlist | `lmcache`, `lmcache-mp`, `kvbm` | | `layer_wise_benchmarks_config.calibration_layer_indices` | `Optional[List[int]]` | `value` | | | | `layer_wise_benchmarks_config.calibration_mode` | `Literal['NONE', 'MARK', 'COLLECT']` | `categorical` | | `NONE`, `MARK`, `COLLECT` | -| `load_format` | `Union[str, tensorrt_llm.llmapi.llm_args.LoadFormat]` | `categorical` | allowlist | `auto`, `dummy`, `vision_only`, `gms` | +| `load_format` | `Union[str, tensorrt_llm.llmapi.llm_args.LoadFormat]` | `categorical` | allowlist | `auto`, `dummy`, `vision_only`, `gms`, `lazy_safetensors` | | `lora_config.lora_ckpt_source` | `Literal['hf', 'nemo']` | `categorical` | | `hf`, `nemo` | | `lora_config.max_cpu_loras` | `Optional[int]` | `value` | | | | `lora_config.max_lora_rank` | `` | `value` | | | diff --git a/tensorrt_llm/_torch/models/checkpoints/hf/checkpoint_loader.py b/tensorrt_llm/_torch/models/checkpoints/hf/checkpoint_loader.py index c495230779cc..0e626bd2ccd3 100644 --- a/tensorrt_llm/_torch/models/checkpoints/hf/checkpoint_loader.py +++ b/tensorrt_llm/_torch/models/checkpoints/hf/checkpoint_loader.py @@ -91,6 +91,7 @@ def build_checkpoint_catalog(self, checkpoint_dir: str, return self.weight_loader.build_checkpoint_catalog( checkpoint_dir, use_consolidated=kwargs.get("use_consolidated", False), + load_lazily=kwargs.get("load_lazily", False), ) @property diff --git a/tensorrt_llm/_torch/models/checkpoints/hf/weight_loader.py b/tensorrt_llm/_torch/models/checkpoints/hf/weight_loader.py index 3eeae6ceb74a..ea7b247d50ac 100644 --- a/tensorrt_llm/_torch/models/checkpoints/hf/weight_loader.py +++ b/tensorrt_llm/_torch/models/checkpoints/hf/weight_loader.py @@ -13,7 +13,6 @@ # See the License for the specific language governing permissions and # limitations under the License. import glob -import json import multiprocessing import os import threading @@ -56,9 +55,6 @@ _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. @@ -360,27 +356,6 @@ def _with_weight_cache(self, self._cache_loaded_weights(cache_key, weights) return weights - @staticmethod - 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 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 _LAZY_SAFETENSORS_MODEL_TYPES - def _load_lazy_safetensors( self, checkpoint_dir: str, @@ -390,7 +365,7 @@ def _load_lazy_safetensors( Values are ``safetensors`` PySafeSlice objects: ``v[:]`` (or any indexing) materializes only the requested bytes from the mmapped file. This lets a model's ``load_weights`` stream a huge checkpoint - and read only its rank-local shard (e.g. Kimi K3 expert-parallel + and read only its rank-local shard (e.g. expert-parallel routed expert slices) without ever holding the full checkpoint in RAM. """ weight_files = sorted(glob.glob(f"{checkpoint_dir}/*.safetensors")) @@ -421,7 +396,7 @@ def _load_lazy_safetensors( # transient source weights when ModelLoader.load() returns. lazy_weights = _LazySafetensorsWeights(weights, handles) # A lazy slice does not carry the file it came from, and a model that - # wants to re-open shards itself (Kimi K3 streams rank-local experts + # wants to re-open shards itself (e.g. streaming rank-local experts # per shard file, precisely to avoid holding this mapping open) has no # other reliable source: transformers no longer sets # ``PretrainedConfig._name_or_path``. @@ -432,8 +407,12 @@ def load_weights(self, checkpoint_dir: str, mapping: Mapping, use_consolidated: bool = False, + load_lazily: bool = False, **kwargs) -> dict[str, Any]: - """Load synchronously without activating session-scoped read-ahead.""" + """Load synchronously without activating session-scoped read-ahead. + + See :meth:`_load_weights_native` for what ``load_lazily`` selects. + """ self._reset_checkpoint_io_status() if self._checkpoint_io_policy != _NATIVE_IO_POLICY: status = self._last_checkpoint_io_status @@ -446,8 +425,11 @@ def load_weights(self, f"load: requested={status.requested}, " f"selected={status.selected}, reason={status.fallback_reason}.") logger.info(message) - weights = self._load_weights_native(checkpoint_dir, mapping, - use_consolidated, **kwargs) + weights = self._load_weights_native(checkpoint_dir, + mapping, + use_consolidated, + load_lazily=load_lazily, + **kwargs) self._last_checkpoint_io_status.effective = _NATIVE_IO_POLICY self._log_checkpoint_io_status() return weights @@ -457,12 +439,21 @@ def open_weight_session(self, checkpoint_dir: str, mapping: Mapping, use_consolidated: bool = False, + load_lazily: bool = False, **kwargs) -> Iterator[dict[str, Any]]: - """Keep opt-in read-ahead alive through model materialization.""" + """Keep opt-in read-ahead alive through model materialization. + + See :meth:`_load_weights_native` for what ``load_lazily`` selects. A + lazy load never reads whole shards up front, so it is not eligible for + read-ahead; the flag is forwarded so every path below agrees on it. + """ self._reset_checkpoint_io_status() if self._checkpoint_io_policy == _NATIVE_IO_POLICY: - weights = self._load_weights_native(checkpoint_dir, mapping, - use_consolidated, **kwargs) + weights = self._load_weights_native(checkpoint_dir, + mapping, + use_consolidated, + load_lazily=load_lazily, + **kwargs) self._last_checkpoint_io_status.effective = _NATIVE_IO_POLICY self._log_checkpoint_io_status() yield weights @@ -492,6 +483,7 @@ def open_weight_session(self, str(coordinated_mapping_error), active_communicator=active_communicator, allow_native_prefetch=False, + load_lazily=load_lazily, **kwargs, ) yield weights @@ -506,6 +498,7 @@ def open_weight_session(self, # Ray workers intentionally disable MPI. Preserve the native # path's per-process read-ahead in that supported mode. allow_native_prefetch=mpi_disabled(), + load_lazily=load_lazily, **kwargs, ) yield weights @@ -516,6 +509,7 @@ def open_weight_session(self, mapping, use_consolidated, active_communicator, + load_lazily=load_lazily, **kwargs, ) if session is None: @@ -583,9 +577,10 @@ def _selected_safetensors_files(checkpoint_dir: str, def build_checkpoint_catalog( self, checkpoint_dir: str, - use_consolidated: bool = False) -> CheckpointCatalog | None: + use_consolidated: bool = False, + load_lazily: bool = False) -> CheckpointCatalog | None: """Inspect the selected eager SafeTensors files without reading payloads.""" - if self._requires_lazy_safetensors(checkpoint_dir): + if load_lazily: return None if (self._partial_model_loading or int(os.environ.get("TLLM_OVERRIDE_LAYER_NUM", "0")) != 0): @@ -629,6 +624,7 @@ def _fallback_to_native(self, session=None, *, allow_native_prefetch: bool | None = None, + load_lazily: bool = False, **kwargs) -> dict[str, Any]: status = self._last_checkpoint_io_status status.activated = False @@ -671,6 +667,7 @@ def load_native() -> dict[str, Any]: checkpoint_dir, mapping, use_consolidated, + load_lazily=load_lazily, _local_communicator=fallback_communicator, _allow_prefetch=allow_native_prefetch, **kwargs) @@ -722,6 +719,7 @@ def _start_rank_striped_read_ahead( mapping: Mapping, use_consolidated: bool, active_communicator, + load_lazily: bool = False, **kwargs, ) -> tuple[dict[str, Any], RankStripedReadAheadSession | None]: """Load weights via rank-striped read-ahead, or fall back to native.""" @@ -742,10 +740,13 @@ def _start_rank_striped_read_ahead( None, node_communicator, active_communicator, "rank-striped node communicator cleanup") self._last_checkpoint_io_status.selected = _NATIVE_IO_POLICY - return self._fallback_to_native(checkpoint_dir, mapping, + return self._fallback_to_native(checkpoint_dir, + mapping, use_consolidated, str(coordinated_split_error), - active_communicator, **kwargs), None + active_communicator, + load_lazily=load_lazily, + **kwargs), None weight_files = [] stats = [] @@ -753,10 +754,9 @@ def _start_rank_striped_read_ahead( eligibility_reason = None preflight_error = None try: - if self._requires_lazy_safetensors(checkpoint_dir): + if load_lazily: eligibility_reason = ( - "the checkpoint requires model-specific lazy SafeTensors loading" - ) + "the checkpoint is loaded as lazy SafeTensors slices") weight_files = self._selected_safetensors_files( checkpoint_dir, use_consolidated) stats = [(path, os.stat(path)) for path in weight_files] @@ -776,11 +776,14 @@ def _start_rank_striped_read_ahead( active_communicator, "rank-striped preflight", preflight_error) if coordinated_preflight_error is not None: self._last_checkpoint_io_status.selected = _NATIVE_IO_POLICY - return self._fallback_to_native(checkpoint_dir, mapping, + return self._fallback_to_native(checkpoint_dir, + mapping, use_consolidated, str(coordinated_preflight_error), active_communicator, - node_communicator, **kwargs), None + node_communicator, + load_lazily=load_lazily, + **kwargs), None local_rank = 0 local_size = 1 @@ -810,10 +813,13 @@ def _start_rank_striped_read_ahead( None, node_communicator, active_communicator, "rank-striped node preflight cleanup") self._last_checkpoint_io_status.selected = _NATIVE_IO_POLICY - return self._fallback_to_native(checkpoint_dir, mapping, + return self._fallback_to_native(checkpoint_dir, + mapping, use_consolidated, str(coordinated_node_error), - active_communicator, **kwargs), None + active_communicator, + load_lazily=load_lazily, + **kwargs), None file_sizes = [(path, stat.st_size) for path, stat in stats] checkpoint_bytes = sum(size for _, size in file_sizes) @@ -845,11 +851,14 @@ def _start_rank_striped_read_ahead( if fallback_reasons: rank, reason = fallback_reasons[0] self._last_checkpoint_io_status.selected = _NATIVE_IO_POLICY - return self._fallback_to_native(checkpoint_dir, mapping, + return self._fallback_to_native(checkpoint_dir, + mapping, use_consolidated, f"rank {rank}: {reason}", active_communicator, - node_communicator, **kwargs), None + node_communicator, + load_lazily=load_lazily, + **kwargs), None self._last_checkpoint_io_status.selected = _RANK_STRIPED_IO_POLICY session = None @@ -865,11 +874,14 @@ def _start_rank_striped_read_ahead( "rank-striped reader setup", setup_error) if coordinated_setup_error is not None: - return self._fallback_to_native(checkpoint_dir, mapping, + return self._fallback_to_native(checkpoint_dir, + mapping, use_consolidated, str(coordinated_setup_error), active_communicator, - node_communicator, session, + node_communicator, + session, + load_lazily=load_lazily, **kwargs), None assert session is not None @@ -912,11 +924,27 @@ def _load_weights_native(self, mapping: Mapping, use_consolidated: bool = False, *, + load_lazily: bool = False, _local_communicator=None, _allow_prefetch: bool = True, **kwargs) -> dict[str, Any]: - """Load weights with the native (no read-ahead) I/O policy.""" - if self._requires_lazy_safetensors(checkpoint_dir): + """Load weights with the native (no read-ahead) I/O policy. + + When ``load_lazily`` is set (selected via ``LoadFormat.LAZY_SAFETENSORS`` + on the ``load_format`` surface, or a model's declared default), the + safetensors shards are opened lazily so only the rank-local slices each + model reads are materialized -- never the whole checkpoint in CPU + memory. This is required for checkpoints too large to fit in host RAM. + + Otherwise the checkpoint is loaded eagerly: safetensors shards may be + prefetched in parallel to warm up the OS file cache if the CPU memory is + large enough, before their tensors are loaded via mmap. When + `_WEIGHT_CACHE_ENV` is on, eager loads can also use a CPU weight cache to + accelerate repeated loading under the same process. + + Returns a `ConsumableWeightsDict` mapping checkpoint tensor names to tensors. + """ + if load_lazily: 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/tensorrt_llm/_torch/models/kimi_k3_knobs.py b/tensorrt_llm/_torch/models/kimi_k3_knobs.py new file mode 100644 index 000000000000..4d3ee870b638 --- /dev/null +++ b/tensorrt_llm/_torch/models/kimi_k3_knobs.py @@ -0,0 +1,176 @@ +# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Resolution of the Kimi K3 FP8 weight-read knobs. + +These knobs decide whether a replicated K3 projection is read from an FP8 +(e4m3, 128x128 block-scale) copy of its weights instead of BF16. They are +consumed on the checkpoint-loading path — ``load_weights`` keeps the FP8 +checkpoint pairs only when the read is enabled, and the post-load conversion +swaps the modules — so they live on +:class:`~tensorrt_llm.models.modeling_utils.QuantConfig`. + +They used to be read straight from the ``KIMI_K3_*`` environment variables +inside :mod:`modeling_kimi_linear`. Resolution precedence for every knob: + +1. an explicit config value (not ``None``) wins; +2. else the deprecated environment variable, if set, is honored (emitting a + one-time deprecation warning); otherwise +3. the historical default (some are computed at runtime from the arch / + parallelism). + +This keeps the shipped env-var behavior working for back-compat while making the +config surface authoritative. Kept intentionally lightweight (no torch model +imports) so the resolution logic is unit-testable on CPU. +""" + +from __future__ import annotations + +import os +from dataclasses import dataclass +from typing import Any, Callable, Optional + +from ..._utils import is_sm_100f +from ...logger import logger + +# Deprecated env var -> the config path that now owns the knob. Used only to +# make the one-time deprecation warning actionable. +_ENV_TO_CONFIG_PATH = { + "KIMI_K3_FP8_WEIGHT_READ": "quant_config.kimi_k3_fp8_weight_read", + "KIMI_K3_FP8_WEIGHT_READ_KDA": "quant_config.kimi_k3_fp8_weight_read_kda", + "KIMI_K3_FP8_WEIGHT_READ_MLA": "quant_config.kimi_k3_fp8_weight_read_mla", + "KIMI_K3_FP8_WEIGHT_READ_GATE_UP": "quant_config.kimi_k3_fp8_weight_read_gate_up", + "KIMI_K3_KDA_GLUE_FP8": "quant_config.kimi_k3_kda_glue_fp8", +} + +# FP8 weight-read knob field names on ``QuantConfig``. These are carried from the +# user-facing ``llm_args.quant_config`` onto the checkpoint-derived +# ``model_config.quant_config`` by :func:`carry_user_quant_knobs`. +KIMI_K3_QUANT_KNOB_FIELDS = ( + "kimi_k3_fp8_weight_read", + "kimi_k3_fp8_weight_read_kda", + "kimi_k3_fp8_weight_read_mla", + "kimi_k3_fp8_weight_read_gate_up", + "kimi_k3_kda_glue_fp8", +) + + +def _resolve( + config_value: Optional[Any], env_name: str, env_parser: Callable[[str], Any], default: Any +) -> Any: + """Resolve one knob: config wins; else deprecated env (warn-once); else default. + + ``config_value`` is ``None`` when the knob was not set on the config surface. + Whenever the deprecated env var is present a one-time warning is emitted, + even if a config value overrides it, so users learn to migrate. + """ + env_raw = os.environ.get(env_name) + if env_raw is not None: + logger.warning_once( + f"Environment variable '{env_name}' is deprecated and will be " + f"removed; set '{_ENV_TO_CONFIG_PATH[env_name]}' on the config " + f"surface instead (via extra_llm_api_options). It is still honored " + f"for now, but the config value takes precedence when both are set.", + key=f"kimi_k3_deprecated_env::{env_name}", + ) + if config_value is not None: + return config_value + if env_raw is not None: + return env_parser(env_raw) + return default + + +def _knob(config: Optional[Any], name: str) -> Optional[Any]: + """Read ``name`` off a (possibly ``None``) config object; ``None`` if absent.""" + if config is None: + return None + return getattr(config, name, None) + + +@dataclass(frozen=True) +class Fp8WeightReadGates: + """Resolved FP8 weight-read gates. + + ``master`` folds in the ``is_sm_100f()`` arch gate: it is ``False`` off + Blackwell regardless of the requested value. The sub-gates only ever narrow + an enabled master, so all of them are ``False`` when ``master`` is ``False``. + """ + + master: bool + kda: bool + kda_glue: bool + mla: bool + gate_up: bool + + +def resolve_fp8_weight_read_gates( + quant_config: Optional[Any], *, enable_attention_dp: bool +) -> Fp8WeightReadGates: + """Resolve the 5 FP8 weight-read gates from ``quant_config`` (+ deprecated env). + + The master switch is opt-in (FP8 weight reads are lossy relative to BF16, so + a default run keeps BF16 and matches the published accuracy numbers) and is + additionally gated on ``is_sm_100f()`` — the DeepGEMM ``fp8_swap_ab_gemm`` + kernel is Blackwell-only. The KDA / KDA-glue / MLA / gate-up sub-gates only + narrow an already-enabled master and stay default-on. + """ + master_requested = _resolve( + _knob(quant_config, "kimi_k3_fp8_weight_read"), + "KIMI_K3_FP8_WEIGHT_READ", + lambda s: s not in ("", "0"), + False, + ) + master = bool(is_sm_100f() and master_requested) + + kda = master and bool( + _resolve( + _knob(quant_config, "kimi_k3_fp8_weight_read_kda"), + "KIMI_K3_FP8_WEIGHT_READ_KDA", + lambda s: s != "0", + True, + ) + ) + kda_glue = kda and bool( + _resolve( + _knob(quant_config, "kimi_k3_kda_glue_fp8"), + "KIMI_K3_KDA_GLUE_FP8", + lambda s: s != "0", + True, + ) + ) + mla = master and bool( + _resolve( + _knob(quant_config, "kimi_k3_fp8_weight_read_mla"), + "KIMI_K3_FP8_WEIGHT_READ_MLA", + lambda s: s != "0", + True, + ) + ) + gate_up = master and bool( + _resolve( + _knob(quant_config, "kimi_k3_fp8_weight_read_gate_up"), + "KIMI_K3_FP8_WEIGHT_READ_GATE_UP", + lambda s: s != "0", + enable_attention_dp, + ) + ) + return Fp8WeightReadGates(master=master, kda=kda, kda_glue=kda_glue, mla=mla, gate_up=gate_up) + + +def carry_user_quant_knobs( + src_quant_config: Optional[Any], dst_quant_config: Optional[Any] +) -> None: + """Copy the K3 FP8 weight-read knobs from the user's ``quant_config`` + (``llm_args.quant_config``) onto the checkpoint-derived + ``model_config.quant_config``. + + ``model_config.quant_config`` is built from the checkpoint's + ``hf_quant_config.json`` and does not carry the user's ``extra_llm_api_options`` + ``quant_config`` values, so the K3 FP8-read knobs are threaded across here. + A no-op for knobs the user did not set (``None``) and for non-K3 runs. + """ + if src_quant_config is None or dst_quant_config is None: + return + for name in KIMI_K3_QUANT_KNOB_FIELDS: + value = getattr(src_quant_config, name, None) + if value is not None: + setattr(dst_quant_config, name, value) diff --git a/tensorrt_llm/_torch/models/modeling_kimi_k3_vl.py b/tensorrt_llm/_torch/models/modeling_kimi_k3_vl.py index d00236574b83..edad8a749df2 100644 --- a/tensorrt_llm/_torch/models/modeling_kimi_k3_vl.py +++ b/tensorrt_llm/_torch/models/modeling_kimi_k3_vl.py @@ -37,7 +37,7 @@ """ import copy -from typing import Optional, Tuple +from typing import TYPE_CHECKING, Dict, Optional, Tuple import torch import torch.nn as nn @@ -77,6 +77,9 @@ register_vision_encoder, ) +if TYPE_CHECKING: + from tensorrt_llm.llmapi.llm_args import TorchLlmArgs + # --------------------------------------------------------------------------- # Native MoonViT3d Vision Encoder Components (K3 deltas) # --------------------------------------------------------------------------- @@ -460,6 +463,18 @@ class KimiK3ForConditionalGeneration(KimiK25ForConditionalGeneration): _VISION_MODEL_CLS = KimiK3VisionModel mamba_metadata_cls = KimiLinearForCausalLM.mamba_metadata_cls + @classmethod + def get_model_defaults(cls, llm_args: "TorchLlmArgs") -> Dict[str, str]: + """Default this model to the lazy safetensors load format. + + The K3 checkpoint (~1.5 TB) must be streamed shard-by-shard, so K3 + declares the lazy ``LoadFormat`` as its default now that the shared + loader no longer sniffs K3 by model type. This wrapper subclasses HF + ``PreTrainedModel``, not the TRT-LLM base, so there is no inherited + ``get_model_defaults`` to extend. A user-set ``load_format`` still wins. + """ + return {"load_format": "lazy_safetensors"} + def __init__( self, model_config: ModelConfig[PretrainedConfig], diff --git a/tensorrt_llm/_torch/models/modeling_kimi_linear.py b/tensorrt_llm/_torch/models/modeling_kimi_linear.py index 38dad076d57f..64988e42d50d 100644 --- a/tensorrt_llm/_torch/models/modeling_kimi_linear.py +++ b/tensorrt_llm/_torch/models/modeling_kimi_linear.py @@ -103,7 +103,6 @@ from safetensors import safe_open from torch import nn -from ..._utils import is_sm_100f from ...logger import logger from ...mapping import Mapping from ...models.modeling_utils import QuantAlgo, QuantConfig @@ -121,6 +120,7 @@ from ..moe.fused_moe import ConfigurableMoE, SiTuActivation, TRTLLMGenFusedMoE, create_moe from ..moe.fused_moe.routing import DeepSeekV3MoeRoutingMethod from ..utils import AuxStreamType +from .kimi_k3_knobs import resolve_fp8_weight_read_gates from .modeling_speculative import SpecDecOneEngineForCausalLM from .modeling_utils import DecoderModel, register_auto_model, run_concurrently @@ -142,61 +142,14 @@ # bounds. KIMI_K3_MLA_MAX_POSITIONS overrides the size for short-context # deployments. _KIMI_K3_MLA_MAX_POSITIONS_ENV = "KIMI_K3_MLA_MAX_POSITIONS" +# The FP8 block-scale weight-read knobs (master + KDA/MLA/gate_up/KDA-glue +# sub-gates) live on ``QuantConfig`` and are resolved in ``kimi_k3_knobs`` +# (config value > deprecated env, warn-once > historical default). _KIMI_K3_MLA_DERIVED_PARAM_SUFFIXES = ( ".self_attn.mixer.k_b_proj_trans", ".self_attn.mixer.v_b_proj", ) -# Serve the replicated MoE-layer MLP projections (shared-expert gate/up/down -# and the latent up/down projection) from an FP8 copy of their weights instead -# of BF16. Under attention data-parallelism every rank re-reads these dense -# weights in full on every decode step, so decode is bound by that HBM read; -# an FP8 (e4m3) weight with 128x128 block scales roughly halves those bytes. -# The MLA projections and the routed MXFP4 experts are left untouched (the KDA -# q/k/v/g/o projections have their own switch below). The FP8 weight read is -# lossy relative to BF16, so it is opt-in: set this to "1" to trade accuracy -# for decode bandwidth. Default "0" keeps BF16, which is what the published -# accuracy numbers are measured against. -_KIMI_K3_FP8_WEIGHT_READ_ENV = "KIMI_K3_FP8_WEIGHT_READ" - -# Also read the KDA linear-attention q/k/v/g/o projections at FP8 block-scale. -# These are the largest single replicated weight read (~61 GB/rank of the -# ~109 GB BF16 read per decode step). They use the same FP8 path as the MLP -# projections above but are gated separately: the recurrent linear-attention -# core is more accuracy-sensitive than the feed-forward MLPs, so set this to -# "0" to keep the KDA projections in BF16 while still reading the MLPs at FP8. -# The master KIMI_K3_FP8_WEIGHT_READ switch and the SM100 gate still apply. -_KIMI_K3_FP8_WEIGHT_READ_KDA_ENV = "KIMI_K3_FP8_WEIGHT_READ_KDA" - -# Also read the MLA (full-attention) q_a/q_b/o and output-gate projections at -# FP8 block-scale. These are the replicated attention weights the MLP pass and -# the KDA pass above leave in BF16, and they are re-read in full by every rank -# each decode step under attention data-parallelism. Two MLA projections are -# deliberately kept in BF16: kv_a_proj_with_mqa outputs kv_lora_rank + -# qk_rope_head_dim (576, not a multiple of 128, so no exact 128x128 block -# scale), and kv_b_proj's weight is consumed directly (not through its forward) -# by the absorbed-decode _kv_b_absorb_split to build the k/v absorb matrices, -# which has no FP8 dequant path. The master KIMI_K3_FP8_WEIGHT_READ switch and -# the SM100 gate still apply. -_KIMI_K3_FP8_WEIGHT_READ_MLA_ENV = "KIMI_K3_FP8_WEIGHT_READ_MLA" - -# Expert override (prototype): set to "0" to drop the -# KimiKDALinearAttention decode -# fast path — fused qkvg and [f_a | b] projections, persistent conv staging, -# and precomputed kernel-layout constants (``forward_decode``) — when the -# KDA projections are read at FP8 block-scale. With the fast path kept (the -# default on an enabled master), decode issues the loader's fused FP8 -# ``qkvg_proj`` GEMM for q/k/v/g plus one small BF16 GEMV for [f_a | b] -# (``finalize_decode_weights_fp8``), so FP8 weight storage and the decode -# glue savings coexist. Requires the FP8 KDA read to be active; no effect -# otherwise. Default on ("0" disables). -_KIMI_K3_KDA_GLUE_FP8_ENV = "KIMI_K3_KDA_GLUE_FP8" - -# FP8 read for the fused shared-expert gate_up_proj. Default follows the -# parallel layout (on under attention DP, off under TP — see the conversion -# helper's comment); set 0/1 to force either. -_KIMI_K3_FP8_WEIGHT_READ_GATE_UP_ENV = "KIMI_K3_FP8_WEIGHT_READ_GATE_UP" - class KimiK3MoEGate(nn.Module): """Kimi K3 gate weights and routing method for ``ConfigurableMoE``.""" @@ -290,23 +243,6 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: return self.weight * hidden_states_float.to(input_dtype) -def _resolve_fp8_weight_read_gates() -> tuple[bool, bool, bool]: - """Resolve the FP8 weight-read switches into (master, kda, kda_glue). - - The master switch is opt-in: FP8 weight reads are lossy relative to BF16, - so a default run keeps BF16 and matches the published accuracy numbers. - The KDA and KDA-glue switches only narrow an already-enabled master, so - they stay default-on and are inert while the master is off. - """ - fp8_weight_read = is_sm_100f() and os.environ.get(_KIMI_K3_FP8_WEIGHT_READ_ENV, "0") not in ( - "", - "0", - ) - kda_fp8 = fp8_weight_read and os.environ.get(_KIMI_K3_FP8_WEIGHT_READ_KDA_ENV, "1") != "0" - kda_glue_fp8 = kda_fp8 and os.environ.get(_KIMI_K3_KDA_GLUE_FP8_ENV, "1") != "0" - return fp8_weight_read, kda_fp8, kda_glue_fp8 - - def _resolve_kimi_situ_betas(cfg: Any) -> tuple[float, float]: """Return the finite SiTu betas required by the routed-expert kernels.""" config_situ_beta = getattr(cfg, "activation_situ_beta", None) @@ -2076,6 +2012,13 @@ def _setup_helix_mappings( @classmethod def get_model_defaults(cls, llm_args) -> dict: + # - load_format=lazy_safetensors: the K3 checkpoint (~1.5 TB) must be + # streamed shard-by-shard (rank-local slices only); eagerly loading + # it into host RAM OOM-kills the job. Declaring the lazy LoadFormat + # as the model default is how K3 opts into that path without any + # model-name check in the shared weight loader, and without the user + # having to set load_format. A user-set load_format still wins + # (apply_model_defaults_to_llm_args honors explicit overrides). # - enable_block_reuse defaults off: reuse is supported as an # explicit opt-in (routes to CppMambaHybridCacheManager with # per-block KDA state snapshots); the default stays on the @@ -2085,10 +2028,11 @@ def get_model_defaults(cls, llm_args) -> dict: # the fallback C++ path requires num_heads % 64 == 0, which K3's # 96 query heads violate. return { + "load_format": "lazy_safetensors", "kv_cache_config": { "enable_block_reuse": False, "tokens_per_block": 64, - } + }, } @classmethod @@ -2274,7 +2218,10 @@ def _load_trunk_params( ) # Keep each FP8_PB_WO checkpoint pair alongside the BF16 # parameter only when the later weight-read conversion consumes it. - stash_ckpt_fp8 = _resolve_fp8_weight_read_gates()[0] + stash_ckpt_fp8 = resolve_fp8_weight_read_gates( + self.model_config.quant_config, + enable_attention_dp=self.model_config.mapping.enable_attention_dp, + ).master # KDA head-shard (attention-DP off): rank r loads head rows/cols # [r*local : (r+1)*local] of every head-major KDA tensor. kda_tp_size, kda_tp_rank = 1, 0 @@ -2662,18 +2609,21 @@ def _finalize_weight_load(self, num_params: int, num_moe_layers: int) -> None: # the bf16 wrapper fast path (finalize_decode_weights) is NOT built: # both fuse the same projections and the wrapper path — checked first # at decode — would bypass the FP8 modules entirely, leaving the FP8 - # copies resident but inert. KIMI_K3_FP8_WEIGHT_READ_KDA=0 restores - # the bf16 wrapper fast path; KIMI_K3_KDA_GLUE_FP8=1 instead rebuilds - # the wrapper fast path on top of the FP8 modules after the + # copies resident but inert. quant_config.kimi_k3_fp8_weight_read_kda=0 + # restores the bf16 wrapper fast path; kimi_k3_kda_glue_fp8=1 instead + # rebuilds the wrapper fast path on top of the FP8 modules after the # conversion (finalize_decode_weights_fp8), so neither is traded away. - fp8_weight_read, kda_fp8, kda_glue_fp8 = _resolve_fp8_weight_read_gates() + gates = resolve_fp8_weight_read_gates( + self.model_config.quant_config, + enable_attention_dp=self.model_config.mapping.enable_attention_dp, + ) # Build the KDA fused projection views and decode kernel constants. # This must run after every KDA parameter is loaded and sharded. num_kda_fused = 0 for layer in self.model.layers: if getattr(layer, "is_kda", False) and _has_weights(layer): - if not kda_fp8: + if not gates.kda: layer.linear_attn.finalize_decode_weights() num_kda_fused += int( layer.linear_attn._qkvg_proj_weight is not None @@ -2696,14 +2646,10 @@ def _finalize_weight_load(self, num_params: int, num_moe_layers: int) -> None: # FP8 block-scale weight read for the replicated MoE-layer MLPs. The # DeepGEMM fp8_swap_ab_gemm kernel is Blackwell-only; keep BF16 on any # other SM or when explicitly disabled. - if fp8_weight_read: - gate_up_default = "1" if self.model_config.mapping.enable_attention_dp else "0" + if gates.master: n_fp8 = _convert_moe_mlps_to_fp8_weight_read( self.model, - include_fused_gate_up=os.environ.get( - _KIMI_K3_FP8_WEIGHT_READ_GATE_UP_ENV, gate_up_default - ) - != "0", + include_fused_gate_up=gates.gate_up, ) logger.info( f"Kimi K3: reading {n_fp8} MoE-layer MLP projections " @@ -2713,14 +2659,14 @@ def _finalize_weight_load(self, num_params: int, num_moe_layers: int) -> None: # weight read; convert them to the same FP8 block-scale read unless # kept in BF16 for accuracy (their own switch — the recurrent core # is the most precision-sensitive slice). - if os.environ.get(_KIMI_K3_FP8_WEIGHT_READ_KDA_ENV, "1") != "0": + if gates.kda: n_kda = _convert_kda_projections_to_fp8_weight_read(self.model) logger.info( f"Kimi K3: reading {n_kda} KDA q/k/v/g/o projections " f"at FP8 block-scale (q/k/v/g fused into one prefill/decode/verify GEMM " f"per layer)" ) - if kda_glue_fp8: + if gates.kda_glue: # Rebuild the fused projection path on top of the FP8 # modules. This must run after the conversion above so # fused FP8 qkvg_proj exists and only [f_a | b] is fused @@ -2739,7 +2685,7 @@ def _finalize_weight_load(self, num_params: int, num_moe_layers: int) -> None: # leave in BF16; convert them to the same FP8 block-scale read # (kv_a/kv_b stay BF16 — see the switch's comment) unless kept in # BF16 for accuracy. - if os.environ.get(_KIMI_K3_FP8_WEIGHT_READ_MLA_ENV, "1") != "0": + if gates.mla: n_mla = _convert_mla_projections_to_fp8_weight_read(self.model) logger.info( f"Kimi K3: reading {n_mla} MLA q_a/q_b/o/g projections at FP8 block-scale" diff --git a/tensorrt_llm/_torch/pyexecutor/model_loader.py b/tensorrt_llm/_torch/pyexecutor/model_loader.py index e1ee1d5ee08f..1405c8f6aaca 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_loader.py +++ b/tensorrt_llm/_torch/pyexecutor/model_loader.py @@ -909,11 +909,13 @@ def load( getattr(module, "_requires_standard_hf_loading", False) for module in model.modules()) # Pinned-host parameters must be materialized and filled in place. - # AUTO uses the HF mapper that preserves their stable addresses; - # AUTO may still resolve to MX, so check the resolved format too. - uses_standard_hf_loader = (load_format == LoadFormat.AUTO - and checkpoint_loader.checkpoint_format - != "MX") + # AUTO and LAZY_SAFETENSORS both go through the HF mapper that + # preserves their stable addresses -- the lazy format only changes + # how shards are opened. AUTO may still resolve to MX, so check the + # resolved format too. + uses_standard_hf_loader = ( + load_format in (LoadFormat.AUTO, LoadFormat.LAZY_SAFETENSORS) + and checkpoint_loader.checkpoint_format != "MX") if requires_standard_hf_loading and not uses_standard_hf_loader: raise ValueError( "Host-resident model weights currently require the standard " @@ -971,8 +973,9 @@ def allocate_buffer_on_cuda(t: torch.Tensor): _apply_to_buffers_only(model, allocate_buffer_on_cuda) - need_initialized_weights = load_format not in (LoadFormat.AUTO, - LoadFormat.DUMMY) + need_initialized_weights = load_format not in ( + LoadFormat.AUTO, LoadFormat.DUMMY, + LoadFormat.LAZY_SAFETENSORS) def allocate_weights_on_cuda(t: torch.Tensor): if t not in memo: @@ -1033,7 +1036,13 @@ def init_meta_tensor(t: torch.Tensor): # prior to zero-copy mapping, then refreshes derived state after # real GMS tensors are bound. gms_post_load_handled = False - if load_format == LoadFormat.AUTO: + if load_format in (LoadFormat.AUTO, LoadFormat.LAZY_SAFETENSORS): + # LAZY_SAFETENSORS shares the AUTO disk-load path; it differs + # only in that the checkpoint weight loader opens shards lazily + # (streaming rank-local slices) instead of materializing the + # whole checkpoint. The load_lazily flag carries that choice + # down to HfWeightLoader so the shared loader dispatches on the + # resolved LoadFormat, not on any model name. # Pass model= so format-specific loaders (e.g. MX) can # write weights directly into parameter buffers via P2P. # Generic loaders ignore model=; loaders that can consume a @@ -1043,6 +1052,7 @@ def init_meta_tensor(t: torch.Tensor): "model": model, # Generic loaders ignore it; MXCheckpointLoader pops it. "source_identity": self._source_identity, + "load_lazily": load_format == LoadFormat.LAZY_SAFETENSORS, } if checkpoint_loader.checkpoint_format == "MX": load_weights_kwargs["model_config"] = config @@ -1858,6 +1868,12 @@ def _load_and_validate_config( 'nvfp4_gemm_allowed_backends'] = config.nvfp4_gemm_allowed_backends config.extra_attrs[ 'kv_cache_dtype'] = self.llm_args.kv_cache_config.dtype + # Thread the Kimi K3 FP8 weight-read knobs onto the checkpoint-derived + # quant_config, which is built from hf_quant_config.json and does not + # otherwise carry the user's quant_config. A no-op for non-K3 runs. + # See _torch/models/kimi_k3_knobs.py. + from ..models.kimi_k3_knobs import carry_user_quant_knobs + carry_user_quant_knobs(self.llm_args.quant_config, config.quant_config) # Store allreduce pre-allocation config for AllReduce module access. # Use get_text_config() so VLM wrapper configs (e.g. KimiK2VLConfig, # KimiK25Config) that store the text config under .text_config are diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index c47f5e107984..84a03b5b8b11 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -5278,6 +5278,11 @@ class LoadFormat(Enum): VISION_ONLY = 2 # Load weights through GPU Memory Service. GMS = 3 + # Open safetensors shards lazily and stream only the rank-local slices, + # never materializing the full checkpoint in host RAM. A model opts in by + # declaring this as its default load format (get_model_defaults), or a user + # can select it via load_format="lazy_safetensors" for any HF checkpoint. + LAZY_SAFETENSORS = 4 class ModelExpressConfig(StrictBaseModel): @@ -5799,7 +5804,7 @@ def validate_mla_skip_correction_config(self) -> 'TorchLlmArgs': description= "How to load the model weights. By default, detect the weight type from the model checkpoint.", telemetry=TelemetryField.categorical("auto", "dummy", "vision_only", - "gms")) + "gms", "lazy_safetensors")) enable_min_latency: bool = Field( default=False, diff --git a/tensorrt_llm/models/modeling_utils.py b/tensorrt_llm/models/modeling_utils.py index f5a9f27364f7..34dea2d784cc 100644 --- a/tensorrt_llm/models/modeling_utils.py +++ b/tensorrt_llm/models/modeling_utils.py @@ -167,6 +167,40 @@ class QuantConfig(StrictBaseModel): "Number of Philox rounds for stochastic rounding PRNG. Higher values give better randomness." ) + # --- Kimi K3 FP8 weight-read knobs (opt-in decode-bandwidth optimization) --- + # Opt-in FP8 (e4m3, 128x128 block-scale) reads for selected replicated K3 + # projections; lossy vs BF16, so off by default. ``None`` means unset; + # resolution (is_sm_100f() gate, deprecated-env fallback, per-knob defaults) + # lives in ``_torch/models/kimi_k3_knobs.py``. Per-knob detail on each field. + kimi_k3_fp8_weight_read: Optional[bool] = Field( + default=None, + description= + "Kimi K3: master FP8 block-scale weight read for the replicated MoE-layer " + "MLP/KDA/MLA projections (opt-in, additionally gated on is_sm_100f()). " + "None keeps the historical default (off).") + kimi_k3_fp8_weight_read_kda: Optional[bool] = Field( + default=None, + description= + "Kimi K3: read the KDA q/k/v/g/o projections at FP8 block-scale; narrows " + "an enabled master. None keeps the historical default (on).") + kimi_k3_fp8_weight_read_mla: Optional[bool] = Field( + default=None, + description= + "Kimi K3: read the MLA q_a/q_b/o/gate projections at FP8 block-scale; " + "narrows an enabled master. None keeps the historical default (on).") + kimi_k3_fp8_weight_read_gate_up: Optional[bool] = Field( + default=None, + description= + "Kimi K3: read the fused shared-expert gate_up_proj at FP8 block-scale; " + "narrows an enabled master. None follows the parallel layout (on under " + "attention DP, off otherwise).") + kimi_k3_kda_glue_fp8: Optional[bool] = Field( + default=None, + description= + "Kimi K3: rebuild the fused KDA decode-glue projection on top of the FP8 " + "modules; narrows an enabled FP8 KDA read. None keeps the historical " + "default (on).") + @cached_property def quant_mode(self) -> QuantModeWrapper: quant_mode_list = [ diff --git a/tensorrt_llm/usage/llm_args_golden_manifest.json b/tensorrt_llm/usage/llm_args_golden_manifest.json index aef78b3a38ce..f3738d8e2c78 100644 --- a/tensorrt_llm/usage/llm_args_golden_manifest.json +++ b/tensorrt_llm/usage/llm_args_golden_manifest.json @@ -1030,7 +1030,8 @@ "auto", "dummy", "vision_only", - "gms" + "gms", + "lazy_safetensors" ], "annotation": "Union[str, tensorrt_llm.llmapi.llm_args.LoadFormat]", "converter": "allowlist", diff --git a/tests/unittest/_torch/modeling/test_kimi_k3_fp8_weight_read_gates.py b/tests/unittest/_torch/modeling/test_kimi_k3_fp8_weight_read_gates.py index d28aeb03a0ac..f3e279abc454 100644 --- a/tests/unittest/_torch/modeling/test_kimi_k3_fp8_weight_read_gates.py +++ b/tests/unittest/_torch/modeling/test_kimi_k3_fp8_weight_read_gates.py @@ -12,71 +12,245 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -"""Pin the Kimi K3 FP8 weight-read switch defaults. +"""Pin the Kimi K3 FP8 weight-read gate resolution and its defaults. The FP8 weight read is lossy relative to BF16, so the master switch must stay opt-in — a default run keeps BF16 and matches the published accuracy numbers. These tests exist so that default cannot drift silently: it did once, and the only symptom was accuracy measured against a configuration nobody selected. +The 5 gates now resolve from ``QuantConfig`` (an explicit config value wins), +else the deprecated ``KIMI_K3_*`` env var (still honored, warn-once), else the +historical default. Both surfaces are exercised here, along with the migration's +acceptance criteria at the config/unit level: each of the 5 former env vars has +a config equivalent, a run with neither env vars nor config values resolves a +complete and valid set of gates, and a set env var is honored exactly once with +a deprecation warning. + CPU-only: ``is_sm_100f`` is patched so the gates can be exercised without Blackwell hardware. """ import pytest -from tensorrt_llm._torch.models import modeling_kimi_linear -from tensorrt_llm._torch.models.modeling_kimi_linear import _resolve_fp8_weight_read_gates +from tensorrt_llm._torch.models import kimi_k3_knobs +from tensorrt_llm._torch.models.kimi_k3_knobs import resolve_fp8_weight_read_gates +from tensorrt_llm.models.modeling_utils import QuantConfig ENVS = ( "KIMI_K3_FP8_WEIGHT_READ", "KIMI_K3_FP8_WEIGHT_READ_KDA", + "KIMI_K3_FP8_WEIGHT_READ_MLA", + "KIMI_K3_FP8_WEIGHT_READ_GATE_UP", "KIMI_K3_KDA_GLUE_FP8", ) @pytest.fixture def sm100f(monkeypatch): - """Report Blackwell so the SM gate never masks the env behavior.""" - monkeypatch.setattr(modeling_kimi_linear, "is_sm_100f", lambda: True) + """Report Blackwell so the SM gate never masks env/config behavior; clear env.""" + monkeypatch.setattr(kimi_k3_knobs, "is_sm_100f", lambda: True) for env in ENVS: monkeypatch.delenv(env, raising=False) +def _gates(quant_config=None, *, enable_attention_dp=False): + return resolve_fp8_weight_read_gates(quant_config, enable_attention_dp=enable_attention_dp) + + +def _tuple(g): + return (g.master, g.kda, g.kda_glue, g.mla, g.gate_up) + + +# --- Defaults (config surface) --------------------------------------------- + + def test_master_switch_defaults_off(sm100f): """With nothing set, no FP8 weight read anywhere.""" - assert _resolve_fp8_weight_read_gates() == (False, False, False) + assert _tuple(_gates(QuantConfig())) == (False, False, False, False, False) -def test_master_switch_opt_in(sm100f, monkeypatch): +def test_master_switch_opt_in_via_config(sm100f): """Setting the master switch enables it and the default-on sub-gates.""" + g = _gates(QuantConfig(kimi_k3_fp8_weight_read=True), enable_attention_dp=True) + assert _tuple(g) == (True, True, True, True, True) + + +def test_gate_up_default_follows_attention_dp(sm100f): + """gate_up defaults on under attention-DP, off otherwise (master on).""" + on = _gates(QuantConfig(kimi_k3_fp8_weight_read=True), enable_attention_dp=True) + off = _gates(QuantConfig(kimi_k3_fp8_weight_read=True), enable_attention_dp=False) + assert on.gate_up is True and off.gate_up is False + + +@pytest.mark.parametrize( + "field,expected", + [ + # kda off -> kda_glue collapses too; mla/gate_up independent. + ("kimi_k3_fp8_weight_read_kda", (True, False, False, True, True)), + ("kimi_k3_kda_glue_fp8", (True, True, False, True, True)), + ("kimi_k3_fp8_weight_read_mla", (True, True, True, False, True)), + ("kimi_k3_fp8_weight_read_gate_up", (True, True, True, True, False)), + ], +) +def test_sub_gates_narrow_enabled_master_via_config(sm100f, field, expected): + """Sub-gates only ever narrow an enabled master; they never enable alone.""" + g = _gates( + QuantConfig(kimi_k3_fp8_weight_read=True, **{field: False}), enable_attention_dp=True + ) + assert _tuple(g) == expected + + +@pytest.mark.parametrize( + "field", + [ + "kimi_k3_fp8_weight_read_kda", + "kimi_k3_fp8_weight_read_mla", + "kimi_k3_fp8_weight_read_gate_up", + "kimi_k3_kda_glue_fp8", + ], +) +def test_sub_gates_inert_while_master_off_via_config(sm100f, field): + """A sub-gate set on must not turn on FP8 reads by itself.""" + g = _gates(QuantConfig(**{field: True}), enable_attention_dp=True) + assert _tuple(g) == (False, False, False, False, False) + + +def test_non_blackwell_never_reads_fp8_config(monkeypatch): + """The SM gate wins even with the master switch explicitly on (config).""" + monkeypatch.setattr(kimi_k3_knobs, "is_sm_100f", lambda: False) + for env in ENVS: + monkeypatch.delenv(env, raising=False) + g = _gates(QuantConfig(kimi_k3_fp8_weight_read=True), enable_attention_dp=True) + assert _tuple(g) == (False, False, False, False, False) + + +# --- Deprecated env var back-compat (still honored) ------------------------ + + +def test_master_switch_opt_in_via_deprecated_env(sm100f, monkeypatch): + """The deprecated master env var is still honored when config is unset.""" monkeypatch.setenv("KIMI_K3_FP8_WEIGHT_READ", "1") - assert _resolve_fp8_weight_read_gates() == (True, True, True) + g = _gates(QuantConfig(), enable_attention_dp=False) + # gate_up follows attention_dp=False -> off; the rest default on. + assert _tuple(g) == (True, True, True, True, False) @pytest.mark.parametrize( "env,expected", [ - ("KIMI_K3_FP8_WEIGHT_READ_KDA", (True, False, False)), - ("KIMI_K3_KDA_GLUE_FP8", (True, True, False)), + ("KIMI_K3_FP8_WEIGHT_READ_KDA", (True, False, False, True, True)), + ("KIMI_K3_KDA_GLUE_FP8", (True, True, False, True, True)), + ("KIMI_K3_FP8_WEIGHT_READ_MLA", (True, True, True, False, True)), + ("KIMI_K3_FP8_WEIGHT_READ_GATE_UP", (True, True, True, True, False)), ], ) -def test_sub_gates_narrow_an_enabled_master(sm100f, monkeypatch, env, expected): - """Sub-gates only ever narrow; they never enable on their own.""" +def test_sub_gates_narrow_via_deprecated_env(sm100f, monkeypatch, env, expected): monkeypatch.setenv("KIMI_K3_FP8_WEIGHT_READ", "1") monkeypatch.setenv(env, "0") - assert _resolve_fp8_weight_read_gates() == expected + assert _tuple(_gates(QuantConfig(), enable_attention_dp=True)) == expected -@pytest.mark.parametrize("env", ENVS[1:]) -def test_sub_gates_are_inert_while_master_is_off(sm100f, monkeypatch, env): - """A sub-gate set to 1 must not turn on FP8 reads by itself.""" - monkeypatch.setenv(env, "1") - assert _resolve_fp8_weight_read_gates() == (False, False, False) +def test_non_blackwell_never_reads_fp8_env(monkeypatch): + """The SM gate wins even with the master env var on.""" + monkeypatch.setattr(kimi_k3_knobs, "is_sm_100f", lambda: False) + for env in ENVS: + monkeypatch.delenv(env, raising=False) + monkeypatch.setenv("KIMI_K3_FP8_WEIGHT_READ", "1") + assert _tuple(_gates(QuantConfig(), enable_attention_dp=True)) == ( + False, + False, + False, + False, + False, + ) + + +# --- Precedence: config wins over the deprecated env var ------------------- + + +def test_config_master_off_beats_env_master_on(sm100f, monkeypatch): + """Explicit config value takes precedence over the deprecated env var.""" + monkeypatch.setenv("KIMI_K3_FP8_WEIGHT_READ", "1") + g = _gates(QuantConfig(kimi_k3_fp8_weight_read=False), enable_attention_dp=True) + assert _tuple(g) == (False, False, False, False, False) + + +def test_config_subgate_off_beats_env_subgate_on(sm100f, monkeypatch): + """A config sub-gate overrides a conflicting deprecated env sub-gate.""" + monkeypatch.setenv("KIMI_K3_FP8_WEIGHT_READ", "1") + monkeypatch.setenv("KIMI_K3_FP8_WEIGHT_READ_KDA", "1") + g = _gates( + QuantConfig(kimi_k3_fp8_weight_read=True, kimi_k3_fp8_weight_read_kda=False), + enable_attention_dp=True, + ) + assert g.master is True and g.kda is False and g.kda_glue is False + + +# --- Config surface: every former env var has a config equivalent ----------- + + +def test_fp8_knobs_exist_on_quant_config(): + """Each of the 5 former KIMI_K3_* env vars has a QuantConfig field.""" + fields = QuantConfig.model_fields + for name in kimi_k3_knobs.KIMI_K3_QUANT_KNOB_FIELDS: + assert name in fields, f"QuantConfig missing {name}" + assert fields[name].default is None, f"{name} must default to None (unset)" + assert len(kimi_k3_knobs.KIMI_K3_QUANT_KNOB_FIELDS) == len(ENVS) + + +def test_knobs_settable_from_extra_llm_api_options_dict(): + """The knobs round-trip through a plain dict, as extra_llm_api_options does.""" + qc = QuantConfig.model_validate( + {"kimi_k3_fp8_weight_read": True, "kimi_k3_fp8_weight_read_kda": False} + ) + assert qc.kimi_k3_fp8_weight_read is True + assert qc.kimi_k3_fp8_weight_read_kda is False + + +def test_zero_env_zero_config_resolves_complete_valid(sm100f): + """No env vars + a default QuantConfig -> every gate resolves to a concrete + value (the historical BF16 default), with no missing-knob error.""" + g = _gates(QuantConfig(), enable_attention_dp=False) + assert _tuple(g) == (False, False, False, False, False) + assert all(isinstance(v, bool) for v in _tuple(g)) + + +# --- Deprecation warning: emitted once, and even when config overrides ------ + + +class _RecordingLogger: + """Mirrors ``logger.warning_once`` dedup so warn-once is deterministic.""" + + def __init__(self): + self._seen = set() + self.emitted = [] + + def warning_once(self, *msg, key): + if key not in self._seen: + self._seen.add(key) + self.emitted.append((key, " ".join(str(m) for m in msg))) + + +def test_deprecated_env_warns_once_and_is_honored(sm100f, monkeypatch): + rec = _RecordingLogger() + monkeypatch.setattr(kimi_k3_knobs, "logger", rec) + monkeypatch.setenv("KIMI_K3_FP8_WEIGHT_READ", "1") + # Resolve several times; the deprecation warning must appear exactly once. + for _ in range(3): + assert _gates(QuantConfig(), enable_attention_dp=True).master is True + keys = [k for k, _ in rec.emitted] + assert keys == ["kimi_k3_deprecated_env::KIMI_K3_FP8_WEIGHT_READ"] + assert "extra_llm_api_options" in rec.emitted[0][1] -def test_non_blackwell_never_reads_fp8(monkeypatch): - """The SM gate wins even with the master switch explicitly on.""" - monkeypatch.setattr(modeling_kimi_linear, "is_sm_100f", lambda: False) +def test_warn_once_even_when_config_overrides(sm100f, monkeypatch): + """The deprecation is about the env var existing; it warns even when a + config value overrides the env value.""" + rec = _RecordingLogger() + monkeypatch.setattr(kimi_k3_knobs, "logger", rec) monkeypatch.setenv("KIMI_K3_FP8_WEIGHT_READ", "1") - assert _resolve_fp8_weight_read_gates() == (False, False, False) + assert _gates(QuantConfig(kimi_k3_fp8_weight_read=False), enable_attention_dp=True).master is ( + False + ) + assert [k for k, _ in rec.emitted] == ["kimi_k3_deprecated_env::KIMI_K3_FP8_WEIGHT_READ"] 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 ccf1973bb7f2..7043b4006797 100644 --- a/tests/unittest/_torch/models/checkpoints/hf/test_weight_loader.py +++ b/tests/unittest/_torch/models/checkpoints/hf/test_weight_loader.py @@ -13,18 +13,15 @@ # 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 @@ -496,51 +493,140 @@ def test_prefetch_files_emits_progress_heartbeat(tmp_path, monkeypatch): assert len(progress_logs) >= 12 -def test_kimi_k3_lazy_load_records_the_checkpoint_dir(tmp_path): +def test_lazy_load_records_the_checkpoint_dir(tmp_path): """A model that re-opens shards itself needs the directory back. - Kimi K3 streams its rank-local experts per shard file to avoid holding - the whole mapping open. A lazy slice does not carry its file, and - transformers 5.x no longer sets ``PretrainedConfig._name_or_path``, so - without this the model silently fell back to the shared mapping and the - step was OOM-killed. + Expert-parallel models stream their rank-local experts per shard file to + avoid holding the whole mapping open. A lazy slice does not carry its + file, and transformers 5.x no longer sets ``PretrainedConfig._name_or_path``, + so without this the model silently fell back to the shared mapping and the + step was OOM-killed. The lazy ``LoadFormat`` is selected here via the + ``load_lazily`` flag (what ``LoadFormat.LAZY_SAFETENSORS`` resolves to in + the model loader) -- no checkpoint model-name detection is involved. """ - import safetensors.torch import torch - (tmp_path / "config.json").write_text(json.dumps({"model_type": "kimi_k3"})) safetensors.torch.save_file( {"w": torch.zeros(2, 2)}, tmp_path / "model-00001-of-00001.safetensors" ) loader = HfWeightLoader() + # ``weights`` is bound before the try so a load_weights() failure surfaces + # as itself instead of an UnboundLocalError raised from the cleanup path. + weights = None try: - weights = loader.load_weights(str(tmp_path), Mapping()) + weights = loader.load_weights(str(tmp_path), Mapping(), load_lazily=True) assert isinstance(weights, ConsumableWeightsDict) assert weights.checkpoint_dir == str(tmp_path) finally: + if weights is not None: + weights.clear() 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 +def _materialize_lazy_slice(value): + """Realize a lazy safetensors slice the way a model's load_weights does. + + Mirrors ``modeling_kimi_linear._materialize``: ``[:]`` realizes a normal + slice, but a 0-dim (scalar) entry needs ``[()]``. + """ + import torch + + if isinstance(value, torch.Tensor): + return value + get_shape = getattr(value, "get_shape", None) + if get_shape is not None and len(get_shape()) == 0: + return value[()] + return value[:] + +def _write_multi_shard_fixture(checkpoint_dir): + """Write a small multi-shard safetensors fixture (no big checkpoint). -@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 + Two sharded files carry distinct keys; a consolidated file carries a + different key so which file set was selected is directly observable from + the loaded key set. Includes a quantized (fp8) weight and a companion + fp32 block-scale tensor so dtype and scale tensors are exercised. + Returns (sharded_keys, consolidated_keys). + """ + import safetensors.torch + import torch + + sharded_a = { + "layer.0.weight": torch.arange(16, dtype=torch.bfloat16).reshape(4, 4), + # FP8 block-scale companion (kept BF16-free on purpose: a real scale). + "layer.0.weight_scale_inv": torch.rand(2, 2, dtype=torch.float32), + } + sharded_b = { + # A genuinely quantized tensor: proves the fp8 dtype survives the + # lazy path byte-for-byte. + "layer.1.weight": torch.arange(16).reshape(4, 4).to(torch.float8_e4m3fn), + } + consolidated = { + "consolidated.only.weight": torch.ones(2, 2, dtype=torch.bfloat16), + } + safetensors.torch.save_file(sharded_a, str(checkpoint_dir / "model-00001-of-00002.safetensors")) + safetensors.torch.save_file(sharded_b, str(checkpoint_dir / "model-00002-of-00002.safetensors")) + safetensors.torch.save_file(consolidated, str(checkpoint_dir / "consolidated.safetensors")) + return (set(sharded_a) | set(sharded_b)), set(consolidated) + + +@pytest.mark.parametrize("use_consolidated", [False, True]) +def test_lazy_format_load_matches_eager_load(tmp_path, use_consolidated): + """The lazy ``LoadFormat`` equals an eager load on the same fixture. + + Covers criterion s1.g1: a general lazy format whose sharded-vs-consolidated + selection matches the eager path, whose materialized values/dtype/scale + tensors match the eager load exactly, and whose rank-local sliced reads do + not materialize the full tensor set. Runs on CPU, no 1.5 TB checkpoint. + """ + import torch -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 + checkpoint_dir = tmp_path / "ckpt" + checkpoint_dir.mkdir() + sharded_keys, consolidated_keys = _write_multi_shard_fixture(checkpoint_dir) + expected_keys = consolidated_keys if use_consolidated else sharded_keys + + loader = HfWeightLoader() + + # Eager reference load. Mock only the OS-cache prefetch; the real file + # selection and the real safetensors load both run. + with mock.patch.object(loader, "prefetch_files"): + eager = loader.load_weights( + str(checkpoint_dir), Mapping(), use_consolidated=use_consolidated + ) + eager_values = {key: eager[key] for key in eager.keys()} + + # Lazy-format load of the same fixture. + lazy = loader.load_weights( + str(checkpoint_dir), Mapping(), use_consolidated=use_consolidated, load_lazily=True + ) + try: + # (1) use_consolidated selection matches the eager path: same key set. + assert set(eager.keys()) == expected_keys + assert set(lazy.keys()) == expected_keys + + # (2) Lazy values are slice handles, not materialized tensors: merely + # opening the checkpoint never realizes the full tensor set. + for key in lazy.keys(): + assert not isinstance(lazy[key], torch.Tensor), key + + # (3) A rank-local sliced read materializes ONLY the requested bytes. + row_key = "consolidated.only.weight" if use_consolidated else "layer.0.weight" + partial = lazy[row_key][0:1] + assert isinstance(partial, torch.Tensor) + assert partial.shape[0] == 1 # only the requested row, not all rows + + # (4) Materialized lazy values equal the eager values in dtype, shape, + # and bytes -- scale and fp8 tensors included. + for key in expected_keys: + lazy_val = _materialize_lazy_slice(lazy[key]) + eager_val = eager_values[key] + assert lazy_val.dtype == eager_val.dtype, key + assert lazy_val.shape == eager_val.shape, key + assert torch.equal(lazy_val.to(torch.float32), eager_val.to(torch.float32)), key + finally: + lazy.clear() + loader.cleanup() diff --git a/tests/unittest/api_stability/references/quant_config.yaml b/tests/unittest/api_stability/references/quant_config.yaml index 9aa0f6b5366f..36c1bf150734 100644 --- a/tests/unittest/api_stability/references/quant_config.yaml +++ b/tests/unittest/api_stability/references/quant_config.yaml @@ -13,6 +13,21 @@ methods: has_zero_point: annotation: bool default: false + kimi_k3_fp8_weight_read: + annotation: Optional[bool] + default: null + kimi_k3_fp8_weight_read_gate_up: + annotation: Optional[bool] + default: null + kimi_k3_fp8_weight_read_kda: + annotation: Optional[bool] + default: null + kimi_k3_fp8_weight_read_mla: + annotation: Optional[bool] + default: null + kimi_k3_kda_glue_fp8: + annotation: Optional[bool] + default: null kv_cache_quant_algo: annotation: Optional[tensorrt_llm.quantization.mode.QuantAlgo] default: null