From d54f0a10b4945f5bf4fde3b127ba1b3c40f0cdfd Mon Sep 17 00:00:00 2001 From: fzyzcjy <5236035+fzyzcjy@users.noreply.github.com> Date: Thu, 8 Jan 2026 20:19:25 +0800 Subject: [PATCH] Support bitwise weight checksum verifier (#16729) --- python/sglang/srt/model_loader/loader.py | 7 + python/sglang/srt/server_args.py | 9 + .../sglang/srt/utils/model_file_verifier.py | 167 ++++++++++++++++++ .../utils/test_model_file_verifier.py | 157 ++++++++++++++++ 4 files changed, 340 insertions(+) create mode 100644 python/sglang/srt/utils/model_file_verifier.py create mode 100644 test/registered/utils/test_model_file_verifier.py diff --git a/python/sglang/srt/model_loader/loader.py b/python/sglang/srt/model_loader/loader.py index deaa1fff2..2477ade45 100644 --- a/python/sglang/srt/model_loader/loader.py +++ b/python/sglang/srt/model_loader/loader.py @@ -422,6 +422,13 @@ class DefaultModelLoader(BaseModelLoader): else: hf_folder = model_name_or_path + server_args = get_global_server_args() + if server_args and server_args.model_checksum is not None: + from sglang.srt.utils.model_file_verifier import verify + + checksums_source = server_args.model_checksum or model_name_or_path + verify(model_path=hf_folder, checksums_source=checksums_source) + hf_weights_files: List[str] = [] for pattern in allow_patterns: hf_weights_files += glob.glob(os.path.join(hf_folder, pattern)) diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index a76486a3b..ca32e98c4 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -327,6 +327,7 @@ class ServerArgs: soft_watchdog_timeout: Optional[float] = None dist_timeout: Optional[int] = None # timeout for torch.distributed download_dir: Optional[str] = None + model_checksum: Optional[str] = None base_gpu_id: int = 0 gpu_id_step: int = 1 sleep_on_idle: bool = False @@ -2963,6 +2964,14 @@ class ServerArgs: default=ServerArgs.download_dir, help="Model download directory for huggingface.", ) + parser.add_argument( + "--model-checksum", + type=str, + nargs="?", + const="", + default=None, + help="Model file integrity verification. If provided without value, uses model-path as HF repo ID. Otherwise, provide checksums JSON file path or HuggingFace repo ID.", + ) parser.add_argument( "--base-gpu-id", type=int, diff --git a/python/sglang/srt/utils/model_file_verifier.py b/python/sglang/srt/utils/model_file_verifier.py new file mode 100644 index 000000000..24fbc999d --- /dev/null +++ b/python/sglang/srt/utils/model_file_verifier.py @@ -0,0 +1,167 @@ +""" +Model File Verifier - Verify model file integrity using SHA256 checksums. + +Example command: + python -m sglang.srt.utils.model_file_verifier verify --model-path /path/to/model --model-checksum Qwen/Qwen3-0.6B +""" + +import argparse +import hashlib +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path +from typing import Dict, List, Optional, Tuple + +IGNORE_PATTERNS = [ + "checksums.json", + ".DS_Store", + "*.lock", + ".gitattributes", + "LICENSE", + "LICENSE.*", + "README.md", + "README.*", + "NOTICE", +] + + +# ======== Verify ======== + + +def verify(*, model_path: str, checksums_source: str, max_workers: int = 4) -> None: + model_path = Path(model_path).resolve() + expected = _load_checksums_from_hf(repo_id=checksums_source) + actual = _compute_checksums_from_folder( + model_path=model_path, filenames=list(expected.keys()), max_workers=max_workers + ) + _compare_checksums(expected=expected, actual=actual) + print(f"[ModelFileVerifier] All {len(expected)} files verified successfully.") + + +def _compare_checksums(*, expected: Dict[str, str], actual: Dict[str, str]) -> None: + errors = [] + for filename, expected_hash in expected.items(): + if filename not in actual: + errors.append(f"{filename}: missing") + elif actual[filename] != expected_hash: + errors.append( + f"{filename}: mismatch (expected={expected_hash[:16]}..., actual={actual[filename][:16]}...)" + ) + + if errors: + raise IntegrityError("Integrity check failed: " + "; ".join(errors)) + + +# ======== Load Checksums ======== + + +def _load_checksums_from_hf(*, repo_id: str) -> Dict[str, str]: + from huggingface_hub import HfFileSystem + + fs = HfFileSystem() + files = fs.ls(repo_id, detail=True) + + checksums = dict( + r + for r in map(lambda f: _get_filename_and_checksum_from_hf_file(fs, f), files) + if r + ) + if not checksums: + raise IntegrityError(f"No files found in HF repo {repo_id}.") + + return checksums + + +def _get_filename_and_checksum_from_hf_file(fs, file_info): + import fnmatch + + if file_info.get("type") != "file": + return None + + filename = Path(file_info.get("name", "")).name + if any(fnmatch.fnmatch(filename, pat) for pat in IGNORE_PATTERNS): + return None + + lfs_info = file_info.get("lfs") + if lfs_info and "sha256" in lfs_info: + return filename, lfs_info["sha256"] + + if "sha256" in file_info: + return filename, file_info["sha256"] + + content = fs.read_bytes(file_info.get("name", "")) + return filename, hashlib.sha256(content).hexdigest() + + +# ======== Compute Checksums ======== + + +def _compute_checksums_from_folder( + *, model_path: Path, filenames: List[str], max_workers: int +) -> Dict[str, str]: + from tqdm import tqdm + + def compute_one(filename: str) -> Tuple[str, Optional[str]]: + full_path = model_path / filename + if not full_path.exists(): + return filename, None + sha256 = compute_sha256(file_path=full_path) + return filename, sha256 + + with ThreadPoolExecutor(max_workers=max_workers) as executor: + results = list( + tqdm( + executor.map(compute_one, filenames), + total=len(filenames), + desc="Computing checksums", + ) + ) + + return {k: v for k, v in results if v is not None} + + +def compute_sha256(*, file_path) -> str: + sha256 = hashlib.sha256() + with open(file_path, "rb") as f: + while chunk := f.read(64 * 1024): + sha256.update(chunk) + return sha256.hexdigest() + + +# ======== Exceptions ======== + + +class IntegrityError(Exception): + pass + + +# ======== CLI ======== + + +def main(): + parser = argparse.ArgumentParser( + description="Model File Verifier - Verify model file integrity using checksums" + ) + parser.add_argument( + "--model-path", + required=True, + help="Local model directory", + ) + parser.add_argument( + "--model-checksum", + required=True, + help="HuggingFace repo ID for checksums", + ) + parser.add_argument( + "--workers", type=int, default=4, help="Number of parallel workers" + ) + + args = parser.parse_args() + verify( + model_path=args.model_path, + checksums_source=args.model_checksum, + max_workers=args.workers, + ) + + +if __name__ == "__main__": + main() diff --git a/test/registered/utils/test_model_file_verifier.py b/test/registered/utils/test_model_file_verifier.py new file mode 100644 index 000000000..f8ed0747a --- /dev/null +++ b/test/registered/utils/test_model_file_verifier.py @@ -0,0 +1,157 @@ +import hashlib +import os +import shutil +import tempfile +import unittest +from contextlib import nullcontext +from io import StringIO + +import requests +from huggingface_hub import snapshot_download + +from sglang.srt.utils import kill_process_tree +from sglang.srt.utils.model_file_verifier import compute_sha256, verify +from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.test_utils import ( + DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + DEFAULT_URL_FOR_TEST, + popen_launch_server, +) + +register_cuda_ci(est_time=120, suite="nightly-1-gpu", nightly=True) + +MODEL_NAME = "Qwen/Qwen3-0.6B" + + +# ======== Base Test Classes ======== + + +class _FakeModelTestCase(unittest.TestCase): + + FAKE_FILES = { + "model.safetensors": b"fake safetensors content " * 100, + "config.json": b'{"model_type": "llama"}', + "tokenizer.json": b'{"version": "1.0"}', + } + + def setUp(self): + self.test_dir = tempfile.mkdtemp() + for filename, content in self.FAKE_FILES.items(): + _create_test_file(self.test_dir, filename, content) + + def tearDown(self): + shutil.rmtree(self.test_dir, ignore_errors=True) + + +class _RealModelTestCase(unittest.TestCase): + + @classmethod + def setUpClass(cls): + cls.original_model_path = snapshot_download(MODEL_NAME) + + def setUp(self): + self.test_dir = tempfile.mkdtemp() + shutil.copytree(self.original_model_path, self.test_dir, dirs_exist_ok=True) + + def tearDown(self): + shutil.rmtree(self.test_dir, ignore_errors=True) + + +# ======== Unit Tests ======== + + +class TestModelFileVerifier(_FakeModelTestCase): + + def test_compute_sha256(self): + test_file = os.path.join(self.test_dir, "test.bin") + content = b"hello world" + with open(test_file, "wb") as f: + f.write(content) + + result = compute_sha256(file_path=test_file) + expected = hashlib.sha256(content).hexdigest() + self.assertEqual(result, expected) + + +# ======== HuggingFace Tests ======== + + +class TestModelFileVerifierHF(_RealModelTestCase): + + def test_verify_with_hf_checksums_source(self): + verify(model_path=self.test_dir, checksums_source=MODEL_NAME) + + +# ======== Real Model E2E Tests ======== + + +class TestModelFileVerifierWithRealModel(_RealModelTestCase): + + def _run_server_test(self, *, corrupt_weights: bool): + corrupted_file = None + if corrupt_weights: + safetensors_files = [ + f for f in os.listdir(self.test_dir) if f.endswith(".safetensors") + ] + self.assertTrue(len(safetensors_files) > 0, "No safetensors files found") + corrupted_file = safetensors_files[0] + _flip_bit_in_file(os.path.join(self.test_dir, corrupted_file)) + + stdout_io, stderr_io = StringIO(), StringIO() + ctx = self.assertRaises(Exception) if corrupt_weights else nullcontext() + with ctx: + process = popen_launch_server( + model=self.test_dir, + base_url=DEFAULT_URL_FOR_TEST, + timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + other_args=["--model-checksum", MODEL_NAME], + return_stdout_stderr=(stdout_io, stderr_io), + ) + + if corrupt_weights: + output = stdout_io.getvalue() + stderr_io.getvalue() + self.assertIn(corrupted_file, output) + self.assertIn("mismatch", output.lower()) + else: + try: + response = requests.post( + f"{DEFAULT_URL_FOR_TEST}/generate", + json={"text": "Hello", "sampling_params": {"max_new_tokens": 8}}, + ) + self.assertEqual(response.status_code, 200) + self.assertIn("text", response.json()) + finally: + kill_process_tree(process.pid) + + def test_server_launch_with_checksum_intact(self): + self._run_server_test(corrupt_weights=False) + + def test_server_launch_fails_with_corrupted_weights(self): + self._run_server_test(corrupt_weights=True) + + +# ======== Test Utilities ======== + + +def _create_test_file(directory: str, filename: str, content: bytes) -> str: + path = os.path.join(directory, filename) + with open(path, "wb") as f: + f.write(content) + return path + + +def _flip_bit_in_file(file_path: str, byte_offset: int = 100, bit_position: int = 0): + file_size = os.path.getsize(file_path) + assert ( + byte_offset < file_size + ), f"byte_offset {byte_offset} >= file_size {file_size}" + + with open(file_path, "r+b") as f: + f.seek(byte_offset) + original_byte = f.read(1)[0] + f.seek(byte_offset) + f.write(bytes([original_byte ^ (1 << bit_position)])) + + +if __name__ == "__main__": + unittest.main()