chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

[Auto Sync] Update weight_utils.py (20260212) (#18692)

Browse Source 
Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>
Co-authored-by: Dan Zheng <dzheng@x.ai>
This commit is contained in:
Lianmin Zheng
2026-02-12 16:26:05 -08:00
committed by GitHub
parent b48fe1d95e
commit 9815ee934c
3 changed files with 126 additions and 16 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
python/sglang/srt/layers/logits_processor.py
View File

@@ -55,7 +55,7 @@ from sglang.srt.model_executor.forward_batch_info import (
ForwardMode,
)
from sglang.srt.server_args import get_global_server_args
from sglang.srt.utils import is_npu, use_intel_amx_backend
from sglang.srt.utils.common import is_npu, use_intel_amx_backend
logger = logging.getLogger(__name__)
@@ -89,8 +89,8 @@ class LogitsProcessorOutput:
# The logprobs of input tokens. shape: [#token]
input_token_logprobs: Optional[torch.Tensor] = None
# The logprobs and ids of the top-k tokens in input positions. shape: [#seq, #token, k]
input_top_logprobs_val: List = None
input_top_logprobs_idx: List = None
input_top_logprobs_val: Optional[List] = None
input_top_logprobs_idx: Optional[List] = None
# The logprobs and ids of the requested token ids in input positions. shape: [#seq, n] (n is the number of requested token ids)
# Can contain either lists or GPU tensors (for delayed GPU-to-CPU transfer optimization)
input_token_ids_logprobs_val: Optional[List[Union[List[float], torch.Tensor]]] = (

28
python/sglang/srt/model_loader/weight_utils.py
View File

@@ -751,8 +751,12 @@ def safetensors_weights_iterator(
if disable_mmap:
with open(st_file, "rb") as f:
result = safetensors.torch.load(f.read())
for name, param in result.items():
yield name, param
# NOTE(xiuyu): safetensors.torch.load() may return keys in a
# different order per TP-shard file. Sort to ensure that cross-rank
# collectives (e.g. all_gather in sync_quantize_weight) stay
# aligned across TP ranks.
for name in sorted(result.keys()):
yield name, result[name]
else:
with safetensors.safe_open(st_file, framework="pt", device="cpu") as f:
for name in f.keys():
@@ -846,23 +850,31 @@ def multi_thread_safetensors_weights_iterator(
return safetensors.torch.load_file(st_file, device="cpu")
with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor:
futures = [executor.submit(_load_file, st_file) for st_file in hf_weights_files]
# Yield in deterministic order to keep cross-rank collectives aligned.
future_to_idx = {
executor.submit(_load_file, st_file): idx
for idx, st_file in enumerate(hf_weights_files)
}
results_by_idx: Dict[int, dict] = {}
if enable_tqdm:
futures_iter = tqdm(
concurrent.futures.as_completed(futures),
concurrent.futures.as_completed(future_to_idx),
total=len(hf_weights_files),
desc="Multi-thread loading shards",
disable=not enable_tqdm,
bar_format=BAR_FORMAT,
)
else:
futures_iter = concurrent.futures.as_completed(futures)
futures_iter = concurrent.futures.as_completed(future_to_idx)
for future in futures_iter:
state_dict = future.result()
for name, param in state_dict.items():
yield name, param
results_by_idx[future_to_idx[future]] = future.result()
for idx in range(len(hf_weights_files)):
state_dict = results_by_idx[idx]
for name in sorted(state_dict.keys()):
yield name, state_dict[name]
def _load_pt_file(bin_file: str) -> dict:

108
scripts/code_sync/copy_to_oss.py
View File

@@ -205,6 +205,97 @@ def get_oss_repo(dry_run):
return oss_root, temp_dir
def _apply_patch(patch_file, dry_run):
"""
Try to apply a patch, falling back to --3way merge if a clean apply fails.
Returns True if the patch was applied successfully (clean or via --3way),
or raises an exception with full diagnostics if both methods fail.
"""
# --- Attempt 1: clean git apply ---
apply_cmd = ["git", "apply", patch_file]
print(f"Run: {' '.join(apply_cmd)}")
if dry_run:
return True
result = subprocess.run(apply_cmd, capture_output=True, text=True)
if result.returncode == 0:
print("✅ Patch applied cleanly.")
return True
print(f"⚠️ Clean apply failed:\n{result.stderr.strip()}")
print("Falling back to git apply --3way ...\n")
# --- Attempt 2: three-way merge ---
threeway_cmd = ["git", "apply", "--3way", patch_file]
print(f"Run: {' '.join(threeway_cmd)}")
result_3way = subprocess.run(threeway_cmd, capture_output=True, text=True)
if result_3way.returncode == 0:
print("✅ Patch applied via --3way merge (no conflicts).")
return True
# --- Both attempts failed — gather diagnostics ---
print(f"❌ --3way merge also failed:\n{result_3way.stderr.strip()}\n")
# Show which hunks conflict
check_cmd = ["git", "apply", "--check", "--verbose", patch_file]
print(f"Run: {' '.join(check_cmd)}")
check_result = subprocess.run(check_cmd, capture_output=True, text=True)
conflict_details = (check_result.stdout + check_result.stderr).strip()
print(
f"\n--- Conflict details ---\n{conflict_details}\n--- End conflict details ---\n"
)
# Show git diff if --3way left conflict markers
diff_result = subprocess.run(["git", "diff"], capture_output=True, text=True)
if diff_result.stdout.strip():
print(
f"\n--- git diff (conflict markers) ---\n"
f"{diff_result.stdout.strip()}\n"
f"--- End git diff ---\n"
)
# Read the patch content for the summary
with open(patch_file, "r", encoding="utf-8") as pf:
patch_content = pf.read()
# Print the patch to stdout so it's visible in the CI logs
separator = "=" * 72
print(
f"\n{separator}\n"
f"PATCH CONTENT (apply this manually):\n"
f"{separator}\n"
f"{patch_content}\n"
f"{separator}\n"
)
# Write a rich summary to the GitHub Actions step summary
summary_lines = [
"\n## ❌ Patch could not be applied automatically\n",
"### Conflict details\n",
f"```\n{conflict_details}\n```\n",
]
if diff_result.stdout.strip():
summary_lines.append("### git diff (conflict markers)\n")
summary_lines.append(f"```diff\n{diff_result.stdout.strip()}\n```\n")
summary_lines.append("### Patch to apply manually\n")
summary_lines.append(
"<details><summary>Click to expand full patch</summary>\n\n"
f"```diff\n{patch_content}\n```\n"
"</details>\n"
)
write_github_step_summary("".join(summary_lines))
# Reset the working tree so the finally-block cleanup is safe
subprocess.run(["git", "checkout", "."], capture_output=True, text=True)
raise RuntimeError(
"Patch could not be applied. See the CI log and GitHub Actions "
"step summary for the full patch and conflict details."
)
def apply_patch_and_push(oss_root, patch_file, branch_name, commit_message, dry_run):
"""
In the OSS repo, create a branch, apply the patch, commit, and push.
@@ -216,10 +307,17 @@ def apply_patch_and_push(oss_root, patch_file, branch_name, commit_message, dry_
os.chdir(oss_root)
try:
# Define commands as lists to avoid shell injection issues
commands_to_run = [
["git", "checkout", "-b", branch_name],
["git", "apply", patch_file],
# Create a new branch
checkout_cmd = ["git", "checkout", "-b", branch_name]
print(f"Run: {' '.join(checkout_cmd)}")
if not dry_run:
subprocess.run(checkout_cmd, check=True, capture_output=True, text=True)
# Apply the patch (with --3way fallback and diagnostics)
_apply_patch(patch_file, dry_run)
# Configure git user and stage changes
post_apply_commands = [
["git", "config", "user.name", "github-actions[bot]"],
[
"git",
@@ -230,7 +328,7 @@ def apply_patch_and_push(oss_root, patch_file, branch_name, commit_message, dry_
["git", "add", "."],
]
for cmd_list in commands_to_run:
for cmd_list in post_apply_commands:
print(f"Run: {' '.join(cmd_list)}")
if not dry_run:
subprocess.run(cmd_list, check=True, capture_output=True, text=True)