[DSV4] Cherry pick missing commits from deepseek_v4 branch and enhance tests (#24793)

Co-authored-by: Xinyuan Tong <115166877+JustinTong0323@users.noreply.github.com>
Co-authored-by: yueming-yuan <yym022502@gmail.com>
This commit is contained in:
Baizhou Zhang
2026-05-09 04:15:37 -07:00
committed by GitHub
co-authored by Xinyuan Tong yueming-yuan
parent 4b23f6bdc5
commit ef5e9f8aba
15 changed files with 481 additions and 87 deletions
@@ -633,13 +633,16 @@ class ChatCompletionRequest(BaseModel):
return_hidden_states: bool = False
return_routed_experts: bool = False
return_cached_tokens_details: bool = False
reasoning_effort: Optional[Literal["none", "low", "medium", "high"]] = Field(
reasoning_effort: Optional[Literal["none", "low", "medium", "high", "max"]] = Field(
default=None,
description="Constrains effort on reasoning for reasoning models. "
"'none' disables reasoning entirely, 'low' is the least effort, 'high' is the most effort. "
"Reducing reasoning effort can result in faster responses and fewer tokens used on reasoning "
"in a response. 'none' defaults thinking and enable_thinking to false in "
"chat_template_kwargs (unless explicitly overridden). Not supported in the harmony path.",
"chat_template_kwargs (unless explicitly overridden). Not supported in the harmony path."
"'max' is an sglang extension to the OpenAI schema for "
"models that expose a maximum-effort tier above 'high'; models that don't "
"support it treat it the same as 'high'.",
)
task: Optional[
Literal["action", "query", "authority", "domain", "title", "read_url"]
@@ -81,8 +81,13 @@ class DeepSeekV32Detector(BaseFormatDetector):
self.function_calls_regex = (
r"<|DSML|function_calls>(.*?)</|DSML|function_calls>"
)
# Long-form `<|DSML|invoke name="x">...</|DSML|invoke>` and the
# self-closing `<|DSML|invoke name="x"/>` shape V4 emits for zero-arg
# tools. The `end` group is empty when the closer hasn't streamed in.
self.invoke_regex = (
r'<|DSML|invoke\s+name="([^"]+)"\s*>(.*?)(</|DSML|invoke>|$)'
r'<|DSML|invoke\s+name="(?P<name>[^"]+)"\s*'
r"(?:(?P<self_close>/>)"
r"|>(?P<body>.*?)(?P<end>(?:</|DSML|invoke>|$)))"
)
self.prefix_parameter_end_call = ["</", "|DSML|", "parameter"]
self.prefix_invoke_end_call = ["</", "|DSML|", "inv", "oke"]
@@ -92,6 +97,20 @@ class DeepSeekV32Detector(BaseFormatDetector):
"""Check if the text contains a deepseek v32 format tool call."""
return self.bot_token in text or "<|DSML|invoke" in text
@staticmethod
def _unpack_invoke_match(m: "re.Match[str]") -> tuple[str, str, bool]:
"""Returns (name, body, is_complete) for an invoke_regex match.
Self-closing invokes have empty body and are always complete.
Long-form bodies are always strings (possibly empty); they're
incomplete when matched against `$` because the closing tag
hasn't streamed in yet.
"""
name = m.group("name").strip()
if m.group("self_close"):
return name, "", True
return name, m.group("body"), bool(m.group("end"))
def _parse_parameters_from_xml(
self, invoke_content: str, allow_partial: bool = False
) -> str:
@@ -192,12 +211,10 @@ class DeepSeekV32Detector(BaseFormatDetector):
function_calls_content = function_calls_match.group(1)
# Find all invoke blocks
invoke_matches = re.findall(
for invoke_match in re.finditer(
self.invoke_regex, function_calls_content, re.DOTALL
)
for func_name, invoke_content, _ in invoke_matches:
# Parse parameters from XML format
):
func_name, invoke_content, _ = self._unpack_invoke_match(invoke_match)
func_args = self._parse_parameters_from_xml(invoke_content)
# construct match_result for parse_base_json
match_result = {"name": func_name, "parameters": json.loads(func_args)}
@@ -254,10 +271,9 @@ class DeepSeekV32Detector(BaseFormatDetector):
if not invoke_match:
break
func_name = invoke_match.group(1).strip()
invoke_content = invoke_match.group(2)
# group(3) is either "</|DSML|invoke>" (complete) or "" (incomplete, matched with $)
is_tool_end = bool(invoke_match.group(3))
func_name, invoke_content, is_tool_end = self._unpack_invoke_match(
invoke_match
)
# Initialize state if this is the first tool call
if self.current_tool_id == -1:
+5
View File
@@ -528,6 +528,9 @@ class DefaultModelLoader(BaseModelLoader):
weight_loader_disable_mmap = server_args.weight_loader_disable_mmap
weight_loader_prefetch = server_args.weight_loader_prefetch_checkpoints
prefetch_num_threads = server_args.weight_loader_prefetch_num_threads
weight_loader_drop_cache_after_load = (
server_args.weight_loader_drop_cache_after_load
)
if self.load_config.load_format == LoadFormat.FASTSAFETENSORS:
weights_iterator = fastsafetensors_weights_iterator(
@@ -542,6 +545,7 @@ class DefaultModelLoader(BaseModelLoader):
disable_mmap=weight_loader_disable_mmap,
prefetch=weight_loader_prefetch,
prefetch_num_threads=prefetch_num_threads,
drop_cache_after_load=weight_loader_drop_cache_after_load,
)
else:
weights_iterator = safetensors_weights_iterator(
@@ -549,6 +553,7 @@ class DefaultModelLoader(BaseModelLoader):
disable_mmap=weight_loader_disable_mmap,
prefetch=weight_loader_prefetch,
prefetch_num_threads=prefetch_num_threads,
drop_cache_after_load=weight_loader_drop_cache_after_load,
)
else:
+33 -3
View File
@@ -875,11 +875,30 @@ def _prefetch_all_checkpoints(
threading.Thread(target=_run_prefetch, daemon=True).start()
def _drop_file_cache_after_load(path: str) -> None:
"""Release of checkpoint pages after weights have been copied out. Used to avoid CPU OOM in RL."""
posix_fadvise = getattr(os, "posix_fadvise", None)
dontneed = getattr(os, "POSIX_FADV_DONTNEED", None)
if posix_fadvise is None or dontneed is None:
return
fd = None
try:
fd = os.open(path, os.O_RDONLY)
posix_fadvise(fd, 0, 0, dontneed)
except OSError as e:
logger.debug("Failed to drop file cache for %s: %s", path, e)
finally:
if fd is not None:
os.close(fd)
def safetensors_weights_iterator(
hf_weights_files: List[str],
disable_mmap: bool = False,
prefetch: bool = False,
prefetch_num_threads: int = 4,
drop_cache_after_load: bool = False,
) -> Generator[Tuple[str, torch.Tensor], None, None]:
"""Iterate over the weights in the model safetensor files."""
enable_tqdm = (
@@ -907,6 +926,8 @@ def safetensors_weights_iterator(
with safetensors.safe_open(st_file, framework="pt", device="cpu") as f:
for name in f.keys():
yield name, f.get_tensor(name)
if drop_cache_after_load:
_drop_file_cache_after_load(st_file)
def fastsafetensors_weights_iterator(
@@ -968,6 +989,7 @@ def multi_thread_safetensors_weights_iterator(
hf_weights_files: List[str],
max_workers: int,
disable_mmap: bool = False,
drop_cache_after_load: bool = False,
) -> Generator[Tuple[str, torch.Tensor], None, None]:
"""Multi-Thread iterate over the weights in the model safetensor files."""
enable_tqdm = (
@@ -977,8 +999,12 @@ def multi_thread_safetensors_weights_iterator(
def _load_file(st_file: str):
if disable_mmap:
with open(st_file, "rb") as f:
return safetensors.torch.load(f.read())
return safetensors.torch.load_file(st_file, device="cpu")
result = safetensors.torch.load(f.read())
else:
with safetensors.safe_open(st_file, framework="pt", device="cpu") as f:
result = {k: f.get_tensor(k) for k in f.keys()}
return st_file, result
with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor:
futures = [executor.submit(_load_file, st_file) for st_file in hf_weights_files]
@@ -995,9 +1021,12 @@ def multi_thread_safetensors_weights_iterator(
futures_iter = concurrent.futures.as_completed(futures)
for future in futures_iter:
state_dict = future.result()
st_file, state_dict = future.result()
for name, param in state_dict.items():
yield name, param
del state_dict
if drop_cache_after_load:
_drop_file_cache_after_load(st_file)
def buffered_multi_thread_safetensors_weights_iterator(
@@ -1006,6 +1035,7 @@ def buffered_multi_thread_safetensors_weights_iterator(
disable_mmap: bool = False,
prefetch: bool = False,
prefetch_num_threads: int = 4,
drop_cache_after_load: bool = False,
) -> Generator[Tuple[str, torch.Tensor], None, None]:
"""Multi-threaded safetensor loader with bounded memory via a sliding window.
+6
View File
@@ -791,6 +791,7 @@ class ServerArgs:
weight_loader_disable_mmap: bool = False
weight_loader_prefetch_checkpoints: bool = False
weight_loader_prefetch_num_threads: int = 4
weight_loader_drop_cache_after_load: bool = False
remote_instance_weight_loader_seed_instance_ip: Optional[str] = None
remote_instance_weight_loader_seed_instance_service_port: Optional[int] = None
remote_instance_weight_loader_send_weights_group_ports: Optional[List[int]] = None
@@ -6658,6 +6659,11 @@ class ServerArgs:
default=ServerArgs.weight_loader_prefetch_num_threads,
help="Number of threads per rank for checkpoint prefetching (default: 4).",
)
parser.add_argument(
"--weight-loader-drop-cache-after-load",
action="store_true",
help="Call posix_fadvise(DONTNEED) on each safetensors shard after loading it.",
)
parser.add_argument(
"--remote-instance-weight-loader-seed-instance-ip",
type=str,