sglang 0.3.1.post3__py3-none-any.whl → 0.3.3__py3-none-any.whl

sglang/srt/server_args.py

@@ -19,9 +19,10 @@ import argparse
 import dataclasses
 import logging
 import random
-from typing import List, Optional, Union
+import tempfile
+from typing import List, Optional
 
-from sglang.srt.utils import is_hip
+from sglang.srt.utils import is_flashinfer_available, is_ipv6, is_port_available
 
 logger = logging.getLogger(__name__)
 
@@ -46,7 +47,6 @@ class ServerArgs:
     # Port
     host: str = "127.0.0.1"
     port: int = 30000
-    additional_ports: Optional[Union[List[int], int]] = None
 
     # Memory and scheduling
     mem_fraction_static: Optional[float] = None
@@ -78,9 +78,9 @@ class ServerArgs:
     load_balance_method: str = "round_robin"
 
     # Distributed args
-    nccl_init_addr: Optional[str] = None
+    dist_init_addr: Optional[str] = None
     nnodes: int = 1
-    node_rank: Optional[int] = None
+    node_rank: int = 0
 
     # Model override args in JSON
     json_model_override_args: str = "{}"
@@ -134,11 +134,6 @@ class ServerArgs:
             else:
                 self.mem_fraction_static = 0.88
 
-        if isinstance(self.additional_ports, int):
-            self.additional_ports = [self.additional_ports]
-        elif self.additional_ports is None:
-            self.additional_ports = []
-
         if self.random_seed is None:
             self.random_seed = random.randint(0, 1 << 30)
 
@@ -156,8 +151,7 @@ class ServerArgs:
             )
             self.sampling_backend = "pytorch"
 
-        # ROCm: flashinfer available later
-        if is_hip():
+        if not is_flashinfer_available():
             self.attention_backend = "triton"
             self.sampling_backend = "pytorch"
 
@@ -199,13 +193,6 @@ class ServerArgs:
         parser.add_argument(
             "--port", type=int, default=ServerArgs.port, help="The port of the server."
         )
-        parser.add_argument(
-            "--additional-ports",
-            type=int,
-            nargs="*",
-            default=[],
-            help="The additional ports specified for the server.",
-        )
         parser.add_argument(
             "--tokenizer-mode",
             type=str,
@@ -279,7 +266,6 @@ class ServerArgs:
                 "marlin",
                 "gptq_marlin",
                 "awq_marlin",
-                "squeezellm",
                 "bitsandbytes",
             ],
             help="The quantization method.",
@@ -426,14 +412,17 @@ class ServerArgs:
 
         # Multi-node distributed serving args
         parser.add_argument(
-            "--nccl-init-addr",
+            "--dist-init-addr",
+            "--nccl-init-addr",  # For backward compatbility. This will be removed in the future.
             type=str,
-            help="The nccl init address of multi-node server.",
+            help="The host address for initializing distributed backend (e.g., `192.168.0.2:25000`).",
         )
         parser.add_argument(
             "--nnodes", type=int, default=ServerArgs.nnodes, help="The number of nodes."
         )
