[weight checker] refactor: add precision branch; allow ULP quant err; used chunked compare (#28974)

python/sglang/srt/managers/io_struct.py

@@ -1712,6 +1712,7 @@ class ResumeMemoryOccupationReqOutput(BaseReq, kw_only=True):
 
 class CheckWeightsReqInput(BaseReq, kw_only=True):
     action: str = "checksum"
+    allow_quant_error: bool = False
 
 
 class CheckWeightsReqOutput(BaseReq, kw_only=True):

python/sglang/srt/managers/scheduler_components/weight_updater.py

@@ -268,12 +268,17 @@ class SchedulerWeightUpdaterManager:
 
     def check_weights(self, recv_req: CheckWeightsReqInput):
         try:
-            payload = self.tp_worker.model_runner.check_weights(action=recv_req.action)
+            payload = self.tp_worker.model_runner.check_weights(
+                action=recv_req.action, allow_quant_error=recv_req.allow_quant_error
+            )
 
             if self.draft_worker is not None:
                 draft_runner = _get_draft_model_runner(self.draft_worker)
                 if draft_runner is not None:
-                    draft_payload = draft_runner.check_weights(action=recv_req.action)
+                    draft_payload = draft_runner.check_weights(
+                        action=recv_req.action,
+                        allow_quant_error=recv_req.allow_quant_error,
+                    )
                     if payload is not None and draft_payload is not None:
                         payload = _merge_checksum_payloads(payload, draft_payload)
 

python/sglang/srt/model_executor/model_runner.py

@@ -3205,8 +3205,10 @@ class ModelRunner(ModelRunnerKVCacheMixin):
         )
         ShardedStateLoader.save_model(self.model, path, pattern, max_size)
 
-    def check_weights(self, action: str):
-        return self._weight_checker.handle(action=action)
+    def check_weights(self, action: str, allow_quant_error: bool = False):
+        return self._weight_checker.handle(
+            action=action, allow_quant_error=allow_quant_error
+        )
 
     def update_weights_from_ipc(self, recv_req):
         """Update weights from IPC for checkpoint-engine integration."""

python/sglang/srt/utils/weight_checker.py

@@ -1,17 +1,20 @@
 import hashlib
 import logging
 import time
-from typing import Dict, Iterable, Optional, Set, Tuple
+from typing import Dict, Iterable, NamedTuple, Optional, Set
 
 import torch
 import torch.distributed as dist
 from pydantic import BaseModel, ConfigDict
 
-from sglang.srt.layers.quantization.fp8_utils import (
-    block_quant_dequant,
-    inverse_transform_scale_ue8m0,
-)
 from sglang.srt.managers.mm_utils import tensor_hash
+from sglang.srt.utils.weight_checker_comparator import (
+    CHUNK_NUMEL,
+    ComparableWeight,
+    RawComparable,
+    compare_weights,
+    select_comparable_weight,
+)
 
 logger = logging.getLogger(__name__)
 
@@ -37,6 +40,17 @@ class ChecksumInfo(_StrictBaseModel):
     parallelism_info: ParallelismInfo
 
 
