c781db0f6c
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> Co-authored-by: cy <chenyang08056032@163.com>
82 lines
2.2 KiB
Python
82 lines
2.2 KiB
Python
"""
|
|
Usage:
|
|
python3 -m unittest test_ascend_memory_consumption.TestMemoryConsumptionAscend.test_memory_consumption
|
|
"""
|
|
|
|
import os
|
|
import unittest
|
|
|
|
import torch
|
|
|
|
from sglang.srt.utils import kill_process_tree
|
|
from sglang.test.ci.ci_register import register_npu_ci
|
|
from sglang.test.test_utils import (
|
|
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
|
DEFAULT_URL_FOR_TEST,
|
|
CustomTestCase,
|
|
popen_launch_server,
|
|
)
|
|
|
|
register_npu_ci(
|
|
est_time=400,
|
|
suite="nightly-1-npu-a3",
|
|
nightly=True,
|
|
disabled="run failed",
|
|
)
|
|
|
|
if "ASCEND_RT_VISIBLE_DEVICES" not in os.environ:
|
|
os.environ["ASCEND_RT_VISIBLE_DEVICES"] = "0,1"
|
|
DEFAULT_PORT_FOR_SRT_TEST_RUNNER = (
|
|
8000 + int(os.environ.get("ASCEND_RT_VISIBLE_DEVICES", "0")[0]) * 100
|
|
)
|
|
DEFAULT_URL_FOR_TEST = f"http://127.0.0.1:{DEFAULT_PORT_FOR_SRT_TEST_RUNNER + 1000}"
|
|
|
|
|
|
class TestMemoryConsumptionAscend(CustomTestCase):
|
|
|
|
def test_memory_consumption(self):
|
|
|
|
model = "nytopop/Qwen3-30B-A3B.w8a8"
|
|
base_url = DEFAULT_URL_FOR_TEST
|
|
|
|
### Calculate initial used memory
|
|
free_npu_memory, total_npu_memory = torch.npu.mem_get_info()
|
|
initial_used_memory = total_npu_memory - free_npu_memory
|
|
|
|
process = popen_launch_server(
|
|
model,
|
|
base_url,
|
|
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
|
other_args=[
|
|
"--trust-remote-code",
|
|
"--device",
|
|
"npu",
|
|
"--attention-backend",
|
|
"ascend",
|
|
"--tp-size",
|
|
"2",
|
|
"--mem-fraction-static",
|
|
"0.8",
|
|
"--cuda-graph-bs",
|
|
"1",
|
|
"--max-total-tokens",
|
|
"1024",
|
|
"--disable-radix-cache",
|
|
"--disable-cuda-graph",
|
|
],
|
|
)
|
|
|
|
### Calculate initial used memory
|
|
free_npu_memory, total_npu_memory = torch.npu.mem_get_info()
|
|
used_memory_after_server_starting = (
|
|
total_npu_memory - free_npu_memory - initial_used_memory
|
|
) / (1 << 30)
|
|
self.assertLessEqual(float(used_memory_after_server_starting), 16.00)
|
|
|
|
# Clean up everything
|
|
kill_process_tree(process.pid)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|