# Copyright 2023-2024 SGLang Team # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== """End-to-end test for the /weights_checker HTTP endpoint. Exercises the full HTTP -> tokenizer_manager -> scheduler -> model_runner -> WeightChecker chain on a real engine. Unit tests in test/registered/unit/utils/test_weight_checker.py cover the in-module logic; this file is the thin integration cover plus interaction with update_weights_from_tensor.""" import unittest from typing import List, Tuple import requests import torch from sglang.srt.utils import MultiprocessingSerializer, kill_process_tree 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, CustomTestCase, popen_launch_server, ) register_cuda_ci(est_time=150, suite="nightly-1-gpu", nightly=True) _MODEL_NAME = "Qwen/Qwen3-0.6B" # We address the up half via the HF-style unfused name "up_proj.weight". sglang's # stacked_params_mapping rewrites this to "gate_up_proj.weight" with shard_id=1, # so the upload writes only the up half of the fused tensor. Sending the fused # name directly hits a name.replace() collision (gate_up_proj contains up_proj), # producing a malformed key like "gate_gate_up_proj.weight" and crashing load. _UP_PROJ_SHAPE = (3072, 1024) # intermediate_size, hidden_size for Qwen3-0.6B class TestWeightCheckerE2E(CustomTestCase): """All cases share one launched server (setUpClass). The reset case mutates weights to random; it is named to sort last so any case that needs intact weights runs first. The server is torn down right after, so leaving the engine in a corrupted state is harmless.""" @classmethod def setUpClass(cls): cls.url = DEFAULT_URL_FOR_TEST # --mem-fraction-static 0.7 leaves enough free GPU for _check_tensors's # CPU->GPU round trip: snapshot lives on CPU, then _compare moves each # snapshot tensor back to GPU for byte equality. With the default 0.88, # sglang holds ~29GB on a 32GB GPU and only ~200MB is free, so the # vocab-embedding round-trip (~600MB) OOMs the snapshot/reset/compare # cycle in test_z_*. cls.process = popen_launch_server( _MODEL_NAME, cls.url, timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, other_args=["--mem-fraction-static", "0.7"], ) @classmethod def tearDownClass(cls): kill_process_tree(cls.process.pid) def _post(self, action: str) -> requests.Response: return requests.post( f"{self.url}/weights_checker", json={"action": action}, timeout=120 ) def _update_weights( self, named_tensors: List[Tuple[str, torch.Tensor]] ) -> requests.Response: return requests.post( f"{self.url}/update_weights_from_tensor", json={ "serialized_named_tensors": [ MultiprocessingSerializer.serialize(named_tensors, output_str=True) ], "flush_cache": True, }, timeout=120, ) def test_a_snapshot_then_compare_unchanged_succeeds(self): resp = self._post("snapshot") self.assertEqual(resp.status_code, 200) self.assertTrue(resp.json()["success"]) resp = self._post("compare") self.assertEqual(resp.status_code, 200) self.assertTrue(resp.json()["success"]) def test_b_unknown_action_returns_400(self): resp = self._post("nonsense_action") self.assertEqual(resp.status_code, 400) self.assertIn("Unsupported", resp.json()["message"]) def test_c_update_with_diff_tensor_makes_compare_fail(self): """A snapshot then an update with new bytes must make compare fail.""" self.assertEqual(self._post("snapshot").status_code, 200) # The unfused HF name "up_proj" is what update_weights_from_tensor accepts; # sglang's loader rewrites it onto the fused gate_up_proj tensor. upload_name = "model.layers.5.mlp.up_proj.weight" new_tensor = torch.full(_UP_PROJ_SHAPE, 1.5, device="cuda") update_resp = self._update_weights([(upload_name, new_tensor)]) self.assertEqual(update_resp.status_code, 200) self.assertTrue(update_resp.json()["success"]) resp = self._post("compare") self.assertEqual(resp.status_code, 400) body = resp.json() self.assertFalse(body["success"]) # The error references the fused on-device parameter name, not the upload alias. self.assertIn("model.layers.5.mlp.gate_up_proj.weight", body["message"]) self.assertIn("max_abs_err", body["message"]) def test_d_update_with_same_tensor_keeps_compare_passing(self): """Prime a param, snapshot, push the same bytes again, compare must pass.""" param_name = "model.layers.6.mlp.up_proj.weight" same_tensor = torch.full(_UP_PROJ_SHAPE, 0.25, device="cuda") # Step 1: prime the param to a known value. self.assertTrue( self._update_weights([(param_name, same_tensor)]).json()["success"] ) # Step 2: snapshot the now-primed state. self.assertEqual(self._post("snapshot").status_code, 200) # Step 3: push the exact same bytes again — should be a byte-perfect no-op. self.assertTrue( self._update_weights([(param_name, same_tensor)]).json()["success"] ) # Step 4: compare passes. resp = self._post("compare") self.assertEqual(resp.status_code, 200) self.assertTrue(resp.json()["success"]) def test_e_checksum_returns_ranks_with_hashes(self): """checksum action must yield a ranks list with hex hashes per rank.""" resp = self._post("checksum") self.assertEqual(resp.status_code, 200) body = resp.json() self.assertTrue(body["success"]) self.assertIn("ranks", body) ranks = body["ranks"] self.assertIsInstance(ranks, list) self.assertGreaterEqual(len(ranks), 1) first = ranks[0] self.assertIn("checksums", first) self.assertIn("parallelism_info", first) info = first["parallelism_info"] for key in ( "tp_rank", "tp_size", "dp_rank", "dp_size", "pp_rank", "pp_size", "rank", "size", ): self.assertIn(key, info) checksums = first["checksums"] self.assertGreater(len(checksums), 0) for name, h in checksums.items(): self.assertIsInstance(h, str) self.assertEqual(len(h), 16, f"unexpected hash length for {name!r}: {h!r}") int(h, 16) def test_e_checksum_is_stable_across_calls(self): """Two consecutive checksum calls with no weight update must match.""" first = self._post("checksum").json()["ranks"] second = self._post("checksum").json()["ranks"] self.assertEqual(first, second) def test_e_checksum_changes_after_weight_update(self): """Updating a tensor must change its corresponding hash.""" param_name = "model.layers.7.mlp.up_proj.weight" fused_name = "model.layers.7.mlp.gate_up_proj.weight" before = self._post("checksum").json()["ranks"][0]["checksums"] before_hash = before.get(fused_name) self.assertIsNotNone(before_hash, f"missing {fused_name!r} in checksum keys") new_tensor = torch.full(_UP_PROJ_SHAPE, 0.5, device="cuda") self.assertTrue( self._update_weights([(param_name, new_tensor)]).json()["success"] ) after = self._post("checksum").json()["ranks"][0]["checksums"] self.assertNotEqual(after[fused_name], before_hash) def test_e_checksum_skips_non_persistent_buffers(self): """No checksum entry should contain a non-persistent-buffer substring.""" ranks = self._post("checksum").json()["ranks"] for rank in ranks: for name in rank["checksums"]: self.assertNotIn("cos_sin_cache", name) self.assertNotIn("inv_freq", name) self.assertNotIn("freqs_cis", name) self.assertNotIn("_weight_fp32", name) def test_z_snapshot_reset_compare_detects_diff(self): """Destructive: leaves weights randomized. Named test_z_* so it runs last.""" self.assertEqual(self._post("snapshot").status_code, 200) self.assertEqual(self._post("reset_tensors").status_code, 200) resp = self._post("compare") self.assertEqual(resp.status_code, 400) body = resp.json() self.assertFalse(body["success"]) self.assertIn("max_abs_err", body["message"]) if __name__ == "__main__": unittest.main()