sglang 0.1.19__py3-none-any.whl → 0.1.21__py3-none-any.whl

sglang/srt/models/qwen2_moe.py

@@ -8,24 +8,28 @@ import torch
 import torch.nn.functional as F
 from torch import nn
 from transformers import PretrainedConfig
-
 from vllm.config import CacheConfig
-from vllm.distributed import (get_tensor_model_parallel_world_size,
-                              tensor_model_parallel_all_reduce)
+from vllm.distributed import (
+    get_tensor_model_parallel_world_size,
+    tensor_model_parallel_all_reduce,
+)
 from vllm.model_executor.layers.activation import SiluAndMul
 from vllm.model_executor.layers.fused_moe import FusedMoE
 from vllm.model_executor.layers.layernorm import RMSNorm
-from vllm.model_executor.layers.linear import (MergedColumnParallelLinear,
-                                               QKVParallelLinear,
-                                               ReplicatedLinear,
-                                               RowParallelLinear)
+from vllm.model_executor.layers.linear import (
+    MergedColumnParallelLinear,
+    QKVParallelLinear,
+    ReplicatedLinear,
+    RowParallelLinear,
+)
 from vllm.model_executor.layers.logits_processor import LogitsProcessor
-from vllm.model_executor.layers.quantization.base_config import (
-    QuantizationConfig)
+from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
 from vllm.model_executor.layers.rotary_embedding import get_rope
 from vllm.model_executor.layers.sampler import Sampler
 from vllm.model_executor.layers.vocab_parallel_embedding import (
-    ParallelLMHead, VocabParallelEmbedding)
+    ParallelLMHead,
+    VocabParallelEmbedding,
+)
 from vllm.model_executor.model_loader.weight_utils import default_weight_loader
 from vllm.model_executor.sampling_metadata import SamplingMetadata
 from vllm.sequence import IntermediateTensors, SamplerOutput
@@ -34,8 +38,8 @@ from sglang.srt.layers.logits_processor import LogitsProcessor
 from sglang.srt.layers.radix_attention import RadixAttention
 from sglang.srt.managers.controller.model_runner import InputMetadata
 
-class Qwen2MoeMLP(nn.Module):
 
+class Qwen2MoeMLP(nn.Module):
     def __init__(
         self,
         hidden_size: int,
@@ -46,17 +50,20 @@ class Qwen2MoeMLP(nn.Module):
     ) -> None:
         super().__init__()
         self.gate_up_proj = MergedColumnParallelLinear(
-            hidden_size, [intermediate_size] * 2,
+            hidden_size, [intermediate_size] * 2, bias=False, quant_config=quant_config
+        )
+        self.down_proj = RowParallelLinear(
+            intermediate_size,
+            hidden_size,
             bias=False,
-            quant_config=quant_config)
-        self.down_proj = RowParallelLinear(intermediate_size,
-                                           hidden_size,
-                                           bias=False,
-                                           quant_config=quant_config,
-                                           reduce_results=reduce_results)
+            quant_config=quant_config,
+            reduce_results=reduce_results,
+        )
         if hidden_act != "silu":
-            raise ValueError(f"Unsupported activation: {hidden_act}. "
-                             "Only silu is supported for now.")
+            raise ValueError(
+                f"Unsupported activation: {hidden_act}. "
+                "Only silu is supported for now."
+            )
         self.act_fn = SiluAndMul()
 
     def forward(self, x):
@@ -67,7 +74,6 @@ class Qwen2MoeMLP(nn.Module):
 
 
 class Qwen2MoeSparseMoeBlock(nn.Module):
-
     def __init__(
         self,
         config: PretrainedConfig,
@@ -79,20 +85,22 @@ class Qwen2MoeSparseMoeBlock(nn.Module):
         if self.tp_size > config.num_experts:
             raise ValueError(
                 f"Tensor parallel size {self.tp_size} is greater than "
-                f"the number of experts {config.num_experts}.")
-
-        self.experts = FusedMoE(num_experts=config.num_experts,
-                                top_k=config.num_experts_per_tok,
-                                hidden_size=config.hidden_size,
-                                intermediate_size=config.moe_intermediate_size,
-                                reduce_results=False,
-                                renormalize=config.norm_topk_prob,
-                                quant_config=quant_config)
-
-        self.gate = ReplicatedLinear(config.hidden_size,
-                                     config.num_experts,
-                                     bias=False,
-                                     quant_config=None)
+                f"the number of experts {config.num_experts}."
+            )
+
+        self.experts = FusedMoE(
+            num_experts=config.num_experts,
+            top_k=config.num_experts_per_tok,
+            hidden_size=config.hidden_size,
+            intermediate_size=config.moe_intermediate_size,
+            reduce_results=False,
+            renormalize=config.norm_topk_prob,
+            quant_config=quant_config,
+        )
+
+        self.gate = ReplicatedLinear(
+            config.hidden_size, config.num_experts, bias=False, quant_config=None
+        )
         if config.shared_expert_intermediate_size > 0:
             self.shared_expert = Qwen2MoeMLP(
                 hidden_size=config.hidden_size,
@@ -103,9 +111,7 @@ class Qwen2MoeSparseMoeBlock(nn.Module):
             )
         else:
             self.shared_expert = None
