From d41e8c459defe11699678b5ca4299f8ca89d2911 Mon Sep 17 00:00:00 2001 From: Sam Shleifer Date: Fri, 1 May 2026 23:11:40 -0400 Subject: [PATCH] Support RunAI loading for quantized checkpoints (#23850) Co-authored-by: Claude Opus 4.7 (1M context) Co-authored-by: Sam Shleifer --- python/sglang/srt/model_loader/loader.py | 15 +- .../sglang/srt/model_loader/weight_utils.py | 6 +- .../deepseek_common/deepseek_weight_loader.py | 15 +- python/sglang/srt/models/kimi_k25.py | 67 +++++---- .../test_runai_model_streamer_loader.py | 128 ++++++++++++++++++ 5 files changed, 196 insertions(+), 35 deletions(-) create mode 100644 test/registered/unit/model_loader/test_runai_model_streamer_loader.py diff --git a/python/sglang/srt/model_loader/loader.py b/python/sglang/srt/model_loader/loader.py index 94ad0e4fe..26aade419 100644 --- a/python/sglang/srt/model_loader/loader.py +++ b/python/sglang/srt/model_loader/loader.py @@ -3238,11 +3238,13 @@ class RunaiModelStreamerLoader(BaseModelLoader): self.target_device_str = "cpu" target_device = torch.device(device_config.device) + quant_config = _get_quantization_config(model_config, self.load_config) with set_default_torch_dtype(model_config.dtype): with target_device: model = _initialize_model( model_config, self.load_config, + quant_config, ) DefaultModelLoader.load_weights_and_postprocess( @@ -3260,7 +3262,16 @@ def get_model_loader( if load_config.load_format == LoadFormat.DUMMY: return DummyModelLoader(load_config) - if model_config and ( + # ModelOptModelLoader's local-copy quantize-and-export workflow doesn't apply + # to RUNAI_STREAMER, which streams weights directly from object storage. + # RUNAI_STREAMER loads always fall through to the unconditional branch at + # the bottom of this function. This also avoids calling _is_already_quantized() + # on RunAI streamer cache paths, where huggingface_hub raises HFValidationError. + model_optloader_allowed = ( + model_config and load_config.load_format != LoadFormat.RUNAI_STREAMER + ) + + if model_optloader_allowed and ( (hasattr(model_config, "modelopt_quant") and model_config.modelopt_quant) or model_config.quantization in ["modelopt_fp8", "modelopt_fp4", "modelopt_mixed", "modelopt"] @@ -3270,7 +3281,7 @@ def get_model_loader( # Use ModelOptModelLoader for unified quantization flags if ( - model_config + model_optloader_allowed and hasattr(model_config, "quantization") and model_config.quantization in ["modelopt_fp8", "modelopt_fp4", "modelopt_mixed"] diff --git a/python/sglang/srt/model_loader/weight_utils.py b/python/sglang/srt/model_loader/weight_utils.py index 28d655ec3..084fa02ec 100644 --- a/python/sglang/srt/model_loader/weight_utils.py +++ b/python/sglang/srt/model_loader/weight_utils.py @@ -69,6 +69,8 @@ except ImportError as e: logger = logging.getLogger(__name__) +RUNAI_STREAMER_TENSOR_ATTR = "_sglang_runai_streamer_tensor" + # Block size for sequential checkpoint prefetch reads (page cache warming). _PREFETCH_BLOCK_SIZE = None @@ -1317,7 +1319,9 @@ def runai_safetensors_weights_iterator( mininterval=2, ) - yield from tensor_iter + for name, tensor in tensor_iter: + setattr(tensor, RUNAI_STREAMER_TENSOR_ATTR, True) + yield name, tensor def set_runai_streamer_env(load_config: LoadConfig): diff --git a/python/sglang/srt/models/deepseek_common/deepseek_weight_loader.py b/python/sglang/srt/models/deepseek_common/deepseek_weight_loader.py index 8112b321e..754dc1bba 100644 --- a/python/sglang/srt/models/deepseek_common/deepseek_weight_loader.py +++ b/python/sglang/srt/models/deepseek_common/deepseek_weight_loader.py @@ -44,7 +44,10 @@ from sglang.srt.model_loader.utils import ( should_async_load, should_deepgemm_weight_requant_ue8m0, ) -from sglang.srt.model_loader.weight_utils import default_weight_loader +from sglang.srt.model_loader.weight_utils import ( + RUNAI_STREAMER_TENSOR_ATTR, + default_weight_loader, +) from sglang.srt.models.deepseek_common.utils import ( _is_cuda, _is_fp8_fnuz, @@ -67,6 +70,12 @@ logger = logging.getLogger(__name__) NVFP4_CKPT_FP8_ATTN_QUANT_MODULES = ["q_b_proj"] +def _clone_if_runai_streamed_tensor(tensor: torch.Tensor) -> torch.Tensor: + if getattr(tensor, RUNAI_STREAMER_TENSOR_ATTR, False): + return tensor.clone().detach() + return tensor + + @dataclass(frozen=True) class NextNEnabledConfig: num_nextn_layers: int @@ -267,7 +276,9 @@ class DeepseekV2WeightLoaderMixin: if fuse_qkv_a_proj and ( "q_a_proj" in name or "kv_a_proj_with_mqa" in name ): - cached_a_proj[name] = loaded_weight + cached_a_proj[name] = _clone_if_runai_streamed_tensor( + loaded_weight + ) q_a_proj_name = ( name if "q_a_proj" in name diff --git a/python/sglang/srt/models/kimi_k25.py b/python/sglang/srt/models/kimi_k25.py index 571f9fbb4..832ee74dd 100644 --- a/python/sglang/srt/models/kimi_k25.py +++ b/python/sglang/srt/models/kimi_k25.py @@ -743,42 +743,49 @@ class KimiK25ForConditionalGeneration(nn.Module): return hidden_states def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): - """Load weights for the model, separating vision and language weights""" + """Stream weights, loading vision weights inline and yielding language weights. + + The streaming pattern (vs accumulating into lists) is required because RunAI's + iterator reuses backing buffers — collecting tensors before consuming them + would clobber prior tensors. + """ mapper = getattr(self, "hf_to_sglang_mapper", None) if mapper is not None: weights = mapper.apply(weights) - # Separate vision tower weights and language model weights - vision_weights = [] - language_weights = [] + vision_params = ( + None + if self.config.language_only + else dict(self.named_parameters(remove_duplicate=False)) + ) - for name, loaded_weight in weights: - if "vision_tower" in name or "mm_projector" in name: - name = name.replace(r"wqkv.", r"attn.qkv_proj.") - name = name.replace(r"wo.", r"attn.proj.") - name = name.replace("mm_projector.proj.0", "mm_projector.linear_1") - name = name.replace("mm_projector.proj.2", "mm_projector.linear_2") - vision_weights.append((name, loaded_weight)) - else: - name = name.replace("language_model.", "") - # All other weights go to language model - language_weights.append((name, loaded_weight)) + def stream_language_weights(): + for name, loaded_weight in weights: + if "vision_tower" in name or "mm_projector" in name: + if vision_params is None: + continue + vname = ( + name.replace(r"wqkv.", r"attn.qkv_proj.") + .replace(r"wo.", r"attn.proj.") + .replace("mm_projector.proj.0", "mm_projector.linear_1") + .replace("mm_projector.proj.2", "mm_projector.linear_2") + ) + if vname not in vision_params: + raise ValueError(f"Weight {vname} not found in params_dict") + param = vision_params[vname] + weight_loader = getattr( + param, "weight_loader", default_weight_loader + ) + weight_loader(param, loaded_weight) + continue + yield name.replace("language_model.", ""), loaded_weight - if not self.config.language_only: - # Load vision tower weights - vision_state_dict = dict(vision_weights) - params_dict = dict(self.named_parameters(remove_duplicate=False)) - for name, loaded_weight in vision_state_dict.items(): - if name not in params_dict: - raise ValueError(f"Weight {name} not found in params_dict") - param = params_dict[name] - weight_loader = getattr(param, "weight_loader", default_weight_loader) - # loaded_weight = self._pad_vit_attn_dummy_heads(name, loaded_weight) - weight_loader(param, loaded_weight) - - # Load language model weights - if not self.config.encoder_only and language_weights: - self.language_model.load_weights(language_weights) + if self.language_model is not None: + self.language_model.load_weights(stream_language_weights()) + else: + # encoder-only: drain the generator so inline vision-weight loading fires. + for _ in stream_language_weights(): + pass @classmethod def get_model_config_for_expert_location(cls, config: KimiK25Config): diff --git a/test/registered/unit/model_loader/test_runai_model_streamer_loader.py b/test/registered/unit/model_loader/test_runai_model_streamer_loader.py new file mode 100644 index 000000000..8b409c560 --- /dev/null +++ b/test/registered/unit/model_loader/test_runai_model_streamer_loader.py @@ -0,0 +1,128 @@ +import sys +import unittest +from types import SimpleNamespace +from typing import cast +from unittest.mock import patch + +import torch + +import sglang.srt.model_loader.loader as loader_mod +import sglang.srt.model_loader.weight_utils as weight_utils +from sglang.srt.configs.device_config import DeviceConfig +from sglang.srt.configs.load_config import LoadConfig, LoadFormat +from sglang.srt.configs.model_config import ModelConfig +from sglang.srt.models.deepseek_common import deepseek_weight_loader +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=6, suite="stage-a-test-cpu") + + +class _FakeModel: + def eval(self): + return self + + +class TestRunaiModelStreamerLoader(CustomTestCase): + def test_passes_quant_config_to_model_init(self): + quant_config = object() + fake_model = _FakeModel() + + with ( + patch.object( + loader_mod, + "_get_quantization_config", + return_value=quant_config, + ), + patch.object(loader_mod, "_initialize_model") as mock_initialize_model, + patch.object( + loader_mod.DefaultModelLoader, + "load_weights_and_postprocess", + ) as mock_load_weights, + ): + mock_initialize_model.return_value = fake_model + runai_loader = loader_mod.RunaiModelStreamerLoader( + LoadConfig( + load_format=LoadFormat.RUNAI_STREAMER, + model_loader_extra_config={}, + ) + ) + model_config = cast( + ModelConfig, + SimpleNamespace(dtype=torch.float16, modelopt_quant=False), + ) + + model = runai_loader.load_model( + model_config=model_config, + device_config=DeviceConfig("cpu"), + ) + + self.assertIs(model, fake_model) + self.assertIs(mock_load_weights.call_args.args[0], fake_model) + self.assertIs(mock_initialize_model.call_args.args[2], quant_config) + + def test_marks_streamer_tensors(self): + source_tensor = torch.tensor([1], dtype=torch.int32) + + class FakeStreamer: + def __enter__(self): + return self + + def __exit__(self, *_args): + pass + + def stream_files(self, *_args, **_kwargs): + self.files_to_tensors_metadata = {0: [object()]} + + def get_tensors(self): + yield "weight", source_tensor + + with patch.dict( + sys.modules, + {"runai_model_streamer": SimpleNamespace(SafetensorsStreamer=FakeStreamer)}, + ): + weights = list( + weight_utils.runai_safetensors_weights_iterator(["model.safetensors"]) + ) + + self.assertEqual(weights[0][0], "weight") + self.assertTrue(getattr(weights[0][1], weight_utils.RUNAI_STREAMER_TENSOR_ATTR)) + + def test_deepseek_clone_only_clones_marked_tensors(self): + unmarked = torch.tensor([1], dtype=torch.int32) + + self.assertIs( + deepseek_weight_loader._clone_if_runai_streamed_tensor(unmarked), + unmarked, + ) + + marked = torch.tensor([1], dtype=torch.int32) + setattr(marked, weight_utils.RUNAI_STREAMER_TENSOR_ATTR, True) + + cloned = deepseek_weight_loader._clone_if_runai_streamed_tensor(marked) + + self.assertIsNot(cloned, marked) + marked.fill_(2) + self.assertEqual(cloned.item(), 1) + + def test_get_model_loader_uses_runai_for_prequantized_modelopt(self): + load_config = LoadConfig( + load_format=LoadFormat.RUNAI_STREAMER, + model_loader_extra_config={}, + ) + model_config = cast( + ModelConfig, + SimpleNamespace( + quantization="modelopt_fp4", + modelopt_quant=False, + _is_already_quantized=lambda: True, + ), + ) + + model_loader = loader_mod.get_model_loader(load_config, model_config) + + self.assertIsInstance(model_loader, loader_mod.RunaiModelStreamerLoader) + + +if __name__ == "__main__": + unittest.main()