Refactor buffer patterns in weight checker (#24538)

This commit is contained in:
fzyzcjy
2026-05-06 22:52:07 +08:00
committed by GitHub
parent 800deaaefa
commit ae5ae840f6
+18 -15
View File
@@ -11,6 +11,18 @@ from sglang.srt.layers.quantization.fp8_utils import (
logger = logging.getLogger(__name__)
_NON_PERSISTENT_BUFFER_PATTERNS = (
"cos_sin_cache",
"inv_freq",
"freqs_cis",
"_weight_fp32",
)
def _is_non_persistent_buffer_name(name: str) -> bool:
return any(pat in name for pat in _NON_PERSISTENT_BUFFER_PATTERNS)
class WeightChecker:
def __init__(self, model_runner):
self._model_runner = model_runner
@@ -38,7 +50,7 @@ class WeightChecker:
def _reset_tensors(self):
for name, param in self._model_state():
if "cos_sin_cache" in name or "freqs_cis" in name or "_weight_fp32" in name:
if _is_non_persistent_buffer_name(name):
continue
param.copy_(_random_like(param))
@@ -125,21 +137,12 @@ def _postprocess_tensors(
skip_compare_names = []
# Skip non-persistent buffers like cos_sin_cache
# These buffers are registered with persistent=False and are not saved in checkpoints
# They should be recomputed after loading weights, so we don't compare them here
non_persistent_buffer_patterns = [
"cos_sin_cache", # RoPE cache
"inv_freq", # RoPE inverse frequency (if it exists as buffer)
"_weight_fp32", # FP32 cache of gate weight (e.g. Glm4MoeGate)
]
# Skip non-persistent buffers (registered with persistent=False; recomputed
# after weight load and not part of the synced payload).
for name in raw:
for pattern in non_persistent_buffer_patterns:
if pattern in name:
skip_compare_names.append(name)
logger.info(f"[check_tensors] Skipping non-persistent buffer: {name}")
break
if _is_non_persistent_buffer_name(name):
skip_compare_names.append(name)
logger.info(f"[check_tensors] Skipping non-persistent buffer: {name}")
# dequant fp8
quant_names = [