-        self.shared_expert_gate = torch.nn.Linear(config.hidden_size,
-                                                  1,
-                                                  bias=False)
+        self.shared_expert_gate = torch.nn.Linear(config.hidden_size, 1, bias=False)
 
     def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
         num_tokens, hidden_dim = hidden_states.shape
@@ -114,24 +120,24 @@ class Qwen2MoeSparseMoeBlock(nn.Module):
         if self.shared_expert is not None:
             shared_output = self.shared_expert(hidden_states)
             if self.shared_expert_gate is not None:
-                shared_output = F.sigmoid(
-                    self.shared_expert_gate(hidden_states)) * shared_output
+                shared_output = (
+                    F.sigmoid(self.shared_expert_gate(hidden_states)) * shared_output
+                )
 
         # router_logits: (num_tokens, n_experts)
         router_logits, _ = self.gate(hidden_states)
-        final_hidden_states = self.experts(hidden_states=hidden_states,
-                                           router_logits=router_logits)
+        final_hidden_states = self.experts(
+            hidden_states=hidden_states, router_logits=router_logits
+        )
         if shared_output is not None:
             final_hidden_states = final_hidden_states + shared_output
         if self.tp_size > 1:
-            final_hidden_states = tensor_model_parallel_all_reduce(
-                final_hidden_states)
+            final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states)
 
         return final_hidden_states.view(num_tokens, hidden_dim)
 
 
 class Qwen2MoeAttention(nn.Module):
-
     def __init__(
         self,
         hidden_size: int,
@@ -190,17 +196,19 @@ class Qwen2MoeAttention(nn.Module):
             base=rope_theta,
             rope_scaling=rope_scaling,
         )
-        self.attn = RadixAttention(self.num_heads,
-                                   self.head_dim,
-                                   self.scaling,
-                                   num_kv_heads=self.num_kv_heads,
-                                   layer_id=layer_id)
+        self.attn = RadixAttention(
+            self.num_heads,
+            self.head_dim,
+            self.scaling,
+            num_kv_heads=self.num_kv_heads,
+            layer_id=layer_id,
+        )
 
     def forward(
         self,
         positions: torch.Tensor,
         hidden_states: torch.Tensor,
-        input_metadata: InputMetadata
+        input_metadata: InputMetadata,
     ) -> torch.Tensor:
         qkv, _ = self.qkv_proj(hidden_states)
         q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
@@ -211,7 +219,6 @@ class Qwen2MoeAttention(nn.Module):
 
 
 class Qwen2MoeDecoderLayer(nn.Module):