-        parser.add_argument("--node-rank", type=int, help="The node rank.")
+        parser.add_argument(
+            "--node-rank", type=int, default=ServerArgs.node_rank, help="The node rank."
+        )
 
         # Model override args
         parser.add_argument(
@@ -567,7 +556,10 @@ class ServerArgs:
         return cls(**{attr: getattr(args, attr) for attr in attrs})
 
     def url(self):
-        return f"http://{self.host}:{self.port}"
+        if is_ipv6(self.host):
+            return f"http://[{self.host}]:{self.port}"
+        else:
+            return f"http://{self.host}:{self.port}"
 
     def check_server_args(self):
         assert (
@@ -583,6 +575,21 @@ class ServerArgs:
             and (self.lora_paths is None or self.disable_radix_cache)
         ), "compatibility of lora and cuda graph and radix attention is in progress"
 
+        assert self.dp_size == 1, (
+            "The support for data parallelism is temporarily disabled during refactor. "
+            "Please use sglang<=0.3.2 or wait for later updates."
+        )
+
+        if isinstance(self.lora_paths, list):
+            lora_paths = self.lora_paths
+            self.lora_paths = {}
+            for lora_path in lora_paths:
+                if "=" in lora_path:
+                    name, path = lora_path.split("=", 1)
+                    self.lora_paths[name] = path
+                else:
+                    self.lora_paths[lora_path] = lora_path
+
 
 def prepare_server_args(argv: List[str]) -> ServerArgs:
     """
@@ -604,11 +611,31 @@ def prepare_server_args(argv: List[str]) -> ServerArgs:
 
 @dataclasses.dataclass
 class PortArgs:
-    tokenizer_port: int
-    controller_port: int
-    detokenizer_port: int
+    # The ipc filename for tokenizer to receive inputs from detokenizer (zmq)
+    tokenizer_ipc_name: str
+    # The ipc filename for scheduler (rank 0) to receive inputs from tokenizer (zmq)
+    scheduler_input_ipc_name: str
+    # The ipc filename for detokenizer to receive inputs from scheduler (zmq)
+    detokenizer_ipc_name: str
+
+    # The port for nccl initialization for multiple TP groups (torch.dist)
     nccl_ports: List[int]
 
+    @classmethod
+    def init_new(self, server_args):
+        port = server_args.port + 1
+        while True:
+            if is_port_available(port):
+                break
+            port += 1
+
+        return PortArgs(
+            tokenizer_ipc_name=tempfile.NamedTemporaryFile(delete=False).name,
+            scheduler_input_ipc_name=tempfile.NamedTemporaryFile(delete=False).name,
+            detokenizer_ipc_name=tempfile.NamedTemporaryFile(delete=False).name,
+            nccl_ports=[port],
+        )
+
 
 class LoRAPathAction(argparse.Action):
     def __call__(self, parser, namespace, values, option_string=None):

sglang/srt/utils.py

@@ -16,14 +16,16 @@ limitations under the License.
 """Common utilities."""
 
 import base64
-import fcntl
+import ipaddress
+import json
 import logging
 import os
+import pickle
 import random
 import resource
 import socket
-import struct
 import time
+import warnings
 from importlib.metadata import PackageNotFoundError, version
 from io import BytesIO
 from typing import Any, Dict, List, Optional, Union
@@ -36,7 +38,7 @@ import torch.distributed as dist
 from fastapi.responses import JSONResponse
 from packaging import version as pkg_version
 from torch import nn
-from torch.nn.parameter import Parameter
+from torch.profiler import ProfilerActivity, profile, record_function
 from triton.runtime.cache import (
     FileCacheManager,
     default_cache_dir,
@@ -51,11 +53,27 @@ show_time_cost = False
 time_infos = {}
 
 
-# torch flag AMD GPU
 def is_hip() -> bool:
+    """Return whether it is HIP on the AMD ROCm platform."""
     return torch.version.hip is not None
 
 
+def is_flashinfer_available():
+    """
+    Check whether flashinfer is available.
+    As of Oct. 6, 2024, it is only available on NVIDIA GPUs.
+    """
+    return torch.cuda.is_available() and not is_hip()
+
+
+def is_ipv6(address):
+    try:
+        ipaddress.IPv6Address(address)
+        return True
+    except ipaddress.AddressValueError:
+        return False
+
+
 def enable_show_time_cost():
     global show_time_cost
     show_time_cost = True
@@ -170,35 +188,6 @@ def is_port_available(port):
             return False
 
 
-def allocate_init_ports(
-    port: Optional[int] = None,
-    additional_ports: Optional[List[int]] = None,
-    dp_size: int = 1,
-):
-    """Allocate ports for all connections."""
-    if additional_ports:
-        ret_ports = [port] + additional_ports
-    else:
-        ret_ports = [port]
-
-    ret_ports = list(set(x for x in ret_ports if is_port_available(x)))
-    cur_port = ret_ports[-1] + 1 if len(ret_ports) > 0 else 10000
-
-    # HTTP + Tokenizer + Controller + Detokenizer + dp_size * 1 (nccl)
-    num_ports_needed = 4 + dp_size
-    while len(ret_ports) < num_ports_needed:
-        if cur_port not in ret_ports and is_port_available(cur_port):
-            ret_ports.append(cur_port)
-        cur_port += 1
-
-    if port is not None and ret_ports[0] != port:
-        logger.warning(
-            f"WARNING: Port {port} is not available. Use port {ret_ports[0]} instead."
-        )
-
-    return ret_ports[0], ret_ports[1:num_ports_needed]
-
-
 def is_multimodal_model(model_architectures):
     if (
         "LlavaLlamaForCausalLM" in model_architectures
@@ -219,6 +208,8 @@ def is_generation_model(model_architectures, is_embedding: bool = False):
     if (
         "LlamaEmbeddingModel" in model_architectures
         or "MistralModel" in model_architectures
+        or "LlamaForSequenceClassification" in model_architectures
+        or "LlamaForSequenceClassificationWithNormal_Weights" in model_architectures
     ):
         return False
     else:
@@ -345,6 +336,10 @@ def suppress_other_loggers():
     logging.getLogger("vllm.selector").setLevel(logging.WARN)
     logging.getLogger("vllm.utils").setLevel(logging.ERROR)
 
+    warnings.filterwarnings(
+        "ignore", category=UserWarning, message="The given NumPy array is not writable"
+    )
+
 
 def assert_pkg_version(pkg: str, min_version: str, message: str):
     try:
@@ -537,89 +532,6 @@ class CustomCacheManager(FileCacheManager):
                 raise RuntimeError("Could not create or locate cache dir")
 
 
-def get_ip_address(ifname):
-    """
-    Get the IP address of a network interface.
-
-    :param ifname: Name of the network interface (e.g., 'eth0')
-    :return: IP address of the network interface
-    """
-    s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
-    ip_address = fcntl.ioctl(
-        s.fileno(),
-        0x8915,  # SIOCGIFADDR
-        struct.pack("256s", bytes(ifname[:15], "utf-8")),
-    )[20:24]
-    return socket.inet_ntoa(ip_address)
-
-
-def send_addrs_to_rank_0(model_port_args, server_args):
-    assert server_args.node_rank != 0 and server_args.dp_size == 1
-
-    ifname = os.environ.get(
-        "SGLANG_SOCKET_IFNAME", os.environ.get("NCCL_SOCKET_IFNAME", "eth0")
-    )
-    ip_addr = get_ip_address(ifname)
-
-    num_tp_ports = server_args.tp_size // server_args.nnodes
-    model_port_args.model_tp_ips[:num_tp_ports] = [ip_addr] * num_tp_ports
-    ip_addr = [int(x) for x in ip_addr.split(".")]
-    addrs_tensor = torch.tensor(
-        ip_addr + model_port_args.model_tp_ports, dtype=torch.int
-    )
-
-    init_method = f"tcp://{server_args.nccl_init_addr}"
-    dist.init_process_group(
-        backend="gloo",
-        init_method=init_method,
-        rank=server_args.node_rank,
-        world_size=server_args.nnodes,
-    )
-    dist.send(addrs_tensor, dst=0)
-    print(
-        f"Node {server_args.node_rank} sent: ip_address {ip_addr} and ports {model_port_args.model_tp_ports}"
-    )
-
-    dist.barrier()
-    dist.destroy_process_group()
-
-
-def receive_addrs(model_port_args, server_args):
-    assert server_args.node_rank == 0 and server_args.dp_size == 1
-
-    ifname = os.environ.get(
-        "SGLANG_SOCKET_IFNAME", os.environ.get("NCCL_SOCKET_IFNAME", "eth0")
-    )
-    ip_addr = get_ip_address(ifname)
-
-    num_tp_ports = server_args.tp_size // server_args.nnodes
-    model_port_args.model_tp_ips[:num_tp_ports] = [ip_addr] * num_tp_ports
-
-    init_method = f"tcp://{server_args.nccl_init_addr}"
-    dist.init_process_group(
-        backend="gloo",
-        init_method=init_method,
-        rank=server_args.node_rank,
-        world_size=server_args.nnodes,
-    )
-
-    for src_rank in range(1, server_args.nnodes):
-        tensor = torch.zeros(4 + num_tp_ports, dtype=torch.int)
-        dist.recv(tensor, src=src_rank)
-        ip = ".".join([str(x) for x in tensor[:4].tolist()])
-        ports = tensor[4:].tolist()
-        model_port_args.model_tp_ips[
-            num_tp_ports * src_rank : num_tp_ports * (src_rank + 1)
-        ] = [ip] * num_tp_ports
-        model_port_args.model_tp_ports[
-            num_tp_ports * src_rank : num_tp_ports * (src_rank + 1)
-        ] = ports
-        print(f"Node 0 received from rank {src_rank}: {tensor.tolist()}")
-
-    dist.barrier()
-    dist.destroy_process_group()
-
-
 def set_ulimit(target_soft_limit=65535):
     resource_type = resource.RLIMIT_NOFILE
     current_soft, current_hard = resource.getrlimit(resource_type)
@@ -643,24 +555,16 @@ def add_api_key_middleware(app, api_key: str):
         return await call_next(request)
 
 
-def prepare_model(model_path: str):
+def prepare_model_and_tokenizer(model_path: str, tokenizer_path: str):
     if "SGLANG_USE_MODELSCOPE" in os.environ:
         if not os.path.exists(model_path):
             from modelscope import snapshot_download
 
-            return snapshot_download(model_path)
-    return model_path
-
-
-def prepare_tokenizer(tokenizer_path: str):
-    if "SGLANG_USE_MODELSCOPE" in os.environ:
-        if not os.path.exists(tokenizer_path):
-            from modelscope import snapshot_download
-
-            return snapshot_download(
+            model_path = snapshot_download(model_path)
+            tokenizer_path = snapshot_download(
                 tokenizer_path, ignore_patterns=["*.bin", "*.safetensors"]
             )
-    return tokenizer_path
+    return model_path, tokenizer_path
 
 
 def configure_logger(server_args, prefix: str = ""):
@@ -702,3 +606,72 @@ def set_weight_attrs(
     for key, value in weight_attrs.items():
         assert not hasattr(weight, key), f"Overwriting existing tensor attribute: {key}"
         setattr(weight, key, value)
+
+
+def broadcast_pyobj(
+    data: List[Any],
+    rank: int,
+    dist_group: Optional[torch.distributed.ProcessGroup] = None,
+):
+    """Broadcast inputs from rank=0 to all other ranks with torch.dist backend."""
+
+    if rank == 0:
+        if len(data) == 0:
+            tensor_size = torch.tensor([0], dtype=torch.long)
+            dist.broadcast(tensor_size, src=0, group=dist_group)
+        else:
+            serialized_data = pickle.dumps(data)
+            size = len(serialized_data)
+            tensor_data = torch.ByteTensor(
+                np.frombuffer(serialized_data, dtype=np.uint8)
+            )
+            tensor_size = torch.tensor([size], dtype=torch.long)
+
+            dist.broadcast(tensor_size, src=0, group=dist_group)
+            dist.broadcast(tensor_data, src=0, group=dist_group)
+        return data
+    else:
+        tensor_size = torch.tensor([0], dtype=torch.long)
+        dist.broadcast(tensor_size, src=0, group=dist_group)
+        size = tensor_size.item()
+
+        if size == 0:
+            return []
+
+        tensor_data = torch.empty(size, dtype=torch.uint8)
+        dist.broadcast(tensor_data, src=0, group=dist_group)
+
+        serialized_data = bytes(tensor_data.cpu().numpy())
+        data = pickle.loads(serialized_data)
+        return data
+
+
+step_counter = 0
+
+
+def pytorch_profile(name, func, *args, data_size=-1):
+    """
+    Args:
+        name (string): the name of recorded function.
+        func: the function to be profiled.
+        args: the arguments of the profiled function.
+        data_size (int): some measurement of the computation complexity.
+            Usually, it could be the batch size.
+    """
+    global step_counter
+    os.makedirs("trace", exist_ok=True)
+    with profile(
+        activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
+        # schedule=torch.profiler.schedule(wait=1, warmup=1, active=3, repeat=2),
+        # on_trace_ready=tensorboard_trace_handler('./log_dir'),
+        record_shapes=True,
+        profile_memory=True,
+        with_stack=True,
+    ) as prof:
+        with record_function(name):
+            with open(f"trace/size_{step_counter}.json", "w") as f:
+                json.dump({"size": data_size}, f)
+            result = func(*args)
+    prof.export_chrome_trace(f"trace/{name}_{step_counter}.json")
+    step_counter += 1
+    return result

sglang/test/runners.py

@@ -21,19 +21,19 @@ from typing import List, Union
 
 import torch
 import torch.nn.functional as F
-from peft import PeftModel
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import AutoModelForCausalLM
 
+from sglang.srt.hf_transformers_utils import get_tokenizer
 from sglang.srt.server import Runtime
 from sglang.test.test_utils import DEFAULT_PORT_FOR_SRT_TEST_RUNNER
 
 DEFAULT_PROMPTS = [
-    # the output of gemma-2-2b from SRT is unstable on the commented prompt
-    # "The capital of France is",
     "Apple is red. Banana is Yellow. " * 800 + "Apple is",
     "The capital of the United Kingdom is",
     "Today is a sunny day and I like",
     "AI is a field of computer science focused on",
+    # the output of gemma-2-2b from SRT is unstable on the commented prompt
+    # "The capital of France is",
 ]
 
 dirpath = os.path.dirname(__file__)
@@ -65,17 +65,18 @@ class ModelOutput:
     top_input_logprobs: List[torch.Tensor] = None
     top_output_logprobs: List[torch.Tensor] = None
     embed_logits: List[torch.Tensor] = None
+    scores: List[float] = None
 
 
 class HFRunner:
     def __init__(
         self,
-        model_path,
-        torch_dtype,
-        is_generation,
-        output_str_only=False,
+        model_path: str,
+        torch_dtype: torch.dtype,
+        model_type: str = "generation",
+        output_str_only: bool = False,
     ):
-        self.is_generation = is_generation
+        self.model_type = model_type
         self.output_str_only = output_str_only
 
         self.in_queue = mp.Queue()
@@ -92,26 +93,41 @@ class HFRunner:
         )
         self.model_proc.start()
 
+    def needs_trust_remote_code(self, model_path):
+        models_needs_trust_remote = [
+            "LxzGordon/URM-LLaMa-3.1-8B",
+        ]
+        if model_path in models_needs_trust_remote:
+            return True
+        return False
+
     def start_model_process(self, in_queue, out_queue, model_path, torch_dtype):
-        self.tokenizer = AutoTokenizer.from_pretrained(
-            model_path,
-            torch_dtype=torch_dtype,
-        )
+        self.tokenizer = get_tokenizer(model_path, torch_dtype=torch.dtype)
 
-        if self.is_generation:
+        if self.model_type == "generation":
             self.base_model = AutoModelForCausalLM.from_pretrained(
                 model_path,
                 torch_dtype=torch_dtype,
                 trust_remote_code=False,
                 low_cpu_mem_usage=True,
             ).cuda()
-        else:
+        elif self.model_type == "embedding":
             from sentence_transformers import SentenceTransformer
 
             self.model = SentenceTransformer(
                 model_path,
                 model_kwargs={"torch_dtype": torch_dtype},
-            )
+            ).cuda()
+        elif self.model_type == "reward":
+            from transformers import AutoModelForSequenceClassification
+
+            self.model = AutoModelForSequenceClassification.from_pretrained(
+                model_path,
+                torch_dtype=torch_dtype,
+                trust_remote_code=self.needs_trust_remote_code(model_path),
+            ).cuda()
+        else:
+            raise Exception(f"Unrecognized model type {self.model_type}")
 
         while True:
             prompts, max_new_tokens, lora_paths = in_queue.get()
@@ -119,7 +135,7 @@ class HFRunner:
                 assert len(prompts) == len(lora_paths)
 
             if prompts is not None:
-                if self.is_generation:
+                if self.model_type == "generation":
                     output_strs = []
                     top_input_logprobs = []
                     top_output_logprobs = []
@@ -132,6 +148,8 @@ class HFRunner:
                             input_ids = torch.tensor([p], device="cuda")
 
                         if lora_paths is not None and lora_paths[i] is not None:
+                            from peft import PeftModel
+
                             self.model = PeftModel.from_pretrained(
                                 self.base_model,
                                 lora_paths[i],
@@ -181,11 +199,27 @@ class HFRunner:
                         )
                     )
 
-                else:
+                elif self.model_type == "embedding":
                     assert not self.output_str_only
                     logits = self.model.encode(prompts).tolist()
                     out_queue.put(ModelOutput(embed_logits=logits))
 
+                elif self.model_type == "reward":
+                    scores = []
+                    for conv in prompts:
+                        conv_formatted = self.tokenizer.apply_chat_template(
+                            conv, tokenize=False
+                        )
+                        conv_tokenized = self.tokenizer(
+                            conv_formatted, return_tensors="pt"
+                        ).to("cuda")
+                        scores.append(
+                            float(self.model(**conv_tokenized).logits[0][0].item())
+                        )
+                    out_queue.put(ModelOutput(scores=scores))
+                else:
+                    raise Exception(f"Unrecognized model type {self.model_type}")
+
     def forward(
         self,
         prompts: Union[List[str], List[torch.Tensor]] = DEFAULT_PROMPTS,
@@ -210,23 +244,24 @@ class HFRunner:
 class SRTRunner:
     def __init__(
         self,
-        model_path,
-        torch_dtype,
-        is_generation,
-        tp_size=1,
-        port=DEFAULT_PORT_FOR_SRT_TEST_RUNNER,
-        lora_paths=None,
-        max_loras_per_batch=4,
-        disable_cuda_graph=False,
-        disable_radix_cache=False,
+        model_path: str,
+        torch_dtype: torch.dtype,
+        model_type: str,
+        tp_size: int = 1,
+        port: int = DEFAULT_PORT_FOR_SRT_TEST_RUNNER,
+        lora_paths: List[str] = None,
+        max_loras_per_batch: int = 4,
+        disable_cuda_graph: bool = False,
+        disable_radix_cache: bool = False,
     ):
-        self.is_generation = is_generation
+        self.model_type = model_type
+        self.is_generation = model_type == "generation"
         self.runtime = Runtime(
             model_path=model_path,
             tp_size=tp_size,
             dtype=get_dtype_str(torch_dtype),
             port=port,
-            mem_fraction_static=0.69,
+            mem_fraction_static=0.65,
             trust_remote_code=False,
             is_embedding=not self.is_generation,
             lora_paths=lora_paths,
@@ -287,8 +322,12 @@ class SRTRunner:
         else:
             response = self.runtime.encode(prompts)
             response = json.loads(response)
-            logits = [x["embedding"] for x in response]
-            return ModelOutput(embed_logits=logits)
+            if self.model_type == "embedding":
+                logits = [x["embedding"] for x in response]
+                return ModelOutput(embed_logits=logits)
+            else:
+                scores = [x["embedding"][0] for x in response]
+                return ModelOutput(scores=scores)
 
     def batch_forward(
         self,
@@ -318,8 +357,12 @@ class SRTRunner:
         else:
             response = self.runtime.encode(prompts)
             response = json.loads(response)
-            logits = [x["embedding"] for x in response]
-            return ModelOutput(embed_logits=logits)
+            if self.model_type == "embedding":
+                logits = [x["embedding"] for x in response]
+                return ModelOutput(embed_logits=logits)
+            else:
+                scores = [x["embedding"][0] for x in response]
+                return ModelOutput(scores=logits)
 
     def __enter__(self):
         return self