33 lines
925 B
Python
33 lines
925 B
Python
import unittest
|
|
|
|
import torch
|
|
|
|
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):
|
|
engine = sgl.Engine(model_path=DEFAULT_SMALL_MODEL_NAME_FOR_TEST)
|
|
|
|
param_name = "model.layers.2.self_attn.k_proj.weight"
|
|
|
|
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([0.0571, -0.0114, 0.0444, 0.0215, -0.0149])
|
|
|
|
new_tensor = torch.full((3072, 2048), 1.5)
|
|
engine.update_weights_from_tensor(param_name, new_tensor)
|
|
|
|
_check_param([1.5] * 5)
|
|
|
|
engine.shutdown()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|