-
     def __init__(
         self,
         config: PretrainedConfig,
@@ -223,8 +230,7 @@ class Qwen2MoeDecoderLayer(nn.Module):
         self.hidden_size = config.hidden_size
         rope_theta = getattr(config, "rope_theta", 10000)
         rope_scaling = getattr(config, "rope_scaling", None)
-        max_position_embeddings = getattr(config, "max_position_embeddings",
-                                          8192)
+        max_position_embeddings = getattr(config, "max_position_embeddings", 8192)
         self.self_attn = Qwen2MoeAttention(
             hidden_size=self.hidden_size,
             num_heads=config.num_attention_heads,
@@ -239,13 +245,13 @@ class Qwen2MoeDecoderLayer(nn.Module):
 
         # Note: Qwen/Qwen2-57B-A14B-Instruct does not have
         # `mlp_only_layers` in the config.
-        mlp_only_layers = ([] if not hasattr(config, "mlp_only_layers") else
-                           config.mlp_only_layers)
+        mlp_only_layers = (
+            [] if not hasattr(config, "mlp_only_layers") else config.mlp_only_layers
+        )
         if (layer_id not in mlp_only_layers) and (
-                config.num_experts > 0 and
-            (layer_id + 1) % config.decoder_sparse_step == 0):
-            self.mlp = Qwen2MoeSparseMoeBlock(config=config,
-                                              quant_config=quant_config)
+            config.num_experts > 0 and (layer_id + 1) % config.decoder_sparse_step == 0
+        ):
+            self.mlp = Qwen2MoeSparseMoeBlock(config=config, quant_config=quant_config)
         else:
             self.mlp = Qwen2MoeMLP(
                 hidden_size=config.hidden_size,
@@ -253,10 +259,10 @@ class Qwen2MoeDecoderLayer(nn.Module):
                 hidden_act=config.hidden_act,
                 quant_config=quant_config,
             )
-        self.input_layernorm = RMSNorm(config.hidden_size,
-                                       eps=config.rms_norm_eps)
-        self.post_attention_layernorm = RMSNorm(config.hidden_size,
-                                                eps=config.rms_norm_eps)
+        self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.post_attention_layernorm = RMSNorm(
+            config.hidden_size, eps=config.rms_norm_eps
+        )
 
     def forward(
         self,
@@ -270,23 +276,20 @@ class Qwen2MoeDecoderLayer(nn.Module):
             residual = hidden_states
             hidden_states = self.input_layernorm(hidden_states)
         else:
-            hidden_states, residual = self.input_layernorm(
-                hidden_states, residual)
+            hidden_states, residual = self.input_layernorm(hidden_states, residual)
         hidden_states = self.self_attn(
             positions=positions,
             hidden_states=hidden_states,
-            input_metadata=input_metadata
+            input_metadata=input_metadata,
         )
 
         # Fully Connected
-        hidden_states, residual = self.post_attention_layernorm(
-            hidden_states, residual)
+        hidden_states, residual = self.post_attention_layernorm(hidden_states, residual)
         hidden_states = self.mlp(hidden_states)
         return hidden_states, residual
 
 
 class Qwen2MoeModel(nn.Module):
-
     def __init__(
         self,
         config: PretrainedConfig,
@@ -301,13 +304,14 @@ class Qwen2MoeModel(nn.Module):
             config.vocab_size,
             config.hidden_size,
         )
-        self.layers = nn.ModuleList([
-            Qwen2MoeDecoderLayer(config,
-                                 layer_id,
-                                 cache_config,
-                                 quant_config=quant_config)
-            for layer_id in range(config.num_hidden_layers)
-        ])
+        self.layers = nn.ModuleList(
+            [
+                Qwen2MoeDecoderLayer(
+                    config, layer_id, cache_config, quant_config=quant_config
+                )
+                for layer_id in range(config.num_hidden_layers)
+            ]
+        )
         self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
 
     def forward(
@@ -315,7 +319,7 @@ class Qwen2MoeModel(nn.Module):
         input_ids: torch.Tensor,
         positions: torch.Tensor,
         input_metadata: InputMetadata,
-        input_embeds: torch.Tensor = None
+        input_embeds: torch.Tensor = None,
     ) -> torch.Tensor:
         if input_embeds is None:
             hidden_states = self.embed_tokens(input_ids)
@@ -324,10 +328,9 @@ class Qwen2MoeModel(nn.Module):
         residual = None
         for i in range(len(self.layers)):
             layer = self.layers[i]
-            hidden_states, residual = layer(positions,
-                                            hidden_states,
-                                            input_metadata,
-                                            residual)
+            hidden_states, residual = layer(
+                positions, hidden_states, input_metadata, residual
+            )
         hidden_states, _ = self.norm(hidden_states, residual)
         return hidden_states
 
@@ -346,9 +349,9 @@ class Qwen2MoeForCausalLM(nn.Module):
         self.config = config
         self.quant_config = quant_config
         self.model = Qwen2MoeModel(config, cache_config, quant_config)
-        self.lm_head = ParallelLMHead(config.vocab_size,
-                                      config.hidden_size,
-                                      quant_config=quant_config)
+        self.lm_head = ParallelLMHead(
+            config.vocab_size, config.hidden_size, quant_config=quant_config
+        )
         self.logits_processor = LogitsProcessor(config)
         self.sampler = Sampler()
 
@@ -357,17 +360,22 @@ class Qwen2MoeForCausalLM(nn.Module):
         input_ids: torch.Tensor,
         positions: torch.Tensor,
         input_metadata: InputMetadata,
-        input_embeds: torch.Tensor = None
+        input_embeds: torch.Tensor = None,
+    ) -> torch.Tensor:
+        hidden_states = self.model(input_ids, positions, input_metadata, input_embeds)
+        return self.logits_processor(
+            input_ids, hidden_states, self.lm_head.weight, input_metadata
+        )
+
+    def compute_logits(
+        self,
+        input_ids: torch.Tensor,
+        hidden_states: torch.Tensor,
+        input_metadata: InputMetadata,
     ) -> torch.Tensor:
-        hidden_states = self.model(input_ids, positions, input_metadata,
-                                   input_embeds)
-        return self.logits_processor(input_ids, hidden_states, self.lm_head.weight,
-                                     input_metadata)
-
-    def compute_logits(self, input_ids: torch.Tensor, hidden_states: torch.Tensor,
-                       input_metadata: InputMetadata) -> torch.Tensor:
-        logits = self.logits_processor(input_ids, hidden_states, self.lm_head.weight,
-                                       input_metadata)
+        logits = self.logits_processor(
+            input_ids, hidden_states, self.lm_head.weight, input_metadata
+        )
         return logits
 
     def sample(
@@ -391,11 +399,18 @@ class Qwen2MoeForCausalLM(nn.Module):
         expert_params_mapping = [
             # These are the weights for the experts
             # (param_name, weight_name, expert_id, shard_id)
-            ("experts.w13_weight" if weight_name in ["gate_proj", "up_proj"]
-             else "experts.w2_weight",
-             f"experts.{expert_id}.{weight_name}.weight", expert_id, shard_id)
-            for expert_id in range(self.config.num_experts) for shard_id,
-            weight_name in enumerate(["gate_proj", "down_proj", "up_proj"])
+            (
+                "experts.w13_weight"
+                if weight_name in ["gate_proj", "up_proj"]
+                else "experts.w2_weight",
+                f"experts.{expert_id}.{weight_name}.weight",
+                expert_id,
+                shard_id,
+            )
+            for expert_id in range(self.config.num_experts)
+            for shard_id, weight_name in enumerate(
+                ["gate_proj", "down_proj", "up_proj"]
+            )
         ]
 
         params_dict = dict(self.named_parameters())
@@ -433,11 +448,13 @@ class Qwen2MoeForCausalLM(nn.Module):
                     name = name.replace(weight_name, param_name)
                     param = params_dict[name]
                     weight_loader = param.weight_loader
-                    weight_loader(param,
-                                  loaded_weight,
-                                  weight_name,
-                                  shard_id=shard_id,
-                                  expert_id=expert_id)
+                    weight_loader(
+                        param,
+                        loaded_weight,
+                        weight_name,
+                        shard_id=shard_id,
+                        expert_id=expert_id,
+                    )
                     break
                 else:
                     # Skip loading extra bias for GPTQ models.
@@ -447,8 +464,10 @@ class Qwen2MoeForCausalLM(nn.Module):
                         continue
 
                     param = params_dict[name]
-                    weight_loader = getattr(param, "weight_loader",
-                                            default_weight_loader)
+                    weight_loader = getattr(
+                        param, "weight_loader", default_weight_loader
+                    )
                     weight_loader(param, loaded_weight)
 
+
 EntryClass = Qwen2MoeForCausalLM

sglang/srt/server.py

@@ -33,9 +33,9 @@ from sglang.srt.managers.controller.manager_multi import (
     start_controller_process as start_controller_process_multi,
 )
 from sglang.srt.managers.controller.manager_single import (
+    launch_tp_servers,
     start_controller_process as start_controller_process_single,
 )
-from sglang.srt.managers.controller.tp_worker import ModelTpService
 from sglang.srt.managers.detokenizer_manager import start_detokenizer_process
 from sglang.srt.managers.io_struct import GenerateReqInput
 from sglang.srt.managers.tokenizer_manager import TokenizerManager
@@ -53,7 +53,6 @@ from sglang.srt.utils import (
     enable_show_time_cost,
     receive_addrs,
     send_addrs_to_rank_0,
-    start_rpyc_service_process,
 )
 from sglang.utils import get_exception_traceback
 
@@ -146,6 +145,7 @@ def launch_server(server_args: ServerArgs, pipe_finish_writer, model_overide_arg
 
     # Set global environments
     os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"
+    os.environ["NCCL_CUMEM_ENABLE"] = "0"
     if server_args.show_time_cost:
         enable_show_time_cost()
     if server_args.disable_disk_cache:
@@ -191,21 +191,17 @@ def launch_server(server_args: ServerArgs, pipe_finish_writer, model_overide_arg
         model_port_args=model_port_args,
     )
 
-    # TODO multi-node dp is not supported
-    assert not (server_args.dp_size > 1 and server_args.node_rank is not None)
+    # Handle multi-node tp
     if server_args.nnodes > 1:
+        assert server_args.dp_size == 1, "Multi-node dp is not supported."
+
         if server_args.node_rank != 0:
-            send_addrs_to_rank_0(model_port_args[0], server_args)
-        else:
-            receive_addrs(model_port_args[0], server_args)
-        for i in range(tp_size_local):
-            start_rpyc_service_process(
-                ModelTpService, model_port_args[0].model_tp_ports[i]
-            )
-        if server_args.node_rank != 0:
-            logger.info(
-                f"[node_rank={server_args.node_rank}]: Listen for connections..."
-            )
+            tp_size_local = server_args.tp_size // server_args.nnodes
+            gpu_ids = [i for _ in range(server_args.nnodes) for i in range(tp_size_local)]
+            tp_rank_range = list(range(server_args.node_rank * tp_size_local,
+                                  (server_args.node_rank + 1) * tp_size_local))
+            procs = launch_tp_servers(gpu_ids, tp_rank_range, server_args,
+                                      port_args.model_port_args[0], model_overide_args)
             while True:
                 pass
 

sglang/srt/server_args.py

@@ -29,7 +29,7 @@ class ServerArgs:
     max_prefill_tokens: Optional[int] = None
     max_running_requests: Optional[int] = None
     schedule_heuristic: str = "lpm"
-    schedule_conservativeness: float = 1.0
+    schedule_conservativeness: float = 0.8
 
     # Other runtime options
     tp_size: int = 1
@@ -53,6 +53,7 @@ class ServerArgs:
     disable_flashinfer: bool = False
     disable_radix_cache: bool = False
     disable_regex_jump_forward: bool = False
+    disable_cuda_graph: bool = False
     disable_disk_cache: bool = False
     attention_reduce_in_fp32: bool = False
     enable_p2p_check: bool = False
@@ -66,14 +67,16 @@ class ServerArgs:
         if self.tokenizer_path is None:
             self.tokenizer_path = self.model_path
         if self.mem_fraction_static is None:
-            if self.tp_size >= 8:
-                self.mem_fraction_static = 0.80
+            if self.tp_size >= 16:
+                self.mem_fraction_static = 0.74
+            elif self.tp_size >= 8:
+                self.mem_fraction_static = 0.78
             elif self.tp_size >= 4:
                 self.mem_fraction_static = 0.82
             elif self.tp_size >= 2:
                 self.mem_fraction_static = 0.85
             else:
-                self.mem_fraction_static = 0.90
+                self.mem_fraction_static = 0.88
         if isinstance(self.additional_ports, int):
             self.additional_ports = [self.additional_ports]
         elif self.additional_ports is None:
@@ -294,6 +297,11 @@ class ServerArgs:
             action="store_true",
             help="Disable regex jump-forward",
         )
+        parser.add_argument(
+            "--disable-cuda-graph",
+            action="store_true",
+            help="Disable cuda graph.",
+        )
         parser.add_argument(
             "--disable-disk-cache",
             action="store_true",

sglang/srt/utils.py

@@ -474,9 +474,9 @@ def monkey_patch_vllm_dummy_weight_loader():
         DummyModelLoader,
         LoRAConfig,
         ModelConfig,
+        MultiModalConfig,
         ParallelConfig,
         SchedulerConfig,
-        MultiModalConfig,
         _initialize_model,
         initialize_dummy_weights,
         nn,

{sglang-0.1.19.dist-info → sglang-0.1.21.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: sglang
-Version: 0.1.19
+Version: 0.1.21
 Summary: A structured generation langauge for LLMs.
 License: Apache License
                                    Version 2.0, January 2004
@@ -623,6 +623,14 @@ python -m sglang.launch_server --model-path meta-llama/Llama-2-7b-chat-hf --port
 python -m sglang.launch_server --model-path meta-llama/Llama-2-7b-chat-hf --port 30000 --mem-fraction-static 0.7
 ```
 - See [hyperparameter_tuning.md](docs/hyperparameter_tuning.md) on tuning hyperparameters for better performance.
+- Add `--nnodes 2` to run tensor parallelism on multiple nodes. If you have two nodes with two GPUs on each node and want to run TP=4, let `sgl-dev-1` be the hostname of the first node and `50000` be an available port.
+```
+# Node 0
+python -m sglang.launch_server --model-path meta-llama/Llama-2-7b-chat-hf --tp 4 --nccl-init sgl-dev-1:50000 --nnodes 2 --node-rank 0
+
+# Node 1
+python -m sglang.launch_server --model-path meta-llama/Llama-2-7b-chat-hf --tp 4 --nccl-init sgl-dev-1:50000 --nnodes 2 --node-rank 1
+```
 
 ### Supported Models
 - Llama