+class CheckEntry(NamedTuple):
+    name: str
+    should_compare: bool
+    comparable: ComparableWeight
+
+
+class QuantizedWeight(NamedTuple):
+    comparable_cls: type[ComparableWeight]
+    scale_name: str
+
+
 _NON_PERSISTENT_BUFFER_PATTERNS = (
     "cos_sin_cache",
     "inv_freq",
@@ -54,14 +68,16 @@ class WeightChecker:
         self._model_runner = model_runner
         self._snapshot_tensors = None
 
-    def handle(self, action: str) -> Optional[Dict]:
-        logger.info(f"[WeightChecker] handle action={action}")
+    def handle(self, action: str, allow_quant_error: bool = False) -> Optional[Dict]:
+        logger.info(
+            f"[WeightChecker] handle action={action} allow_quant_error={allow_quant_error}"
+        )
         if action == "snapshot":
             return self._snapshot()
         elif action == "reset_tensors":
             return self._reset_tensors()
         elif action == "compare":
-            return self._compare()
+            return self._compare(allow_quant_error=allow_quant_error)
         elif action == "checksum":
             return self._compute_checksum()
         else:
@@ -82,43 +98,44 @@ class WeightChecker:
                 continue
             param.copy_(_random_like(param))
 
-    def _compare(self):
+    def _compare(self, allow_quant_error: bool = False):
         assert self._snapshot_tensors is not None
 
+        quantized_set = _build_quantized_set(self._model_runner.model)
         skip_compare_names = {
             name
             for name, param in self._model_state()
             if getattr(param, "_skip_weight_check", False)
         }
         _check_tensors(
-            expect_tensors=_postprocess_tensors(
-                self._snapshot_tensors, skip_compare_names
+            expect_tensors=_build_check_entries(
+                self._snapshot_tensors, skip_compare_names, quantized_set
             ),
-            actual_tensors=_postprocess_tensors(
-                dict(self._model_state()), skip_compare_names
+            actual_tensors=_build_check_entries(
+                dict(self._model_state()), skip_compare_names, quantized_set
             ),
+            allow_quant_error=allow_quant_error,
         )
 
     def _compute_checksum(self) -> Dict:
         torch.cuda.synchronize()
         start = time.perf_counter()
 
+        quantized_set = _build_quantized_set(self._model_runner.model)
         skip_compare_names = {
             name
             for name, param in self._model_state()
             if getattr(param, "_skip_weight_check", False)
         }
 
-        # Reuse the snapshot/compare postprocess pipeline so fp8 weights are
-        # dequantized to bf16 before hashing — two (qweight, scale) pairs that
-        # produce the same bf16 must produce the same checksum.
-        checksums = {
-            name: _hash_tensor(tensor.data)
-            for name, should_compare, tensor in _postprocess_tensors(
-                dict(self._model_state()), skip_compare_names
-            )
-            if should_compare
-        }
+        # Hash the dequantized weight so two (qweight, scale) pairs with the same
+        # bf16 hash equal.
+        checksums = {}
+        for name, should_compare, comparable in _build_check_entries(
+            dict(self._model_state()), skip_compare_names, quantized_set
+        ):
+            if should_compare:
+                checksums[name] = _hash_tensor(comparable.dequantize().data)
 
         h = hashlib.sha256()
         for name in sorted(checksums):
@@ -162,42 +179,50 @@ def _hash_tensor(t: torch.Tensor) -> str:
 
 
 def _check_tensors(
-    expect_tensors: Iterable[Tuple[str, bool, torch.Tensor]],
-    actual_tensors: Iterable[Tuple[str, bool, torch.Tensor]],
+    expect_tensors: Iterable[CheckEntry],
+    actual_tensors: Iterable[CheckEntry],
+    allow_quant_error: bool = False,
 ):
-    from sglang.srt.debug_utils.dumper import get_tensor_info
-
     good_names = []
     error_messages = []
     info_messages = []
 
-    for (expect_name, expect_should_compare, expect), (
+    for (expect_name, should_compare, expect_comparable), (
         actual_name,
         actual_should_compare,
-        actual,
+        actual_comparable,
     ) in zip(expect_tensors, actual_tensors, strict=True):
         assert expect_name == actual_name, f"{expect_name=} {actual_name=}"
         assert (
-            expect_should_compare == actual_should_compare
-        ), f"{expect_should_compare=} {actual_should_compare=}"
+            should_compare == actual_should_compare
+        ), f"{should_compare=} {actual_should_compare=}"
         name = expect_name
-        should_compare = expect_should_compare
 
-        expect = expect.cuda()
-        actual = actual.cuda()
-
-        if torch.all(expect == actual):
-            good_names.append(name)
-        else:
-            abs_diff = (actual.float() - expect.float()).abs()
-            msg = (
-                f"name={name} "
-                f"max_abs_err={abs_diff.max()} "
-                f"mean_abs_err={abs_diff.mean()} "
-                f"{get_tensor_info(expect)=} "
-                f"{get_tensor_info(actual)=} "
+        try:
+            equal, max_abs_err, mean_abs_err, num_exceed = compare_weights(
+                expect_comparable, actual_comparable
             )
-            (error_messages if should_compare else info_messages).append(msg)
+        except Exception as e:
+            e.add_note(
+                f"when handling {name=} expect={expect_comparable!r} actual={actual_comparable!r}"
+            )
+            raise
+        if equal:
+            good_names.append(name)
+            continue
+        msg = (
+            f"name={name} "
+            f"max_abs_err={max_abs_err} "
+            f"mean_abs_err={mean_abs_err} "
+            f"num_exceed={num_exceed} "
+            f"expect={expect_comparable!r} actual={actual_comparable!r} "
+        )
+        if not should_compare:
+            info_messages.append(msg)
+        elif allow_quant_error and num_exceed == 0:
+            info_messages.append(msg + "(within quantization ULP tolerance)")
+        else:
+            error_messages.append(msg)
 
     logger.info(f"[check_tensors] equal tensors: {good_names}")
     if len(info_messages) > 0:
@@ -212,7 +237,12 @@ def _random_like(t: torch.Tensor):
     dtype = t.dtype
 
     if dtype.is_floating_point:
-        return torch.rand(shape, device=device, dtype=torch.float32).to(dtype)
+        out = torch.empty(shape, device=device, dtype=dtype)
+        for chunk in out.view(-1).split(CHUNK_NUMEL):
+            chunk.copy_(
+                torch.rand(chunk.shape, device=device, dtype=torch.float32).to(dtype)
+            )
+        return out
 
     if dtype == torch.bool:
         return torch.rand(shape, device=device) > 0.5
@@ -223,58 +253,44 @@ def _random_like(t: torch.Tensor):
     )
 
 
-def _postprocess_tensors(
+def _build_quantized_set(model) -> Dict[str, QuantizedWeight]:
+    """Run the router over the model: {weight_name: QuantizedWeight} for each
+    quantized weight; weights absent from the set compare raw."""
+    quantized_set = {}
+    for module_name, module in model.named_modules():
+        comparable_cls = select_comparable_weight(getattr(module, "quant_method", None))
+        if comparable_cls is None:
+            continue
+        prefix = f"{module_name}." if module_name else ""
+        own = {name for name, _ in module.named_parameters(recurse=False)}
+        for name in own:
+            scale = name.replace("weight", "weight_scale_inv")
+            if name.endswith("weight") and scale in own:
+                quantized_set[prefix + name] = QuantizedWeight(
+                    comparable_cls, prefix + scale
+                )
+    return quantized_set
+
+
+def _build_check_entries(
     raw: Dict[str, torch.Tensor],
     skip_compare_names: Set[str],
-) -> Iterable[Tuple[str, bool, torch.Tensor]]:
-    from sglang.srt.debug_utils.dumper import get_tensor_info
-
+    quantized_set: Optional[Dict[str, QuantizedWeight]] = None,
+) -> Iterable[CheckEntry]:
+    """Yields a CheckEntry per weight; quantized weights consume their scale, everything
+    else is raw."""
     skip_compare_names = set(skip_compare_names)
+    quantized_set = quantized_set or {}
+    scale_names = {qw.scale_name for qw in quantized_set.values()}
 
-    # Skip non-persistent buffers (registered with persistent=False; recomputed
-    # after weight load and not part of the synced payload).
-    for name in raw:
-        if _is_non_persistent_buffer_name(name):
-            skip_compare_names.add(name)
-            logger.info(f"[check_tensors] Skipping non-persistent buffer: {name}")
-
-    # dequant fp8
-    quant_names = [
-        name
-        for name in raw
-        # Match: `something.weight`, `something.experts.w2_weight`
-        if name.endswith("weight") and name.replace("weight", "weight_scale_inv") in raw
-    ]
-    quant_scale_names = [
-        name.replace("weight", "weight_scale_inv") for name in quant_names
-    ]
-    skip_compare_names.update(quant_names)
-    skip_compare_names.update(quant_scale_names)
-    for name in quant_names:
-        w_q = raw[name]
-        w_s = raw[name.replace("weight", "weight_scale_inv")]
-
-        try:
-            if w_s.dtype == torch.int32:
-                # UE8M0 packed format (Blackwell DeepGEMM)
-                w_s_for_dequant = inverse_transform_scale_ue8m0(w_s, mn=w_q.shape[-2])
-            else:
-                w_s_for_dequant = w_s
-
-            w_dequant = block_quant_dequant(
-                w_q,
-                w_s_for_dequant,
-                # TODO do not hardcode
-                block_size=[128, 128],
-                dtype=torch.bfloat16,
+    for name, tensor in raw.items():
+        if name in scale_names:
+            continue  # compared via its weight's comparable
+        if name in quantized_set:
+            qw = quantized_set[name]
+            yield CheckEntry(name, True, qw.comparable_cls(tensor, raw[qw.scale_name]))
+        else:
+            should_compare = name not in skip_compare_names and (
+                not _is_non_persistent_buffer_name(name)
             )
-            yield name, True, w_dequant
-        except Exception as e:
-            e.add_note(
-                f"when handling {name=} {get_tensor_info(w_q)=} {get_tensor_info(w_s)=}"
-            )
-            raise
-
-    for name in raw:
-        should_compare = name not in skip_compare_names
-        yield name, should_compare, raw[name]
+            yield CheckEntry(name, should_compare, RawComparable(tensor))

python/sglang/srt/utils/weight_checker_comparator.py

@@ -0,0 +1,166 @@
+from typing import Iterable, NamedTuple, Optional, Tuple
+
+import torch
+
+from sglang.srt.layers.quantization.fp8 import Fp8LinearMethod, Fp8MoEMethod
+from sglang.srt.layers.quantization.fp8_utils import (
+    block_quant_dequant,
+    inverse_transform_scale_ue8m0,
+)
+from sglang.srt.layers.quantization.modelopt_quant import (
+    ModelOptFp4LinearMethod,
+    ModelOptNvFp4FusedMoEMethod,
+)
+
+# chunk to avoid too high GPU memory peak
+CHUNK_NUMEL = 64 * 1024 * 1024
+
+
+class CompareResult(NamedTuple):
+    equal: bool
+    max_abs_err: float
+    mean_abs_err: float
+    num_exceed: int  # elements past the combined per-side tolerance
+
+
+class ComparableWeight:
+    """Base comparable-weight class; one subclass per precision or raw tensor."""
+
+    @staticmethod
+    def _quant_ulp(w_q: torch.Tensor) -> torch.Tensor:
+        """Per-element ULP of w_q in its own dtype."""
+        finfo = torch.finfo(w_q.dtype)
+        x = w_q.to(torch.float32).abs()
+        # frexp: x = m * 2^e, m in [0.5, 1), so 2^(e-1) is x's binade base.
+        _, exponent = torch.frexp(x)
+        binade = torch.exp2((exponent - 1).to(torch.float32))
+        # Zeros and subnormals share the spacing of the smallest normal binade.
+        binade = binade.masked_fill(x < finfo.smallest_normal, finfo.smallest_normal)
+        return binade * finfo.eps
+
+    def iter_chunks(self) -> Iterable[Tuple[torch.Tensor, Optional[torch.Tensor]]]:
+        raise NotImplementedError
+
+    def dequantize(self, dtype: torch.dtype = torch.bfloat16) -> torch.Tensor:
+        raise NotImplementedError
+
+
+class Fp8BlockComparable(ComparableWeight):
+    """Deepseek-style FP8 quantization."""
+
+    def __init__(self, w_q: torch.Tensor, w_s: torch.Tensor):
+        self.w_q = w_q
+        self.w_s = w_s
+
+    def __repr__(self) -> str:
+        return f"fp8_block(shape={tuple(self.w_q.shape)} dtype={self.w_q.dtype})"
+
+    @staticmethod
+    def _normalize_scale(w_q: torch.Tensor, w_s: torch.Tensor) -> torch.Tensor:
+        if w_s.dtype == torch.int32:
+            w_s = inverse_transform_scale_ue8m0(w_s, mn=w_q.shape[-2])
+        return w_s.to(torch.float32)
+
+    @staticmethod
+    def _infer_block_size(w_q: torch.Tensor, w_s: torch.Tensor) -> list:
+        k, s_k = w_q.shape[-1], w_s.shape[-1]
+        assert k % s_k == 0, f"cannot infer block size from {w_q.shape=} {w_s.shape=}"
+        block = k // s_k
+        return [block, block]
+
+    @staticmethod
+    def _iter_quant_chunks(w_q: torch.Tensor, w_s: torch.Tensor, block_n: int):
+        """Yields block-row-aligned (q_slice, s_slice) pairs of bounded size."""
+        q3 = w_q.reshape(-1, *w_q.shape[-2:])
+        s3 = w_s.reshape(-1, *w_s.shape[-2:])
+        n, k = q3.shape[-2:]
+        rows = max(block_n, CHUNK_NUMEL // k // block_n * block_n)
+        for b in range(q3.shape[0]):
+            for r0 in range(0, n, rows):
+                r1 = min(r0 + rows, n)
+                yield q3[b, r0:r1], s3[b, r0 // block_n : -(-r1 // block_n)]
+
+    def _scale_and_block_size(self):
+        s = self._normalize_scale(self.w_q, self.w_s)
+        return s, self._infer_block_size(self.w_q, s)
+
+    def iter_chunks(self):
+        s, block_size = self._scale_and_block_size()
+        for q, s_chunk in self._iter_quant_chunks(self.w_q, s, block_size[0]):
+            q, s_chunk = q.cuda(), s_chunk.cuda()
+            yield (
+                block_quant_dequant(q, s_chunk, block_size, dtype=torch.bfloat16),
+                block_quant_dequant(
+                    self._quant_ulp(q), s_chunk, block_size, dtype=torch.float32
+                ),
+            )
+
+    def dequantize(self, dtype: torch.dtype = torch.bfloat16) -> torch.Tensor:
+        s, block_size = self._scale_and_block_size()
+        return block_quant_dequant(self.w_q, s, block_size, dtype=dtype)
+
+
+class RawComparable(ComparableWeight):
+    """Bitwise equal compare on raw tensor."""
+
+    def __init__(self, tensor: torch.Tensor):
+        self.tensor = tensor
+
+    def __repr__(self) -> str:
+        return f"raw(shape={tuple(self.tensor.shape)} dtype={self.tensor.dtype})"
+
+    def iter_chunks(self):
+        flat = self.tensor.reshape(-1)
+        for start in range(0, flat.numel(), CHUNK_NUMEL):
+            yield flat[start : start + CHUNK_NUMEL].cuda(), None
+
+    def dequantize(self, dtype: torch.dtype = torch.bfloat16) -> torch.Tensor:
+        return self.tensor
+
+
+def compare_weights(
+    expect: ComparableWeight, actual: ComparableWeight
+) -> CompareResult:
+    """Chunked element-wise compare in ComparableWeight space."""
+    equal = True
+    max_abs_err = torch.zeros((), dtype=torch.float32)
+    sum_abs_err = 0.0
+    num_exceed = 0
+    numel = 0
+    for (expect_dq, expect_tol), (actual_dq, actual_tol) in zip(
+        expect.iter_chunks(), actual.iter_chunks(), strict=True
+    ):
+        assert (
+            expect_dq.shape == actual_dq.shape
+        ), f"{expect_dq.shape=} {actual_dq.shape=}"
+        numel += expect_dq.numel()
+        abs_diff = (actual_dq.float() - expect_dq.float()).abs()
+        if torch.all(abs_diff == 0):
+            continue
+        equal = False
+        # |actual_dq - expect_dq| ≤ |actual_dq - w| + |expect_dq - w| ≤ actual_tol + expect_tol
+        tol = (
+            0.0 if expect_tol is None or actual_tol is None else expect_tol + actual_tol
+        )
+        max_abs_err = torch.maximum(max_abs_err, abs_diff.max().cpu())
+        sum_abs_err += abs_diff.sum().item()
+        # `~(diff <= tol)` instead of `diff > tol` so NaN counts as exceeding.
+        num_exceed += int((~(abs_diff <= tol)).sum())
+    return CompareResult(
+        equal, max_abs_err.item(), sum_abs_err / max(numel, 1), num_exceed
+    )
+
+
+def select_comparable_weight(quant_method) -> Optional[type]:
+    """Map a module's quant_method to its ComparableWeight. None means raw (bitwise equal) compare."""
+    if (
+        isinstance(quant_method, (Fp8LinearMethod, Fp8MoEMethod))
+        and quant_method.block_quant
+        and not quant_method.use_mxfp8
+    ):
+        return Fp8BlockComparable
+    if isinstance(quant_method, (ModelOptFp4LinearMethod, ModelOptNvFp4FusedMoEMethod)):
+        raise NotImplementedError(
+            f"weight checker has no ComparableWeight for {type(quant_method).__name__}"
+        )
+    return None