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

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: