EngineX-Hygon/sglang
5
0
Fork 0
Code Issues Pull Requests Actions 7 Projects Releases Wiki Activity

Add update_weights_from_disk endpoint to Engine (#4102)

Browse Source 
Co-authored-by: zhaochenyang20 <zhaochen20@outlook.com>
This commit is contained in:
Jhin
2025-03-05 14:25:18 -06:00
committed by GitHub
parent ef9d3b3c2c
commit 70b3c6eeb1
3 changed files with 239 additions and 31 deletions
Download Patch File Download Diff File Expand all files Collapse all files

22
python/sglang/srt/entrypoints/engine.py
View File

@@ -44,6 +44,7 @@ from sglang.srt.managers.io_struct import (
InitWeightsUpdateGroupReqInput,
ReleaseMemoryOccupationReqInput,
ResumeMemoryOccupationReqInput,
UpdateWeightFromDiskReqInput,
UpdateWeightsFromDistributedReqInput,
UpdateWeightsFromTensorReqInput,
)
@@ -302,6 +303,27 @@ class Engine:
self.tokenizer_manager.update_weights_from_tensor(obj, None)
)
def update_weights_from_disk(
self,
model_path: str,
load_format: Optional[str] = None,
):
"""Update the weights from disk inplace without re-launching the engine.
This method allows updating the model weights from disk without restarting
the engine. It can be used to load a different model or update weights with
new training.
"""
obj = UpdateWeightFromDiskReqInput(
model_path=model_path,
load_format=load_format,
)
loop = asyncio.get_event_loop()
return loop.run_until_complete(
self.tokenizer_manager.update_weights_from_disk(obj, None)
)
def get_weights_by_name(self, name: str, truncate_size: int = 100):
"""Get weights by parameter name."""
obj = GetWeightsByNameReqInput(name=name, truncate_size=truncate_size)

244
test/srt/test_update_weights_from_disk.py
View File

