sglang 0.2.13__py3-none-any.whl → 0.2.14.post1__py3-none-any.whl

sglang/srt/model_executor/model_runner.py

@@ -15,11 +15,11 @@ limitations under the License.
 
 """ModelRunner runs the forward passes of the models."""
 
+import gc
 import importlib
 import importlib.resources
 import logging
 import pkgutil
-import warnings
 from functools import lru_cache
 from typing import Optional, Type
 
@@ -37,7 +37,9 @@ from vllm.distributed import (
     get_tp_group,
     init_distributed_environment,
     initialize_model_parallel,
+    set_custom_all_reduce,
 )
+from vllm.distributed.parallel_state import in_the_same_node_as
 from vllm.model_executor.model_loader import get_model
 from vllm.model_executor.models import ModelRegistry
 
@@ -88,22 +90,35 @@ class ModelRunner:
             {
                 "disable_flashinfer": server_args.disable_flashinfer,
                 "disable_flashinfer_sampling": server_args.disable_flashinfer_sampling,
-                "attention_reduce_in_fp32": server_args.attention_reduce_in_fp32,
+                "triton_attention_reduce_in_fp32": server_args.triton_attention_reduce_in_fp32,
                 "enable_mla": server_args.enable_mla,
             }
         )
 
+        min_per_gpu_memory = self.init_torch_distributed()
+        self.load_model()
+        self.init_memory_pool(
+            min_per_gpu_memory,
+            server_args.max_num_reqs,
+            server_args.max_total_tokens,
+        )
+        self.init_cublas()
+        self.init_flashinfer()
+        self.init_cuda_graphs()
+
+    def init_torch_distributed(self):
         # Init torch distributed
         torch.cuda.set_device(self.gpu_id)
-        logger.info(f"[gpu={self.gpu_id}] Init nccl begin.")
+        logger.info("Init nccl begin.")
 
-        if not server_args.enable_p2p_check:
+        if not self.server_args.enable_p2p_check:
             monkey_patch_vllm_p2p_access_check(self.gpu_id)
 
-        if server_args.nccl_init_addr:
-            nccl_init_method = f"tcp://{server_args.nccl_init_addr}"
+        if self.server_args.nccl_init_addr:
+            nccl_init_method = f"tcp://{self.server_args.nccl_init_addr}"
         else:
             nccl_init_method = f"tcp://127.0.0.1:{self.nccl_port}"
+        set_custom_all_reduce(not self.server_args.disable_custom_all_reduce)
         init_distributed_environment(
             backend="nccl",
             world_size=self.tp_size,
@@ -112,43 +127,45 @@ class ModelRunner:
             distributed_init_method=nccl_init_method,
         )
         initialize_model_parallel(tensor_model_parallel_size=self.tp_size)
