[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:
co-authored by
Xinyuan Tong
yueming-yuan
parent
4b23f6bdc5
commit
ef5e9f8aba
@@ -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:
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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.
|
||||
|
||||
|
||||
@@ -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,
|
||||
|
||||
Reference in New Issue
Block a user