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

Fix OOM in DeepSeek weight loading by deferring dict(weights) materialization (#17744)

Browse Source
This commit is contained in:
jeff
2026-01-31 13:59:00 +08:00
committed by GitHub
parent 45fe51a28e
commit 04efd03dbf
1 changed files with 16 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/models/deepseek_common/deepseek_weight_loader.py
+16 -12
View File
@@ -623,9 +623,9 @@ class DeepseekV2WeightLoaderMixin:
nextn_conf: NextN configuration nextn_conf: NextN configuration
Returns: Returns:
List of (name, tensor) pairs with quantized weights Original weights iterator if no quantization needed,
otherwise list of (name, tensor) pairs with quantized weights
""" """
weights_dict = dict(weights)
weight_block_size = [128, 128] weight_block_size = [128, 128]
partial_names = [] partial_names = []
@@ -655,16 +655,20 @@ class DeepseekV2WeightLoaderMixin:
f"model.layers.{layer_id}.self_attn.{stem}" f"model.layers.{layer_id}.self_attn.{stem}"
) )
if partial_names: # Early return if no quantization needed - avoid materializing all weights into memory
for partial_name in tqdm.tqdm( if not partial_names:
partial_names, desc="quant weights to fp8 ue8m0" return weights
):
original_weight = weights_dict[f"{partial_name}.weight"] # Only materialize weights dict when quantization is actually needed
out_w, out_s = quant_weight_ue8m0( weights_dict = dict(weights)
original_weight, weight_block_size=weight_block_size
) for partial_name in tqdm.tqdm(partial_names, desc="quant weights to fp8 ue8m0"):
weights_dict[f"{partial_name}.weight"] = out_w original_weight = weights_dict[f"{partial_name}.weight"]
weights_dict[f"{partial_name}.weight_scale_inv"] = out_s out_w, out_s = quant_weight_ue8m0(
original_weight, weight_block_size=weight_block_size
)
weights_dict[f"{partial_name}.weight"] = out_w
weights_dict[f"{partial_name}.weight_scale_inv"] = out_s
if isinstance( if isinstance(
nextn_conf, NextNEnabledConfig nextn_conf, NextNEnabledConfig