sglang 0.4.9.post5__py3-none-any.whl → 0.4.10__py3-none-any.whl

sglang/srt/managers/scheduler_update_weights_mixin.py

@@ -0,0 +1,142 @@
+import logging
+from typing import Tuple
+
+import torch
+
+from sglang.srt.constants import GPU_MEMORY_TYPE_KV_CACHE, GPU_MEMORY_TYPE_WEIGHTS
+from sglang.srt.managers.io_struct import (
+    GetWeightsByNameReqInput,
+    GetWeightsByNameReqOutput,
+    InitWeightsUpdateGroupReqInput,
+    InitWeightsUpdateGroupReqOutput,
+    ReleaseMemoryOccupationReqInput,
+    ReleaseMemoryOccupationReqOutput,
+    ResumeMemoryOccupationReqInput,
+    ResumeMemoryOccupationReqOutput,
+    UpdateWeightFromDiskReqInput,
+    UpdateWeightFromDiskReqOutput,
+    UpdateWeightsFromDistributedReqInput,
+    UpdateWeightsFromDistributedReqOutput,
+    UpdateWeightsFromTensorReqInput,
+    UpdateWeightsFromTensorReqOutput,
+)
+
+logger = logging.getLogger(__name__)
+
+
+class SchedulerUpdateWeightsMixin:
+
+    def update_weights_from_disk(self, recv_req: UpdateWeightFromDiskReqInput):
+        """In-place update of the weights from disk."""
+        success, message = self.tp_worker.update_weights_from_disk(recv_req)
+        if success:
+            flush_cache_success = self.flush_cache()
+            assert flush_cache_success, "Cache flush failed after updating weights"
+        else:
+            logger.error(message)
+        return UpdateWeightFromDiskReqOutput(success, message, 0)
+
+    def init_weights_update_group(self, recv_req: InitWeightsUpdateGroupReqInput):
+        """Initialize the online model parameter update group."""
+        success, message = self.tp_worker.init_weights_update_group(recv_req)
+        return InitWeightsUpdateGroupReqOutput(success, message)
+
+    def update_weights_from_distributed(
+        self,
+        recv_req: UpdateWeightsFromDistributedReqInput,
+    ) -> Tuple[bool, str]:
+        """Update the online model parameter."""
+        success, message = self.tp_worker.update_weights_from_distributed(recv_req)
+        if success:
+            if recv_req.flush_cache:
+                flush_cache_success = self.flush_cache()
+                assert flush_cache_success, "Cache flush failed after updating weights"
+        else:
+            logger.error(message)
+        return UpdateWeightsFromDistributedReqOutput(success, message)
+
+    def update_weights_from_tensor(self, recv_req: UpdateWeightsFromTensorReqInput):
+        """Update the online model parameter from tensors."""
+        success, message = self.tp_worker.update_weights_from_tensor(recv_req)
+        # TODO extract common code b/t update_weights_from_distributed and update_weights_from_tensor later
+        if success:
+            if recv_req.flush_cache:
+                flush_cache_success = self.flush_cache()
+                assert flush_cache_success, "Cache flush failed after updating weights"
+        else:
+            logger.error(message)
+        torch.distributed.barrier(group=self.tp_cpu_group)
+        return UpdateWeightsFromTensorReqOutput(success, message)
+
+    def get_weights_by_name(self, recv_req: GetWeightsByNameReqInput):
+        parameter = self.tp_worker.get_weights_by_name(recv_req)
+        return GetWeightsByNameReqOutput(parameter)
+
+    def release_memory_occupation(self, recv_req: ReleaseMemoryOccupationReqInput):
+        tags = recv_req.tags
+
+        if tags is None or len(tags) == 0:
+            tags = [GPU_MEMORY_TYPE_WEIGHTS, GPU_MEMORY_TYPE_KV_CACHE]
+
+        if GPU_MEMORY_TYPE_KV_CACHE in tags:
+            self.memory_saver_adapter.pause(GPU_MEMORY_TYPE_KV_CACHE)
+            self.flush_cache()
+
+        if GPU_MEMORY_TYPE_WEIGHTS in tags:
+            self.stashed_model_static_state = _export_static_state(
+                self.tp_worker.worker.model_runner.model
+            )
+            torch.distributed.barrier(self.tp_cpu_group)
+            self.memory_saver_adapter.pause(GPU_MEMORY_TYPE_WEIGHTS)
+
+        return ReleaseMemoryOccupationReqOutput()
+
+    def resume_memory_occupation(self, recv_req: ResumeMemoryOccupationReqInput):
+        tags = recv_req.tags
+
+        if tags is None or len(tags) == 0:
+            tags = [GPU_MEMORY_TYPE_WEIGHTS, GPU_MEMORY_TYPE_KV_CACHE]
+
+        if GPU_MEMORY_TYPE_WEIGHTS in tags:
+            self.memory_saver_adapter.resume(GPU_MEMORY_TYPE_WEIGHTS)
+            torch.distributed.barrier(self.tp_cpu_group)
+            _import_static_state(
+                self.tp_worker.worker.model_runner.model,
+                self.stashed_model_static_state,
+            )
+            del self.stashed_model_static_state
+
+        if GPU_MEMORY_TYPE_KV_CACHE in tags:
+            self.memory_saver_adapter.resume(GPU_MEMORY_TYPE_KV_CACHE)
+
+        return ResumeMemoryOccupationReqOutput()
+
+    def save_remote_model(self, params):
+        url = params["url"]
+
+        worker = self.tp_worker.worker
+
+        worker.model_runner.save_remote_model(url)
+
+    def save_sharded_model(self, params):
+        worker = self.tp_worker.worker
+
+        worker.model_runner.save_sharded_model(
+            path=params["path"],
+            pattern=params["pattern"],
+            max_size=params["max_size"],
+        )
+
+
+def _export_static_state(model):
+    return dict(
+        buffers=[
+            (name, buffer.detach().clone()) for name, buffer in model.named_buffers()
+        ]
+    )
+
+
+def _import_static_state(model, static_params):
+    self_named_buffers = dict(model.named_buffers())
+    for name, tensor in static_params["buffers"]:
+        self_named_buffers[name][...] = tensor