-        self.tp_group = get_tp_group()
-        total_gpu_memory = get_available_gpu_memory(
+        min_per_gpu_memory = get_available_gpu_memory(
             self.gpu_id, distributed=self.tp_size > 1
         )
+        self.tp_group = get_tp_group()
 
+        # Currently, there is a bug with mulit-node tensor parallelsim + padded cuda graph,
+        # so we disable padding in cuda graph.
+        if not all(in_the_same_node_as(self.tp_group.cpu_group, source_rank=0)):
+            self.server_args.disable_cuda_graph_padding = True
+            logger.info(
+                "Setting disable_cuda_graph_padding to True because of multi-node tensor parallelism."
+            )
+
+        # Check memory for tensor parallelism
         if self.tp_size > 1:
-            total_local_gpu_memory = get_available_gpu_memory(self.gpu_id)
-            if total_local_gpu_memory < total_gpu_memory * 0.9:
+            local_gpu_memory = get_available_gpu_memory(self.gpu_id)
+            if min_per_gpu_memory < local_gpu_memory * 0.9:
                 raise ValueError(
                     "The memory capacity is unbalanced. Some GPUs may be occupied by other processes."
                 )
 
-        # Load the model and create memory pool
-        self.load_model()
-        self.init_memory_pool(
-            total_gpu_memory,
-            server_args.max_num_reqs,
-            server_args.max_total_tokens,
-        )
-        self.init_cublas()
-        self.init_flashinfer()
-
-        if self.is_generation:
-            # FIXME Currently, cuda graph only capture decode steps, which only exists in causal models
-            # Capture cuda graphs
-            self.init_cuda_graphs()
+        return min_per_gpu_memory
 
     def load_model(self):
         logger.info(
-            f"[gpu={self.gpu_id}] Load weight begin. "
-            f"avail mem={get_available_gpu_memory(self.gpu_id):.2f} GB"
+            f"Load weight begin. avail mem={get_available_gpu_memory(self.gpu_id):.2f} GB"
         )
+        if torch.cuda.get_device_capability()[0] < 8:
+            logger.info(
+                "Compute capability below sm80. Use float16 due to lack of bfloat16 support."
+            )
+            self.server_args.dtype = "float16"
+            if torch.cuda.get_device_capability()[1] < 5:
+                raise RuntimeError("SGLang only supports sm75 and above.")
 
         monkey_patch_vllm_dummy_weight_loader()
-        device_config = DeviceConfig()
-        load_config = LoadConfig(load_format=self.server_args.load_format)
-        vllm_model_config = VllmModelConfig(
+        self.device_config = DeviceConfig()
+        self.load_config = LoadConfig(load_format=self.server_args.load_format)
+        self.vllm_model_config = VllmModelConfig(
             model=self.server_args.model_path,
             quantization=self.server_args.quantization,
             tokenizer=None,
@@ -159,43 +176,132 @@ class ModelRunner:
             skip_tokenizer_init=True,
         )
 
+        # A temporary hack to fix the num_heads for meta-llama/Meta-Llama-3.1-405B-FP8 checkpoints
+        # Drop this after Sept, 2024.
         if is_llama3_405b_fp8_head_16(self.model_config) and self.tp_size <= 8:
-            # A temporary hack to fix the num_heads for meta-llama/Meta-Llama-3.1-405B-FP8 checkpoints
             self.model_config.hf_config.num_key_value_heads = 8
-            vllm_model_config.hf_config.num_key_value_heads = 8
+            self.vllm_model_config.hf_config.num_key_value_heads = 8
             monkey_patch_vllm_qvk_linear_loader()
 
-        self.dtype = vllm_model_config.dtype
+        self.dtype = self.vllm_model_config.dtype
         if self.model_config.model_overide_args is not None:
-            vllm_model_config.hf_config.update(self.model_config.model_overide_args)
+            self.vllm_model_config.hf_config.update(
+                self.model_config.model_overide_args
+            )
 
         self.model = get_model(
-            model_config=vllm_model_config,
-            device_config=device_config,
-            load_config=load_config,
-            lora_config=None,
-            multimodal_config=None,
+            model_config=self.vllm_model_config,
+            load_config=self.load_config,
+            device_config=self.device_config,
             parallel_config=None,
             scheduler_config=None,
+            lora_config=None,
             cache_config=None,
         )
         self.sliding_window_size = (
-            self.model.get_window_size()
-            if hasattr(self.model, "get_window_size")
+            self.model.get_attention_sliding_window_size()
+            if hasattr(self.model, "get_attention_sliding_window_size")
             else None
         )
         self.is_generation = is_generation_model(
-            self.model_config.hf_config.architectures
+            self.model_config.hf_config.architectures, self.server_args.is_embedding
         )
 
         logger.info(
-            f"[gpu={self.gpu_id}] Load weight end. "
+            f"Load weight end. "
             f"type={type(self.model).__name__}, "
             f"dtype={self.dtype}, "
             f"avail mem={get_available_gpu_memory(self.gpu_id):.2f} GB"
         )
 
-    def profile_max_num_token(self, total_gpu_memory):
+    def update_weights(self, model_path: str, load_format: str):
+        """Update weights in-place."""
+        from vllm.model_executor.model_loader.loader import (
+            DefaultModelLoader,
+            device_loading_context,
+            get_model_loader,
+        )
+        from vllm.model_executor.model_loader.utils import set_default_torch_dtype
+
+        logger.info(
+            f"Update weights begin. "
+            f"avail mem={get_available_gpu_memory(self.gpu_id):.2f} GB"
+        )
+
+        target_device = torch.device(self.device_config.device)
+
+        try:
+            # TODO: Use a better method to check this
+            vllm_model_config = VllmModelConfig(
+                model=model_path,
+                quantization=self.server_args.quantization,
+                tokenizer=None,
+                tokenizer_mode=None,
+                trust_remote_code=self.server_args.trust_remote_code,
+                dtype=self.server_args.dtype,
+                seed=42,
+                skip_tokenizer_init=True,
+            )
+        except Exception as e:
+            logger.error(f"Failed to load model config: {e}")
+            return False, "Failed to update model weights"
+
+        load_config = LoadConfig(load_format=load_format)
+
+        # Only support vllm DefaultModelLoader for now
+        loader = get_model_loader(load_config)
+        if not isinstance(loader, DefaultModelLoader):
+            logger.error("Failed to get weights iterator: Unsupported loader")
+            return False, "Failed to update model weights"
+
+        def get_weight_iter(config):
+            iter = loader._get_weights_iterator(
+                config.model,
+                config.revision,
+                fall_back_to_pt=getattr(
+                    self.model, "fall_back_to_pt_during_load", True
+                ),
+            )
+            return iter
+
+        def model_load_weights(model, iter):
+            model.load_weights(iter)
+            for _, module in self.model.named_modules():
+                quant_method = getattr(module, "quant_method", None)
+                if quant_method is not None:
+                    with device_loading_context(module, target_device):
+                        quant_method.process_weights_after_loading(module)
+            return model
+
+        with set_default_torch_dtype(vllm_model_config.dtype):
+            try:
+                iter = get_weight_iter(vllm_model_config)
+            except Exception as e:
+                message = f"Failed to get weights iterator: {e}"
+                logger.error(message)
+                return False, message
+            try:
+                model = model_load_weights(self.model, iter)
+            except Exception as e:
+                message = f"Failed to update weights: {e}. \n Rolling back to original weights"
+                logger.error(message)
+                del iter
+                gc.collect()
+                iter = get_weight_iter(self.vllm_model_config)
+                self.model = model_load_weights(self.model, iter)
+                return False, message
+
+        self.model = model
+        self.server_args.model_path = model_path
+        self.server_args.load_format = load_format
+        self.vllm_model_config = vllm_model_config
+        self.load_config = load_config
+        self.model_config.path = model_path
+
+        logger.info("Update weights end.")
+        return True, "Succeeded to update model weights"
+
+    def profile_max_num_token(self, total_gpu_memory: int):
         available_gpu_memory = get_available_gpu_memory(
             self.gpu_id, distributed=self.tp_size > 1
         )
@@ -206,7 +312,7 @@ class ModelRunner:
             cell_size = (
                 (self.model_config.kv_lora_rank + self.model_config.qk_rope_head_dim)
                 * self.model_config.num_hidden_layers
-                * torch._utils._element_size(self.dtype)
+                * torch._utils._element_size(self.kv_cache_dtype)
             )
         else:
             cell_size = (
@@ -214,7 +320,7 @@ class ModelRunner:
                 * self.model_config.head_dim
                 * self.model_config.num_hidden_layers
                 * 2
-                * torch._utils._element_size(self.dtype)
+                * torch._utils._element_size(self.kv_cache_dtype)
             )
         rest_memory = available_gpu_memory - total_gpu_memory * (
             1 - self.mem_fraction_static
@@ -223,12 +329,30 @@ class ModelRunner:
         return max_num_token
 
     def init_memory_pool(
-        self, total_gpu_memory, max_num_reqs=None, max_total_tokens=None
+        self,
+        total_gpu_memory: int,
+        max_num_reqs: int = None,
+        max_total_tokens: int = None,
     ):
+        if self.server_args.kv_cache_dtype == "auto":
+            self.kv_cache_dtype = self.dtype
+        elif self.server_args.kv_cache_dtype == "fp8_e5m2":
+            if self.server_args.disable_flashinfer or self.server_args.enable_mla:
+                logger.warning(
+                    "FP8 KV cache is not supported for Triton kernel now, using auto kv cache dtype"
+                )
+                self.kv_cache_dtype = self.dtype
+            else:
+                self.kv_cache_dtype = torch.float8_e5m2
+        else:
+            raise ValueError(
+                f"Unsupported kv_cache_dtype: {self.server_args.kv_cache_dtype}."
+            )
+
         self.max_total_num_tokens = self.profile_max_num_token(total_gpu_memory)
         if max_total_tokens is not None:
             if max_total_tokens > self.max_total_num_tokens:
-                warnings.warn(
+                logging.warning(
                     f"max_total_tokens={max_total_tokens} is larger than the profiled value "
                     f"{self.max_total_num_tokens}. "
                     f"Use the profiled value instead."
@@ -261,7 +385,7 @@ class ModelRunner:
         ):
             self.token_to_kv_pool = MLATokenToKVPool(
                 self.max_total_num_tokens,
-                dtype=self.dtype,
+                dtype=self.kv_cache_dtype,
                 kv_lora_rank=self.model_config.kv_lora_rank,
                 qk_rope_head_dim=self.model_config.qk_rope_head_dim,
                 layer_num=self.model_config.num_hidden_layers,
@@ -272,13 +396,13 @@ class ModelRunner:
         else:
             self.token_to_kv_pool = MHATokenToKVPool(
                 self.max_total_num_tokens,
-                dtype=self.dtype,
+                dtype=self.kv_cache_dtype,
                 head_num=self.model_config.get_num_kv_heads(self.tp_size),
                 head_dim=self.model_config.head_dim,
                 layer_num=self.model_config.num_hidden_layers,
             )
         logger.info(
-            f"[gpu={self.gpu_id}] Memory pool end. "
+            f"Memory pool end. "
             f"avail mem={get_available_gpu_memory(self.gpu_id):.2f} GB"
         )
 
@@ -292,6 +416,7 @@ class ModelRunner:
         return c
 
     def init_flashinfer(self):
+        """Init flashinfer attention kernel wrappers."""
         if self.server_args.disable_flashinfer:
             assert (
                 self.sliding_window_size is None
@@ -352,20 +477,29 @@ class ModelRunner:
                 )
 
     def init_cuda_graphs(self):
+        """Capture cuda graphs."""
+        if not self.is_generation:
+            # TODO: Currently, cuda graph only captures decode steps, which only exists for generation models
+            return
+
         from sglang.srt.model_executor.cuda_graph_runner import CudaGraphRunner
 
         if self.server_args.disable_cuda_graph or self.server_args.disable_flashinfer:
             self.cuda_graph_runner = None
             return
 
-        logger.info(
-            f"[gpu={self.gpu_id}] Capture cuda graph begin. This can take up to several minutes."
-        )
-        batch_size_list = [1, 2, 4] + [i * 8 for i in range(1, 17)]
+        logger.info("Capture cuda graph begin. This can take up to several minutes.")
+
+        if self.server_args.disable_cuda_graph_padding:
+            batch_size_list = list(range(1, 32)) + [64, 128]
+        else:
+            batch_size_list = [1, 2, 4] + [i * 8 for i in range(1, 21)]
+
         self.cuda_graph_runner = CudaGraphRunner(
             self,
             max_batch_size_to_capture=max(batch_size_list),
             use_torch_compile=self.server_args.enable_torch_compile,
+            disable_padding=self.server_args.disable_cuda_graph_padding,
         )
         try:
             self.cuda_graph_runner.capture(batch_size_list)
@@ -401,9 +535,18 @@ class ModelRunner:
             batch,
             forward_mode=ForwardMode.EXTEND,
         )
-        return self.model.forward(
-            batch.input_ids, input_metadata.positions, input_metadata
-        )
+        if self.is_generation:
+            return self.model.forward(
+                batch.input_ids, input_metadata.positions, input_metadata
+            )
+        else:
+            # Only embedding models have get_embedding parameter
+            return self.model.forward(
+                batch.input_ids,
+                input_metadata.positions,
+                input_metadata,
+                get_embedding=True,
+            )
 
     @torch.inference_mode()
     def forward_extend_multi_modal(self, batch: ScheduleBatch):
@@ -477,4 +620,4 @@ def load_model_cls_srt(model_arch: str) -> Optional[Type[nn.Module]]:
 
 
 # Monkey patch model loader
-setattr(ModelRegistry, "load_model_cls", load_model_cls_srt)
+setattr(ModelRegistry, "_try_load_model_cls", load_model_cls_srt)

sglang/srt/models/gemma2.py

@@ -25,7 +25,6 @@ from vllm.distributed import get_tensor_model_parallel_world_size
 
 # FIXME: temporary solution, remove after next vllm release
 from vllm.model_executor.custom_op import CustomOp
-from vllm.model_executor.layers.activation import GeluAndMul
 
 # from vllm.model_executor.layers.layernorm import GemmaRMSNorm
 from vllm.model_executor.layers.linear import (
@@ -39,6 +38,7 @@ from vllm.model_executor.layers.quantization.base_config import QuantizationConf
 from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding
 from vllm.model_executor.model_loader.weight_utils import default_weight_loader
 
+from sglang.srt.layers.activation import GeluAndMul
 from sglang.srt.layers.logits_processor import LogitsProcessor
 from sglang.srt.layers.radix_attention import RadixAttention
 from sglang.srt.model_executor.forward_batch_info import InputMetadata
@@ -46,7 +46,7 @@ from sglang.srt.model_executor.forward_batch_info import InputMetadata
 
 # Aligned with HF's implementation, using sliding window inclusive with the last token
 # SGLang assumes exclusive
-def get_window_size(config):
+def get_attention_sliding_window_size(config):
     return config.sliding_window - 1
 
 
@@ -135,7 +135,7 @@ class Gemma2MLP(nn.Module):
                 "function. Please set `hidden_act` and `hidden_activation` to "
                 "`gelu_pytorch_tanh`."
             )
-        self.act_fn = GeluAndMul(approximate="tanh")
+        self.act_fn = GeluAndMul()
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         gate_up, _ = self.gate_up_proj(x)
@@ -213,7 +213,11 @@ class Gemma2Attention(nn.Module):
             self.scaling,
             num_kv_heads=self.num_kv_heads,
             layer_id=layer_idx,
-            sliding_window_size=get_window_size(config) if use_sliding_window else None,
+            sliding_window_size=(
+                get_attention_sliding_window_size(config)
+                if use_sliding_window
+                else None
+            ),
             logit_cap=self.config.attn_logit_softcapping,
         )
 
@@ -406,8 +410,8 @@ class Gemma2ForCausalLM(nn.Module):
             input_ids, hidden_states, self.model.embed_tokens.weight, input_metadata
         )
 
-    def get_window_size(self):
-        return get_window_size(self.config)
+    def get_attention_sliding_window_size(self):
+        return get_attention_sliding_window_size(self.config)
 
     def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]):
         stacked_params_mapping = [

sglang/srt/models/gpt_bigcode.py

@@ -23,7 +23,6 @@ from torch import nn
 from transformers import GPTBigCodeConfig
 from vllm.config import CacheConfig, LoRAConfig
 from vllm.distributed import get_tensor_model_parallel_world_size
-from vllm.model_executor.layers.activation import get_act_fn
 from vllm.model_executor.layers.linear import (
     ColumnParallelLinear,
     QKVParallelLinear,
@@ -33,6 +32,7 @@ from vllm.model_executor.layers.quantization.base_config import QuantizationConf
 from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding
 from vllm.model_executor.model_loader.weight_utils import default_weight_loader
 
+from sglang.srt.layers.activation import get_act_fn
 from sglang.srt.layers.logits_processor import LogitsProcessor
 from sglang.srt.layers.radix_attention import RadixAttention
 from sglang.srt.model_executor.forward_batch_info import InputMetadata

sglang/srt/models/grok.py

@@ -300,6 +300,9 @@ class Grok1ModelForCausalLM(nn.Module):
 
         # Monkey patch _prepare_weights to load pre-sharded weights
         setattr(DefaultModelLoader, "_prepare_weights", _prepare_presharded_weights)
+
+        self.use_presharded_weights = True
+
         warnings.filterwarnings("ignore", category=FutureWarning)
 
     def forward(
@@ -355,6 +358,13 @@ class Grok1ModelForCausalLM(nn.Module):
                         continue
                     name = name.replace(weight_name, param_name)
 
+                    if self.use_presharded_weights:
+                        extra_kwargs = {
+                            "use_presharded_weights": self.use_presharded_weights
+                        }
+                    else:
+                        extra_kwargs = {}
+
                     param = params_dict[name]
                     weight_loader = param.weight_loader
                     weight_loader(
@@ -363,7 +373,7 @@ class Grok1ModelForCausalLM(nn.Module):
                         weight_name,
                         shard_id=shard_id,
                         expert_id=expert_id,
-                        pre_sharded=get_tensor_model_parallel_world_size() > 1,
+                        **extra_kwargs,
                     )
                     break
                 else:

sglang/srt/models/llama_embedding.py

@@ -29,7 +29,11 @@ class LlamaEmbeddingModel(nn.Module):
         positions: torch.Tensor,
         input_metadata: InputMetadata,
         input_embeds: torch.Tensor = None,
+        get_embedding: bool = True,
     ) -> EmbeddingPoolerOutput:
+        assert (
+            get_embedding
+        ), "LlamaEmbeddingModel / MistralModel is only used for embedding"
         hidden_states = self.model(input_ids, positions, input_metadata, input_embeds)
         return self.pooler(hidden_states, input_metadata)