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

Speed up update_weights_from_tensor (#2695)

Browse Source
This commit is contained in:
fzyzcjy
2025-01-02 18:05:19 +08:00
committed by GitHub
parent 148254d4db
commit 9183c23eca
6 changed files with 48 additions and 25 deletions
Download Patch File Download Diff File Expand all files Collapse all files

30
test/srt/test_update_weights_from_tensor.py
View File

@@ -1,3 +1,4 @@
import time
import unittest
import torch
@@ -6,27 +7,32 @@ import sglang as sgl
from sglang.test.test_utils import DEFAULT_SMALL_MODEL_NAME_FOR_TEST
class TestReleaseGPUOccupation(unittest.TestCase):
def test_release_and_resume_occupation(self):
class TestUpdateWeightsFromTensor(unittest.TestCase):
def test_update_weights_from_tensor(self):
engine = sgl.Engine(model_path=DEFAULT_SMALL_MODEL_NAME_FOR_TEST)
param_name = "model.layers.2.self_attn.k_proj.weight"
param_names = [f"model.layers.{i}.mlp.up_proj.weight" for i in range(6, 16)]
def _check_param(expect_values):
actual_values = torch.tensor(engine.get_weights_by_name(param_name))[0, :5]
assert torch.allclose(
actual_values, torch.tensor(expect_values), atol=0.001
), f"{actual_values=}"
_check_param(engine, param_names[0], [0.0087, -0.0214, -0.0004, 0.0039, 0.0110])
_check_param([0.0571, -0.0114, 0.0444, 0.0215, -0.0149])
new_tensor = torch.full((16384, 2048), 1.5)
new_tensor = torch.full((3072, 2048), 1.5)
engine.update_weights_from_tensor(param_name, new_tensor)
time_start = time.time()
engine.update_weights_from_tensor([(x, new_tensor) for x in param_names])
print(f"Time delta: {time.time() - time_start:.03f}")
_check_param([1.5] * 5)
for param_name in param_names[:3]:
_check_param(engine, param_name, [1.5] * 5)
engine.shutdown()
def _check_param(engine, param_name, expect_values):
actual_values = torch.tensor(engine.get_weights_by_name(param_name))[0, :5]
assert torch.allclose(
actual_values, torch.tensor(expect_values), atol=0.002
), f"{actual_values=}"
if __name__ == "__main__":
unittest.main()