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

optm(checkpoint-engine): disable multi-thread loading when update weights (#12374)

Browse Source 
Signed-off-by: Yang Kaiyong <yangkaiyong.yky@antgroup.com>
This commit is contained in:
Yang Kaiyong
2025-11-07 21:30:04 +08:00
committed by GitHub
parent c67fce160e
commit 61bfd9fa9b
3 changed files with 110 additions and 30 deletions
Download Patch File Download Diff File Expand all files Collapse all files

43
python/sglang/srt/model_loader/utils.py
View File

@@ -1,9 +1,10 @@
# Adapted from https://github.com/vllm-project/vllm/blob/v0.6.4.post1/vllm/model_executor/model_loader/utils.py
"""Utilities for selecting and loading models."""
import concurrent.futures
import contextlib
import logging
from typing import Tuple, Type
from typing import Any, Callable, Dict, Iterable, List, Optional, Tuple, Type
import torch
import transformers
@@ -116,3 +117,43 @@ def post_load_weights(model: nn.Module, model_config: ModelConfig):
model.post_load_weights(is_nextn=True)
else:
model.post_load_weights()
def should_async_load(weight: torch.Tensor) -> bool:
"""Return True if we should load the given weight asynchronously.
For host (CPU) tensors, using a threadpool can overlap H2D copies
and improve throughput. For device tensors, threading often adds overhead
(e.g., GIL contention) without benefit, so we do it synchronously.
"""
device = getattr(weight, "device", None)
if device is None:
return False
return device.type == "cpu"
def maybe_executor_submit(
*,
executor: concurrent.futures.ThreadPoolExecutor,
futures: List[concurrent.futures.Future],
use_async: bool,
func: Callable[..., Any],
func_args: Iterable[Any] = (),
func_kwargs: Optional[Dict[str, Any]] = None,
) -> None:
"""Submit a task to the executor if async loading is enabled.
Parameters (keyword-only):
- executor: ThreadPoolExecutor used to submit background tasks
- futures: a list collecting the submitted Future objects
- use_async: whether to submit to executor or run inline
- func: the callable to run
- func_args: positional args for the callable (defaults to empty tuple)
- func_kwargs: keyword args for the callable (defaults to empty dict)
"""
if func_kwargs is None:
func_kwargs = {}
if use_async:
futures.append(executor.submit(func, *func_args, **func_kwargs))
else:
func(*func_args, **func_kwargs)

52
python/sglang/srt/models/deepseek_v2.py
View File

@@ -112,6 +112,7 @@ from sglang.srt.layers.vocab_parallel_embedding import (
VocabParallelEmbedding,
)
from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors
from sglang.srt.model_loader.utils import maybe_executor_submit, should_async_load
from sglang.srt.model_loader.weight_utils import default_weight_loader
from sglang.srt.server_args import get_global_server_args
from sglang.srt.single_batch_overlap import SboFlags
@@ -3505,6 +3506,7 @@ class DeepseekV2ForCausalLM(nn.Module):
params_dict = dict(self.named_parameters())
weight_names = []
for name, loaded_weight in weights:
use_async_loading = should_async_load(loaded_weight)
layer_id = get_layer_id(name)
if (
layer_id is not None
@@ -3572,8 +3574,12 @@ class DeepseekV2ForCausalLM(nn.Module):
continue
param = params_dict[name]
weight_loader = param.weight_loader
futures.append(
executor.submit(weight_loader, param, loaded_weight, shard_id)
maybe_executor_submit(
executor=executor,
futures=futures,
use_async=use_async_loading,
func=weight_loader,
func_args=(param, loaded_weight, shard_id),
)
break
else:
@@ -3584,15 +3590,20 @@ class DeepseekV2ForCausalLM(nn.Module):
name = name.replace(weight_name, param_name)
param = params_dict[name]
weight_loader = param.weight_loader
futures.append(
executor.submit(
weight_loader,
maybe_executor_submit(
executor=executor,
futures=futures,
use_async=use_async_loading,
func=weight_loader,
func_args=(
param,
loaded_weight,
name,
shard_id=shard_id,
expert_id=expert_id,
)
),
func_kwargs={
"shard_id": shard_id,
"expert_id": expert_id,
},
)
break
else:
@@ -3652,8 +3663,12 @@ class DeepseekV2ForCausalLM(nn.Module):
weight_loader = getattr(
param, "weight_loader", default_weight_loader
)
futures.append(
executor.submit(weight_loader, param, fused_weight)
maybe_executor_submit(
executor=executor,
futures=futures,
use_async=use_async_loading,
func=weight_loader,
func_args=(param, fused_weight),
)
cached_a_proj.pop(q_a_proj_name)
cached_a_proj.pop(kv_a_proj_name)
@@ -3699,8 +3714,13 @@ class DeepseekV2ForCausalLM(nn.Module):
weight_loader = getattr(
param, "weight_loader", default_weight_loader
)
futures.append(
executor.submit(weight_loader, param, fused_weight)
maybe_executor_submit(
executor,
futures,
use_async_loading,
weight_loader,
param,
fused_weight,
)
cached_wk_and_weights_proj.pop(wk_name)
cached_wk_and_weights_proj.pop(weights_proj_name)
@@ -3725,8 +3745,12 @@ class DeepseekV2ForCausalLM(nn.Module):
weight_loader = getattr(
param, "weight_loader", default_weight_loader
)
futures.append(
executor.submit(weight_loader, param, loaded_weight)
maybe_executor_submit(
executor=executor,
futures=futures,
use_async=use_async_loading,
func=weight_loader,
func_args=(param, loaded_weight),
)
# Wait for all tasks to complete and raise any exceptions.

45
python/sglang/srt/models/longcat_flash.py
View File

@@ -80,6 +80,7 @@ from sglang.srt.layers.vocab_parallel_embedding import (
VocabParallelEmbedding,
)
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
from sglang.srt.model_loader.utils import maybe_executor_submit, should_async_load
from sglang.srt.model_loader.weight_utils import default_weight_loader
from sglang.srt.models.deepseek_v2 import DeepseekV2AttentionMLA
from sglang.srt.server_args import get_global_server_args
@@ -854,6 +855,7 @@ class LongcatFlashForCausalLM(nn.Module):
params_dict = dict(self.named_parameters())
weight_names = []
for name, loaded_weight in weights:
use_async_loading = should_async_load(loaded_weight)
if "mtp" in name:
continue
weight_names.append(name)
@@ -877,8 +879,12 @@ class LongcatFlashForCausalLM(nn.Module):
continue
param = params_dict[name]
weight_loader = param.weight_loader
futures.append(
executor.submit(weight_loader, param, loaded_weight, shard_id)
maybe_executor_submit(
executor=executor,
futures=futures,
use_async=use_async_loading,
func=weight_loader,
func_args=(param, loaded_weight, shard_id),
)
break
else:
@@ -889,15 +895,16 @@ class LongcatFlashForCausalLM(nn.Module):
name = name.replace(weight_name, param_name)
param = params_dict[name]
weight_loader = param.weight_loader
futures.append(
executor.submit(
weight_loader,
param,
loaded_weight,
name,
shard_id=shard_id,
expert_id=expert_id,
)
maybe_executor_submit(
executor=executor,
futures=futures,
use_async=use_async_loading,
func=weight_loader,
func_args=(param, loaded_weight, name),
func_kwargs={
"shard_id": shard_id,
"expert_id": expert_id,
},
)
break
else:
@@ -951,8 +958,12 @@ class LongcatFlashForCausalLM(nn.Module):
weight_loader = getattr(
param, "weight_loader", default_weight_loader
)
futures.append(
executor.submit(weight_loader, param, fused_weight)
maybe_executor_submit(
executor=executor,
futures=futures,
use_async=use_async_loading,
func=weight_loader,
func_args=(param, fused_weight),
)
cached_a_proj.pop(q_a_proj_name)
cached_a_proj.pop(kv_a_proj_name)
@@ -977,8 +988,12 @@ class LongcatFlashForCausalLM(nn.Module):
weight_loader = getattr(
param, "weight_loader", default_weight_loader
)
futures.append(
executor.submit(weight_loader, param, loaded_weight)
maybe_executor_submit(
executor=executor,
futures=futures,
use_async=use_async_loading,
func=weight_loader,
func_args=(param, loaded_weight),
)
# Wait for all tasks to complete and raise any exceptions.