302 lines
9.7 KiB
Python
302 lines
9.7 KiB
Python
import hashlib
|
|
import logging
|
|
import time
|
|
from typing import Any, Callable, Dict, Iterable, NamedTuple, Optional, Set
|
|
|
|
import torch
|
|
import torch.distributed as dist
|
|
from pydantic import BaseModel, ConfigDict
|
|
|
|
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__)
|
|
|
|
|
|
class _StrictBaseModel(BaseModel):
|
|
model_config = ConfigDict(extra="forbid")
|
|
|
|
|
|
class ParallelismInfo(_StrictBaseModel):
|
|
tp_rank: int
|
|
tp_size: int
|
|
dp_rank: int
|
|
dp_size: int
|
|
pp_rank: int
|
|
pp_size: int
|
|
rank: int
|
|
size: int
|
|
|
|
|
|
class ChecksumInfo(_StrictBaseModel):
|
|
checksums: Dict[str, str]
|
|
per_gpu_checksum: str
|
|
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",
|
|
"freqs_cis",
|
|
"expert_mask_gpu",
|
|
)
|
|
|
|
|
|
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, *, get_model: Callable[[], Any], ps: Any):
|
|
self._get_model = get_model
|
|
self._ps = ps
|
|
self._snapshot_tensors = None
|
|
|
|
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(allow_quant_error=allow_quant_error)
|
|
elif action == "checksum":
|
|
return self._compute_checksum()
|
|
else:
|
|
raise Exception(f"Unsupported {action=}")
|
|
|
|
def _snapshot(self):
|
|
named_tensors = [
|
|
(name, param.data.detach().cpu()) for name, param in self._model_state()
|
|
]
|
|
self._snapshot_tensors = dict(named_tensors)
|
|
assert len(self._snapshot_tensors) == len(
|
|
named_tensors
|
|
), f"should not have duplicated tensor name"
|
|
|
|
def _reset_tensors(self):
|
|
for name, param in self._model_state():
|
|
if _is_non_persistent_buffer_name(name):
|
|
continue
|
|
param.copy_(_random_like(param))
|
|
|
|
def _compare(self, allow_quant_error: bool = False):
|
|
assert self._snapshot_tensors is not None
|
|
|
|
quantized_set = _build_quantized_set(self._get_model())
|
|
skip_compare_names = {
|
|
name
|
|
for name, param in self._model_state()
|
|
if getattr(param, "_skip_weight_check", False)
|
|
}
|
|
_check_tensors(
|
|
expect_tensors=_build_check_entries(
|
|
self._snapshot_tensors, skip_compare_names, quantized_set
|
|
),
|
|
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._get_model())
|
|
skip_compare_names = {
|
|
name
|
|
for name, param in self._model_state()
|
|
if getattr(param, "_skip_weight_check", False)
|
|
}
|
|
|
|
# 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):
|
|
h.update(name.encode())
|
|
h.update(checksums[name].encode())
|
|
overall = h.hexdigest()
|
|
|
|
torch.cuda.synchronize()
|
|
elapsed = time.perf_counter() - start
|
|
logger.info(
|
|
f"[WeightChecker] checksum computed for {len(checksums)} tensors in {elapsed:.3f}s"
|
|
)
|
|
|
|
info = ChecksumInfo(
|
|
checksums=checksums,
|
|
per_gpu_checksum=overall,
|
|
parallelism_info=self._parallelism_info(),
|
|
)
|
|
return info.model_dump()
|
|
|
|
def _parallelism_info(self) -> ParallelismInfo:
|
|
ps = self._ps
|
|
return ParallelismInfo(
|
|
tp_rank=ps.tp_rank,
|
|
tp_size=ps.tp_size,
|
|
dp_rank=ps.dp_rank if ps.dp_rank is not None else 0,
|
|
dp_size=ps.attn_dp_size,
|
|
pp_rank=ps.pp_rank,
|
|
pp_size=ps.pp_size,
|
|
rank=dist.get_rank() if dist.is_initialized() else 0,
|
|
size=dist.get_world_size() if dist.is_initialized() else 1,
|
|
)
|
|
|
|
def _model_state(self):
|
|
model = self._get_model()
|
|
yield from model.named_parameters()
|
|
yield from model.named_buffers()
|
|
|
|
|
|
def _hash_tensor(t: torch.Tensor) -> str:
|
|
return f"{tensor_hash(t):016x}"
|
|
|
|
|
|
def _check_tensors(
|
|
expect_tensors: Iterable[CheckEntry],
|
|
actual_tensors: Iterable[CheckEntry],
|
|
allow_quant_error: bool = False,
|
|
):
|
|
good_names = []
|
|
error_messages = []
|
|
info_messages = []
|
|
|
|
for (expect_name, should_compare, expect_comparable), (
|
|
actual_name,
|
|
actual_should_compare,
|
|
actual_comparable,
|
|
) in zip(expect_tensors, actual_tensors, strict=True):
|
|
if ".cos_sin_cache" in expect_name:
|
|
# skip cos/sin cache which is deterministic from shape and dtype and may have different shapes due to different implementations.
|
|
continue
|
|
assert expect_name == actual_name, f"{expect_name=} {actual_name=}"
|
|
assert (
|
|
should_compare == actual_should_compare
|
|
), f"{should_compare=} {actual_should_compare=}"
|
|
name = expect_name
|
|
|
|
try:
|
|
equal, max_abs_err, mean_abs_err, num_exceed = compare_weights(
|
|
expect_comparable, actual_comparable
|
|
)
|
|
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:
|
|
logger.info(f"[check_tensors] info: {info_messages}")
|
|
if len(error_messages) > 0:
|
|
raise Exception(f"check tensor equality failed:\n" + "\n".join(error_messages))
|
|
|
|
|
|
def _random_like(t: torch.Tensor):
|
|
device = t.device
|
|
shape = t.shape
|
|
dtype = t.dtype
|
|
|
|
if dtype.is_floating_point:
|
|
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
|
|
|
|
info = torch.iinfo(dtype)
|
|
return torch.randint(
|
|
low=int(info.min), high=int(info.max), size=shape, device=device, dtype=dtype
|
|
)
|
|
|
|
|
|
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],
|
|
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()}
|
|
|
|
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 CheckEntry(name, should_compare, RawComparable(tensor))
|