@@ -1,18 +1,76 @@
import json
import random
import unittest
import requests
import sglang as sgl
from sglang.srt.utils import kill_process_tree
from sglang.test.test_utils import (
DEFAULT_SMALL_MODEL_NAME_FOR_TEST,
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
DEFAULT_URL_FOR_TEST,
is_in_ci,
popen_launch_server,
)
class TestUpdateWeights(unittest.TestCase):
###############################################################################
# Engine Mode Tests (Single-configuration)
###############################################################################
class TestEngineUpdateWeightsFromDisk(unittest.TestCase):
def setUp(self):
self.model = DEFAULT_SMALL_MODEL_NAME_FOR_TEST
# Initialize the engine in offline (direct) mode.
self.engine = sgl.Engine(model_path=self.model)
def tearDown(self):
self.engine.shutdown()
def run_decode(self):
prompts = ["The capital of France is"]
sampling_params = {"temperature": 0, "max_new_tokens": 32}
outputs = self.engine.generate(prompts, sampling_params)
print("=" * 100)
print(
f"[Engine Mode] Prompt: {prompts[0]}\nGenerated text: {outputs[0]['text']}"
)
return outputs[0]["text"]
def run_update_weights(self, model_path):
ret = self.engine.update_weights_from_disk(model_path)
print(json.dumps(ret))
return ret
def test_update_weights(self):
origin_response = self.run_decode()
# Update weights: use new model (remove "-Instruct")
new_model_path = self.model.replace("-Instruct", "")
ret = self.run_update_weights(new_model_path)
self.assertTrue(ret[0]) # ret is a tuple; index 0 holds the success flag
updated_response = self.run_decode()
self.assertNotEqual(origin_response[:32], updated_response[:32])
# Revert back to original weights
ret = self.run_update_weights(self.model)
self.assertTrue(ret[0])
reverted_response = self.run_decode()
self.assertEqual(origin_response[:32], reverted_response[:32])
def test_update_weights_unexist_model(self):
origin_response = self.run_decode()
new_model_path = self.model.replace("-Instruct", "wrong")
ret = self.run_update_weights(new_model_path)
self.assertFalse(ret[0])
updated_response = self.run_decode()
self.assertEqual(origin_response[:32], updated_response[:32])
###############################################################################
# HTTP Server Mode Tests (Single-configuration)
###############################################################################
class TestServerUpdateWeightsFromDisk(unittest.TestCase):
@classmethod
def setUpClass(cls):
cls.model = DEFAULT_SMALL_MODEL_NAME_FOR_TEST
@@ -30,16 +88,12 @@ class TestUpdateWeights(unittest.TestCase):
self.base_url + "/generate",
json={
"text": "The capital of France is",
"sampling_params": {
"temperature": 0,
"max_new_tokens": 32,
},
"sampling_params": {"temperature": 0, "max_new_tokens": 32},
},
)
print(json.dumps(response.json()))
print("=" * 100)
text = response.json()["text"]
return text
print(f"[Server Mode] Generated text: {response.json()['text']}")
return response.json()["text"]
def get_model_info(self):
response = requests.get(self.base_url + "/get_model_info")
@@ -50,58 +104,188 @@ class TestUpdateWeights(unittest.TestCase):
def run_update_weights(self, model_path):
response = requests.post(
self.base_url + "/update_weights_from_disk",
json={
"model_path": model_path,
},
json={"model_path": model_path},
)
ret = response.json()
print(json.dumps(response.json()))
print(json.dumps(ret))
return ret
def test_update_weights(self):
origin_model_path = self.get_model_info()
print(f"origin_model_path: {origin_model_path}")
print(f"[Server Mode] origin_model_path: {origin_model_path}")
origin_response = self.run_decode()
# update weights
new_model_path = DEFAULT_SMALL_MODEL_NAME_FOR_TEST.replace("-Instruct", "")
ret = self.run_update_weights(new_model_path)
assert ret["success"]
self.assertTrue(ret["success"])
updated_model_path = self.get_model_info()
print(f"updated_model_path: {updated_model_path}")
assert updated_model_path == new_model_path
assert updated_model_path != origin_model_path
print(f"[Server Mode] updated_model_path: {updated_model_path}")
self.assertEqual(updated_model_path, new_model_path)
self.assertNotEqual(updated_model_path, origin_model_path)
updated_response = self.run_decode()
assert origin_response[:32] != updated_response[:32]
self.assertNotEqual(origin_response[:32], updated_response[:32])
# update weights back
ret = self.run_update_weights(origin_model_path)
assert ret["success"]
self.assertTrue(ret["success"])
updated_model_path = self.get_model_info()
assert updated_model_path == origin_model_path
self.assertEqual(updated_model_path, origin_model_path)
updated_response = self.run_decode()
assert origin_response[:32] == updated_response[:32]
self.assertEqual(origin_response[:32], updated_response[:32])
def test_update_weights_unexist_model(self):
origin_model_path = self.get_model_info()
print(f"origin_model_path: {origin_model_path}")
print(f"[Server Mode] origin_model_path: {origin_model_path}")
origin_response = self.run_decode()
# update weights
new_model_path = DEFAULT_SMALL_MODEL_NAME_FOR_TEST.replace("-Instruct", "wrong")
ret = self.run_update_weights(new_model_path)
assert not ret["success"]
self.assertFalse(ret["success"])
updated_model_path = self.get_model_info()
print(f"updated_model_path: {updated_model_path}")
assert updated_model_path == origin_model_path
print(f"[Server Mode] updated_model_path: {updated_model_path}")
self.assertEqual(updated_model_path, origin_model_path)
updated_response = self.run_decode()
assert origin_response[:32] == updated_response[:32]
self.assertEqual(origin_response[:32], updated_response[:32])
###############################################################################
# Parameterized Tests for update_weights_from_disk
# Test coverage is determined based on the value of is_in_ci:
# - In a CI environment: randomly select one mode (Engine or Server) and test only with tp=1, dp=1.
# - In a non-CI environment: test both Engine and Server modes, and enumerate all combinations
# with tp and dp ranging from 1 to 2.
###############################################################################
class TestUpdateWeightsFromDiskParameterized(unittest.TestCase):
def run_common_test(self, mode, tp, dp):
"""
Common test procedure for update_weights_from_disk.
For Engine mode, we instantiate the engine with tp_size=tp.
For Server mode, we launch the server with additional arguments for tp (dp is not used in server launch here).
"""
if mode == "Engine":
# Instantiate engine with additional parameter tp_size.
print(f"[Parameterized Engine] Testing with tp={tp}, dp={dp}")
engine = sgl.Engine(
model_path=DEFAULT_SMALL_MODEL_NAME_FOR_TEST,
random_seed=42,
tp_size=tp,
# dp parameter is not explicitly used in this API.
)
try:
origin_response = self._engine_update_weights_test(engine)
finally:
engine.shutdown()
elif mode == "Server":
print(f"[Parameterized Server] Testing with tp={tp}, dp={dp}")
# Pass additional arguments to launch the server.
base_args = ["--tp-size", str(tp)]
process = popen_launch_server(
DEFAULT_SMALL_MODEL_NAME_FOR_TEST,
DEFAULT_URL_FOR_TEST,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=base_args,
)
try:
origin_response = self._server_update_weights_test(DEFAULT_URL_FOR_TEST)
finally:
kill_process_tree(process.pid)
else:
raise ValueError(f"Unknown mode: {mode}")
def _engine_update_weights_test(self, engine):
# Run the update weights test on the given engine instance.
def run_decode():
prompts = ["The capital of France is"]
sampling_params = {"temperature": 0, "max_new_tokens": 32}
outputs = engine.generate(prompts, sampling_params)
print("=" * 100)
print(
f"[Parameterized Engine] Prompt: {prompts[0]}\nGenerated text: {outputs[0]['text']}"
)
return outputs[0]["text"]
def run_update_weights(model_path):
ret = engine.update_weights_from_disk(model_path)
print(json.dumps(ret))
return ret
origin_response = run_decode()
new_model_path = DEFAULT_SMALL_MODEL_NAME_FOR_TEST.replace("-Instruct", "")
ret = run_update_weights(new_model_path)
self.assertTrue(ret[0])
updated_response = run_decode()
self.assertNotEqual(origin_response[:32], updated_response[:32])
ret = run_update_weights(DEFAULT_SMALL_MODEL_NAME_FOR_TEST)
self.assertTrue(ret[0])
reverted_response = run_decode()
self.assertEqual(origin_response[:32], reverted_response[:32])
return origin_response
def _server_update_weights_test(self, base_url):
def run_decode():
response = requests.post(
base_url + "/generate",
json={
"text": "The capital of France is",
"sampling_params": {"temperature": 0, "max_new_tokens": 32},
},
)
print("=" * 100)
print(f"[Parameterized Server] Generated text: {response.json()['text']}")
return response.json()["text"]
def get_model_info():
response = requests.get(base_url + "/get_model_info")
model_path = response.json()["model_path"]
print(json.dumps(response.json()))
return model_path
def run_update_weights(model_path):
response = requests.post(
base_url + "/update_weights_from_disk",
json={"model_path": model_path},
)
ret = response.json()
print(json.dumps(ret))
return ret
origin_model_path = get_model_info()
origin_response = run_decode()
new_model_path = DEFAULT_SMALL_MODEL_NAME_FOR_TEST.replace("-Instruct", "")
ret = run_update_weights(new_model_path)
self.assertTrue(ret["success"])
updated_model_path = get_model_info()
self.assertEqual(updated_model_path, new_model_path)
self.assertNotEqual(updated_model_path, origin_model_path)
updated_response = run_decode()
self.assertNotEqual(origin_response[:32], updated_response[:32])
ret = run_update_weights(origin_model_path)
self.assertTrue(ret["success"])
updated_model_path = get_model_info()
self.assertEqual(updated_model_path, origin_model_path)
reverted_response = run_decode()
self.assertEqual(origin_response[:32], reverted_response[:32])
return origin_response
def test_parameterized_update_weights(self):
if is_in_ci():
# In CI, choose one random mode (Engine or Server) with tp=1, dp=1.
mode = random.choice(["Engine", "Server"])
test_suits = [(1, 1, mode)]
else:
# Otherwise, test both modes and enumerate tp,dp combinations from 1 to 2.
test_suits = []
for mode in ["Engine", "Server"]:
for tp in [1, 2]:
for dp in [1, 2]:
test_suits.append((tp, dp, mode))
for tp, dp, mode in test_suits:
with self.subTest(mode=mode, tp=tp, dp=dp):
self.run_common_test(mode, tp, dp)
if __name__ == "__main__":

4
test/srt/test_update_weights_from_distributed.py
View File

@@ -15,6 +15,7 @@ distributed setup.
import gc
import os
import random
import time
import unittest
@@ -529,8 +530,9 @@ class TestUpdateWeightsFromDistributed(unittest.TestCase):
assert torch.cuda.device_count() >= 2, "At least 2 GPUs are required"
# test_suits : tp, dp, model_name, backend
if is_in_ci():
mode = random.choice(["Engine", "Server"])
test_suits = [
(1, 1, DEFAULT_SMALL_MODEL_NAME_FOR_TEST, "Engine"),
(1, 1, DEFAULT_SMALL_MODEL_NAME_FOR_TEST, mode),
]
else:
test_suits = [