sglang/srt/managers/template_manager.py

@@ -53,7 +53,7 @@ class TemplateManager:
     def __init__(self):
         self._chat_template_name: Optional[str] = None
         self._completion_template_name: Optional[str] = None
-        self._jinja_template_content_format: Optional[str] = None
+        self._jinja_template_content_format: Optional[str] = "openai"
 
     @property
     def chat_template_name(self) -> Optional[str]:
@@ -71,31 +71,60 @@ class TemplateManager:
         return self._jinja_template_content_format
 
     def load_chat_template(
-        self, tokenizer_manager, chat_template_arg: str, model_path: str
+        self, tokenizer_manager, chat_template_arg: Optional[str], model_path: str
     ) -> None:
         """
         Load a chat template from various sources.
 
         Args:
             tokenizer_manager: The tokenizer manager instance
-            chat_template_arg: Template name or file path
+            chat_template_arg: Template name, file path, or None to auto-detect
             model_path: Path to the model
         """
-        logger.info(f"Loading chat template: {chat_template_arg}")
+        if chat_template_arg:
+            self._load_explicit_chat_template(tokenizer_manager, chat_template_arg)
+        else:
+            # Try HuggingFace template first
+            hf_template = self._resolve_hf_chat_template(tokenizer_manager)
+            if hf_template:
+                self._jinja_template_content_format = (
+                    detect_jinja_template_content_format(hf_template)
+                )
+                logger.info(
+                    f"Using default HuggingFace chat template with detected content format: {self._jinja_template_content_format}"
+                )
+                return
 
-        if not chat_template_exists(chat_template_arg):
-            if not os.path.exists(chat_template_arg):
-                raise RuntimeError(
-                    f"Chat template {chat_template_arg} is not a built-in template name "
-                    "or a valid chat template file path."
+            # Fallback to SGLang template guessing
+            self.guess_chat_template_from_model_path(model_path)
+
+            # Set default format if no template was found
+            if self._chat_template_name is None:
+                self._jinja_template_content_format = "string"
+                logger.info(
+                    "No chat template found, defaulting to 'string' content format"
                 )
 
-            if chat_template_arg.endswith(".jinja"):
-                self._load_jinja_template(tokenizer_manager, chat_template_arg)
-            else:
-                self._load_json_chat_template(chat_template_arg)
-        else:
+    def _load_explicit_chat_template(
+        self, tokenizer_manager, chat_template_arg: str
+    ) -> None:
+        """Load explicitly specified chat template."""
+        logger.info(f"Loading chat template from argument: {chat_template_arg}")
+
+        if chat_template_exists(chat_template_arg):
             self._chat_template_name = chat_template_arg
+            return
+
+        if not os.path.exists(chat_template_arg):
+            raise RuntimeError(
+                f"Chat template {chat_template_arg} is not a built-in template name "
+                "or a valid chat template file path."
+            )
+
+        if chat_template_arg.endswith(".jinja"):
+            self._load_jinja_template(tokenizer_manager, chat_template_arg)
+        else:
+            self._load_json_chat_template(chat_template_arg)
 
     def guess_chat_template_from_model_path(self, model_path: str) -> None:
         """
@@ -146,10 +175,7 @@ class TemplateManager:
             completion_template: Optional completion template name/path
         """
         # Load chat template
-        if chat_template:
-            self.load_chat_template(tokenizer_manager, chat_template, model_path)
-        else:
-            self.guess_chat_template_from_model_path(model_path)
+        self.load_chat_template(tokenizer_manager, chat_template, model_path)
 
         # Load completion template
         if completion_template:
@@ -166,7 +192,7 @@ class TemplateManager:
             chat_template
         )
         logger.info(
-            f"Detected chat template content format: {self._jinja_template_content_format}"
+            f"Detected user specified Jinja chat template with content format: {self._jinja_template_content_format}"
         )
 
     def _load_json_chat_template(self, template_path: str) -> None:
@@ -224,3 +250,20 @@ class TemplateManager:
                 override=True,
             )
         self._completion_template_name = template["name"]
+
+    def _resolve_hf_chat_template(self, tokenizer_manager) -> Optional[str]:
+        """
+        Resolve HuggingFace chat template.
+
+        Returns the chat template string if found, None otherwise.
+        """
+        tokenizer = tokenizer_manager.tokenizer
+
+        # Try to get AutoTokenizer chat template
+        try:
+            return tokenizer.get_chat_template()
+        except Exception as e:
+            logger.debug(f"Error getting chat template via get_chat_template(): {e}")
+
+        logger.debug("No HuggingFace chat template found")
+        return None