From 5b540b16de8eb33dd24ad46c056b2c35640efe18 Mon Sep 17 00:00:00 2001 From: fzyzcjy <5236035+fzyzcjy@users.noreply.github.com> Date: Tue, 14 Jul 2026 15:58:40 +0800 Subject: [PATCH] Extract load_model helpers into a load_model_utils module (#31155) --- .../sglang/srt/model_executor/model_runner.py | 239 ++++------------ .../load_model_utils.py | 266 ++++++++++++++++++ 2 files changed, 321 insertions(+), 184 deletions(-) create mode 100644 python/sglang/srt/model_executor/model_runner_components/load_model_utils.py diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index 16627622b..c97cb9dd0 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -16,22 +16,16 @@ from __future__ import annotations import contextlib -import datetime import inspect import logging -import os -import socket -import threading import time from collections import defaultdict from dataclasses import dataclass from typing import Optional, Union import torch -import torch.distributed as dist -from sglang.srt.configs.device_config import DeviceConfig -from sglang.srt.configs.load_config import LoadConfig, LoadFormat +from sglang.srt.configs.load_config import LoadConfig from sglang.srt.configs.model_config import ( AttentionArch, ModelConfig, @@ -41,20 +35,14 @@ from sglang.srt.configs.model_config import ( is_deepseek_dsa, ) from sglang.srt.configs.update_config import adjust_config_with_unaligned_cpu_tp -from sglang.srt.constants import GPU_MEMORY_TYPE_WEIGHTS from sglang.srt.debug_utils.dumper import dumper -from sglang.srt.debug_utils.tensor_dump_forward_hook import ( - register_forward_hook_for_model, -) from sglang.srt.distributed import ( bootstrap, - get_tp_group, get_world_group, ) from sglang.srt.distributed.device_communicators.pynccl_allocator import ( prealloc_symmetric_memory_pool, ) -from sglang.srt.distributed.parallel_state import monkey_patch_vllm_parallel_state from sglang.srt.dllm.config import DllmConfig from sglang.srt.elastic_ep.elastic_ep import ( ElasticEPStateManager, @@ -129,6 +117,17 @@ from sglang.srt.model_executor.forward_context import ( ) from sglang.srt.model_executor.graph_shared_output import GraphSharedOutput from sglang.srt.model_executor.hook_manager import register_forward_hooks +from sglang.srt.model_executor.model_runner_components.load_model_utils import ( + build_load_config, + dist_barrier_after_load, + load_kv_cache_scales, + load_model_with_memory_saver, + maybe_downgrade_dtype_for_legacy_gpu, + maybe_register_debug_tensor_dump_hook, + maybe_trigger_remote_instance_nccl_send_group, + report_online_quantization, + resolve_sliding_window_size, +) from sglang.srt.model_executor.model_runner_components.ngram_embedding_manager import ( NgramEmbeddingManager, ) @@ -150,11 +149,6 @@ from sglang.srt.model_executor.runner import ( PrefillCudaGraphRunner, get_batch_sizes_to_capture, ) -from sglang.srt.model_loader.loader import get_model_loader -from sglang.srt.model_loader.remote_instance_weight_loader_utils import ( - RemoteInstanceWeightLoaderBackend, - trigger_init_weights_send_group_for_remote_instance_request, -) from sglang.srt.model_loader.utils import resolve_language_model from sglang.srt.platforms import current_platform from sglang.srt.runtime_context import get_flags, get_server_args @@ -196,7 +190,7 @@ from sglang.srt.utils import ( set_cuda_arch, slow_rank_detector, ) -from sglang.srt.utils.network import NetworkAddress, get_local_ip_auto +from sglang.srt.utils.network import get_local_ip_auto from sglang.srt.utils.nvtx_pytorch_hooks import PytHooks from sglang.srt.utils.nvtx_utils import profile_range from sglang.srt.utils.offloader import ( @@ -222,7 +216,6 @@ elif current_platform.is_out_of_tree(): current_platform.init_backend() # Detect stragger ranks in model loading -UNBALANCED_MODEL_LOADING_TIMEOUT_S = 480 # leave more time for post data processing logger = logging.getLogger(__name__) @@ -1115,49 +1108,17 @@ class ModelRunner(ModelRunnerKVCacheMixin): if self.device != "cpu": torch.set_num_threads(1) if self.device == "cuda": - if torch.cuda.get_device_capability()[0] < 8: - logger.info( - "Compute capability below sm80. Use float16 due to lack of bfloat16 support." - ) - from sglang.srt.arg_groups.overrides import ( - declare_load_time_override, - ) - - declare_load_time_override( - "ModelRunner._sm80_dtype_fallback", {"dtype": "float16"} - ) - self.model_config.dtype = torch.float16 - if torch.cuda.get_device_capability()[1] < 5: - raise RuntimeError("SGLang only supports sm75 and above.") + maybe_downgrade_dtype_for_legacy_gpu( + server_args=self.server_args, model_config=self.model_config + ) set_cuda_arch() - # Prepare the model config - from sglang.srt.configs.modelopt_config import ModelOptConfig - - modelopt_config = ModelOptConfig( - quant=self.server_args.modelopt_quant, - checkpoint_restore_path=self.server_args.modelopt_checkpoint_restore_path, - checkpoint_save_path=self.server_args.modelopt_checkpoint_save_path, - export_path=self.server_args.modelopt_export_path, - quantize_and_serve=self.server_args.quantize_and_serve, - ) - - self.load_config = LoadConfig( - load_format=self.server_args.load_format, - download_dir=self.server_args.download_dir, - model_loader_extra_config=self.server_args.model_loader_extra_config, + self.load_config = build_load_config( + server_args=self.server_args, tp_rank=self.tp_rank, - remote_instance_weight_loader_seed_instance_ip=self.server_args.remote_instance_weight_loader_seed_instance_ip, - remote_instance_weight_loader_seed_instance_service_port=self.server_args.remote_instance_weight_loader_seed_instance_service_port, - remote_instance_weight_loader_send_weights_group_ports=self.server_args.remote_instance_weight_loader_send_weights_group_ports, - remote_instance_weight_loader_backend=self.server_args.remote_instance_weight_loader_backend, - remote_instance_weight_loader_transfer_engine=self.remote_instance_weight_transporter.engine, - remote_instance_weight_loader_transfer_engine_session_id=self.remote_instance_weight_transporter.session_id, - modelexpress_url=self.server_args.modelexpress_url, - modelexpress_transport=self.server_args.modelexpress_transport, - modelopt_config=modelopt_config, - rl_quant_profile=self.server_args.rl_quant_profile, + remote_instance_weight_transporter_engine=self.remote_instance_weight_transporter.engine, + remote_instance_weight_transporter_session_id=self.remote_instance_weight_transporter.session_id, draft_model_idx=self.draft_model_idx, ) if self.device == "cpu": @@ -1165,52 +1126,25 @@ class ModelRunner(ModelRunnerKVCacheMixin): self.model_config, self.load_config, self.tp_size ) - if ( - self.server_args.load_format == LoadFormat.REMOTE_INSTANCE - and self.server_args.remote_instance_weight_loader_backend - == RemoteInstanceWeightLoaderBackend.NCCL - ): - if self.tp_rank == 0: - instance_ip = NetworkAddress.resolve_host(socket.gethostname()) - t = threading.Thread( - target=trigger_init_weights_send_group_for_remote_instance_request, - args=( - self.server_args.remote_instance_weight_loader_seed_instance_ip, - self.server_args.remote_instance_weight_loader_seed_instance_service_port, - self.server_args.remote_instance_weight_loader_send_weights_group_ports, - instance_ip, - ), - ) - t.start() - - # Load the model - # Remove monkey_patch when linear.py quant remove dependencies with vllm - monkey_patch_vllm_parallel_state() - - enable_cpu_backup = self.server_args.enable_weights_cpu_backup or ( - self.is_draft_worker and self.server_args.enable_draft_weights_cpu_backup + maybe_trigger_remote_instance_nccl_send_group( + server_args=self.server_args, tp_rank=self.tp_rank ) - with self.memory_saver_adapter.region( - GPU_MEMORY_TYPE_WEIGHTS, - enable_cpu_backup=enable_cpu_backup, - ): - self.loader = get_model_loader( - load_config=self.load_config, - model_config=self.model_config, + + loaded = load_model_with_memory_saver( + server_args=self.server_args, + model_config=self.model_config, + load_config=self.load_config, + device=self.device, + gpu_id=self.gpu_id, + memory_saver_adapter=self.memory_saver_adapter, + is_draft_worker=self.is_draft_worker, + ) + self.loader = loaded.loader + self.model = loaded.model + if loaded.remote_instance_weight_info is not None: + self.remote_instance_weight_transporter.weight_info = ( + loaded.remote_instance_weight_info ) - self.model = self.loader.load_model( - model_config=self.model_config, - device_config=DeviceConfig(self.device, self.gpu_id), - ) - if hasattr(self.loader, "remote_instance_transfer_engine_weight_info"): - self.remote_instance_weight_transporter.weight_info = ( - self.loader.remote_instance_transfer_engine_weight_info - ) - # Cache needs to be cleared after loading model weights (in the self.loader.load_model function). - # To avoid conflict with memory_saver_adapter.region, empty_cache operation is now moved here. - if _is_npu: - torch.npu.empty_cache() - monkey_patch_vllm_parallel_state(reverse=True) if not self.is_draft_worker: get_offloader().post_init() @@ -1220,44 +1154,11 @@ class ModelRunner(ModelRunnerKVCacheMixin): pyt_hooks = PytHooks() pyt_hooks.register_hooks(self.model, module_prefix="model") - if self.server_args.kv_cache_dtype == "fp8_e4m3": - if self.server_args.quantization_param_path is not None: - if callable(getattr(self.model, "load_kv_cache_scales", None)): - self.model.load_kv_cache_scales( - self.server_args.quantization_param_path - ) - logger.info( - "Loaded KV cache scaling factors from %s", - self.server_args.quantization_param_path, - ) - else: - raise RuntimeError( - "Using FP8 KV cache and scaling factors provided but " - "model %s does not support loading scaling factors.", - self.model.__class__, - ) - else: - logger.warning( - "Using FP8 KV cache but no scaling factors " - "provided. Defaulting to scaling factors of 1.0. " - "This may lead to less accurate results!" - ) + load_kv_cache_scales(model=self.model, server_args=self.server_args) - # Parse other args - self.sliding_window_size = None - if hasattr(self.model, "get_attention_sliding_window_size"): - self.sliding_window_size = self.model.get_attention_sliding_window_size() - elif ( - self.model_config.is_hybrid_swa - and self.model_config.sliding_window_size is not None - ): - # sliding window field in model config may have different meaning for different kinds of models (e.g., dllm), here we only consider the sliding window in SWA model - self.sliding_window_size = self.model_config.sliding_window_size - elif self.model_config.attention_chunk_size is not None: - self.sliding_window_size = self.model_config.attention_chunk_size - logger.info( - f"Setting sliding_window_size to be attention_chunk_size: {self.sliding_window_size}" - ) + self.sliding_window_size = resolve_sliding_window_size( + self.model, self.model_config + ) self.prefill_aware_swa = ( hasattr(self.model, "is_prefill_aware_swa") @@ -1281,34 +1182,17 @@ class ModelRunner(ModelRunnerKVCacheMixin): f"mem usage={self.weight_load_mem_usage:.2f} GB." ) - # TODO: Make sure all models have `quant_config` attribute, and all online quantization methods register which layers they actually quantize. - # TODO: Move this online-quantization reporting out of ModelRunner. - quantized_layers = getattr( - getattr(self.model, "quant_config", None), "quantized_layers", None - ) - if ( - self.server_args.quantization is not None - and isinstance(quantized_layers, tuple) - and len(quantized_layers) == 2 - ): - layer_types, quantized_layers_count = quantized_layers - logger.info( - f"Online {self.server_args.quantization} quantization: quantized {quantized_layers_count} layers of types: {layer_types}" - ) + report_online_quantization(model=self.model, server_args=self.server_args) - if self.server_args.debug_tensor_dump_output_folder is not None: - dump_folder = self.server_args.debug_tensor_dump_output_folder - if self.spec_algorithm.is_eagle(): - role = "draft" if self.is_draft_worker else "target" - dump_folder = os.path.join(dump_folder, role) - register_forward_hook_for_model( - self.model, - dump_folder, - self.server_args.debug_tensor_dump_layers, - self.tp_size, - self.tp_rank, - self.pp_rank, - ) + maybe_register_debug_tensor_dump_hook( + model=self.model, + server_args=self.server_args, + spec_algorithm=self.spec_algorithm, + is_draft_worker=self.is_draft_worker, + tp_size=self.tp_size, + tp_rank=self.tp_rank, + pp_rank=self.pp_rank, + ) if dumper.may_enable: dumper.apply_source_patches() @@ -1322,23 +1206,10 @@ class ModelRunner(ModelRunnerKVCacheMixin): logger, ) - if self.server_args.elastic_ep_backend == "mooncake": - # Mooncake does not support `monitored_barrier` - dist.barrier(group=get_tp_group().cpu_group) - else: - # Handle the case where some ranks do not finish loading. - try: - dist.monitored_barrier( - group=get_tp_group().cpu_group, - timeout=datetime.timedelta( - seconds=UNBALANCED_MODEL_LOADING_TIMEOUT_S - ), - wait_all_ranks=True, - ) - except RuntimeError: - raise ValueError( - f"TP rank {self.tp_rank} could finish the model loading, but there are other ranks that didn't finish loading. It is likely due to unexpected failures (e.g., OOM) or a slow node." - ) from None + dist_barrier_after_load( + elastic_ep_backend=self.server_args.elastic_ep_backend, + tp_rank=self.tp_rank, + ) def _prepare_moe_topk(self): balancer_cls = None diff --git a/python/sglang/srt/model_executor/model_runner_components/load_model_utils.py b/python/sglang/srt/model_executor/model_runner_components/load_model_utils.py new file mode 100644 index 000000000..90e08e6a9 --- /dev/null +++ b/python/sglang/srt/model_executor/model_runner_components/load_model_utils.py @@ -0,0 +1,266 @@ +from __future__ import annotations + +import datetime +import logging +import os +import socket +import threading +from typing import TYPE_CHECKING, Any, Optional + +import msgspec +import torch +import torch.distributed as dist + +from sglang.srt.configs.device_config import DeviceConfig +from sglang.srt.configs.load_config import LoadConfig, LoadFormat +from sglang.srt.constants import GPU_MEMORY_TYPE_WEIGHTS +from sglang.srt.debug_utils.tensor_dump_forward_hook import ( + register_forward_hook_for_model, +) +from sglang.srt.distributed import get_tp_group +from sglang.srt.distributed.parallel_state import monkey_patch_vllm_parallel_state +from sglang.srt.model_loader.loader import get_model_loader +from sglang.srt.model_loader.remote_instance_weight_loader_utils import ( + RemoteInstanceWeightLoaderBackend, + trigger_init_weights_send_group_for_remote_instance_request, +) +from sglang.srt.utils.common import is_npu +from sglang.srt.utils.network import NetworkAddress + +if TYPE_CHECKING: + from sglang.srt.configs.model_config import ModelConfig + from sglang.srt.server_args import ServerArgs + from sglang.srt.speculative.spec_info import SpeculativeAlgorithm + +logger = logging.getLogger(__name__) + +_is_npu = is_npu() + + +UNBALANCED_MODEL_LOADING_TIMEOUT_S = 480 # leave more time for post data processing + + +class LoadedModel(msgspec.Struct, frozen=True, kw_only=True): + loader: Any + model: Any + remote_instance_weight_info: Optional[Any] + + +def maybe_downgrade_dtype_for_legacy_gpu( + *, server_args: ServerArgs, model_config: ModelConfig +) -> None: + if torch.cuda.get_device_capability()[0] < 8: + logger.info( + "Compute capability below sm80. Use float16 due to lack of bfloat16 support." + ) + from sglang.srt.arg_groups.overrides import declare_load_time_override + + declare_load_time_override( + "ModelRunner._sm80_dtype_fallback", {"dtype": "float16"} + ) + model_config.dtype = torch.float16 + if torch.cuda.get_device_capability()[1] < 5: + raise RuntimeError("SGLang only supports sm75 and above.") + + +def maybe_trigger_remote_instance_nccl_send_group( + *, server_args: ServerArgs, tp_rank: int +) -> None: + if ( + server_args.load_format == LoadFormat.REMOTE_INSTANCE + and server_args.remote_instance_weight_loader_backend + == RemoteInstanceWeightLoaderBackend.NCCL + ): + if tp_rank == 0: + instance_ip = NetworkAddress.resolve_host(socket.gethostname()) + t = threading.Thread( + target=trigger_init_weights_send_group_for_remote_instance_request, + args=( + server_args.remote_instance_weight_loader_seed_instance_ip, + server_args.remote_instance_weight_loader_seed_instance_service_port, + server_args.remote_instance_weight_loader_send_weights_group_ports, + instance_ip, + ), + ) + t.start() + + +def load_kv_cache_scales(*, model, server_args: ServerArgs) -> None: + if server_args.kv_cache_dtype == "fp8_e4m3": + if server_args.quantization_param_path is not None: + if callable(getattr(model, "load_kv_cache_scales", None)): + model.load_kv_cache_scales(server_args.quantization_param_path) + logger.info( + "Loaded KV cache scaling factors from %s", + server_args.quantization_param_path, + ) + else: + raise RuntimeError( + "Using FP8 KV cache and scaling factors provided but " + "model %s does not support loading scaling factors.", + model.__class__, + ) + else: + logger.warning( + "Using FP8 KV cache but no scaling factors " + "provided. Defaulting to scaling factors of 1.0. " + "This may lead to less accurate results!" + ) + + +def resolve_sliding_window_size(model, model_config: ModelConfig) -> Optional[int]: + # Parse other args + sliding_window_size = None + if hasattr(model, "get_attention_sliding_window_size"): + sliding_window_size = model.get_attention_sliding_window_size() + elif model_config.is_hybrid_swa and model_config.sliding_window_size is not None: + # sliding window field in model config may have different meaning for different kinds of models (e.g., dllm), here we only consider the sliding window in SWA model + sliding_window_size = model_config.sliding_window_size + elif model_config.attention_chunk_size is not None: + sliding_window_size = model_config.attention_chunk_size + logger.info( + f"Setting sliding_window_size to be attention_chunk_size: {sliding_window_size}" + ) + return sliding_window_size + + +def report_online_quantization(*, model, server_args: ServerArgs) -> None: + # TODO: Make sure all models have `quant_config` attribute, and all online quantization methods register which layers they actually quantize. + quantized_layers = getattr( + getattr(model, "quant_config", None), "quantized_layers", None + ) + if ( + server_args.quantization is not None + and isinstance(quantized_layers, tuple) + and len(quantized_layers) == 2 + ): + layer_types, quantized_layers_count = quantized_layers + logger.info( + f"Online {server_args.quantization} quantization: quantized {quantized_layers_count} layers of types: {layer_types}" + ) + + +def maybe_register_debug_tensor_dump_hook( + *, + model, + server_args: ServerArgs, + spec_algorithm: SpeculativeAlgorithm, + is_draft_worker: bool, + tp_size: int, + tp_rank: int, + pp_rank: int, +) -> None: + if server_args.debug_tensor_dump_output_folder is not None: + dump_folder = server_args.debug_tensor_dump_output_folder + if spec_algorithm.is_eagle(): + role = "draft" if is_draft_worker else "target" + dump_folder = os.path.join(dump_folder, role) + register_forward_hook_for_model( + model, + dump_folder, + server_args.debug_tensor_dump_layers, + tp_size, + tp_rank, + pp_rank, + ) + + +def build_load_config( + *, + server_args: ServerArgs, + tp_rank: int, + remote_instance_weight_transporter_engine: Any, + remote_instance_weight_transporter_session_id: str, + draft_model_idx: Optional[int], +) -> LoadConfig: + from sglang.srt.configs.modelopt_config import ModelOptConfig + + modelopt_config = ModelOptConfig( + quant=server_args.modelopt_quant, + checkpoint_restore_path=server_args.modelopt_checkpoint_restore_path, + checkpoint_save_path=server_args.modelopt_checkpoint_save_path, + export_path=server_args.modelopt_export_path, + quantize_and_serve=server_args.quantize_and_serve, + ) + + return LoadConfig( + load_format=server_args.load_format, + download_dir=server_args.download_dir, + model_loader_extra_config=server_args.model_loader_extra_config, + tp_rank=tp_rank, + remote_instance_weight_loader_seed_instance_ip=server_args.remote_instance_weight_loader_seed_instance_ip, + remote_instance_weight_loader_seed_instance_service_port=server_args.remote_instance_weight_loader_seed_instance_service_port, + remote_instance_weight_loader_send_weights_group_ports=server_args.remote_instance_weight_loader_send_weights_group_ports, + remote_instance_weight_loader_backend=server_args.remote_instance_weight_loader_backend, + remote_instance_weight_loader_transfer_engine=remote_instance_weight_transporter_engine, + remote_instance_weight_loader_transfer_engine_session_id=remote_instance_weight_transporter_session_id, + modelexpress_url=server_args.modelexpress_url, + modelexpress_transport=server_args.modelexpress_transport, + modelopt_config=modelopt_config, + rl_quant_profile=server_args.rl_quant_profile, + draft_model_idx=draft_model_idx, + ) + + +def load_model_with_memory_saver( + *, + server_args: ServerArgs, + model_config: ModelConfig, + load_config: LoadConfig, + device: str, + gpu_id: int, + memory_saver_adapter: Any, + is_draft_worker: bool, +) -> LoadedModel: + # Remove monkey_patch when linear.py quant remove dependencies with vllm + monkey_patch_vllm_parallel_state() + + enable_cpu_backup = server_args.enable_weights_cpu_backup or ( + is_draft_worker and server_args.enable_draft_weights_cpu_backup + ) + remote_instance_weight_info = None + with memory_saver_adapter.region( + GPU_MEMORY_TYPE_WEIGHTS, + enable_cpu_backup=enable_cpu_backup, + ): + loader = get_model_loader( + load_config=load_config, + model_config=model_config, + ) + model = loader.load_model( + model_config=model_config, + device_config=DeviceConfig(device, gpu_id), + ) + if hasattr(loader, "remote_instance_transfer_engine_weight_info"): + remote_instance_weight_info = ( + loader.remote_instance_transfer_engine_weight_info + ) + # Cache needs to be cleared after loading model weights (in the loader.load_model function). + # To avoid conflict with memory_saver_adapter.region, empty_cache operation is now moved here. + if _is_npu: + torch.npu.empty_cache() + monkey_patch_vllm_parallel_state(reverse=True) + + return LoadedModel( + loader=loader, + model=model, + remote_instance_weight_info=remote_instance_weight_info, + ) + + +def dist_barrier_after_load(*, elastic_ep_backend: Optional[str], tp_rank: int) -> None: + if elastic_ep_backend == "mooncake": + # Mooncake does not support `monitored_barrier` + dist.barrier(group=get_tp_group().cpu_group) + else: + # Handle the case where some ranks do not finish loading. + try: + dist.monitored_barrier( + group=get_tp_group().cpu_group, + timeout=datetime.timedelta(seconds=UNBALANCED_MODEL_LOADING_TIMEOUT_S), + wait_all_ranks=True, + ) + except RuntimeError: + raise ValueError( + f"TP rank {tp_rank} could finish the model loading, but there are other ranks that didn't finish loading. It is likely due to unexpected failures (e.g., OOM) or a slow node." + ) from None