chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

ci: migrate RL tests to test/registered/rl/ (#16417)

Browse Source
This commit is contained in:
Alison Shao
2026-01-05 22:13:13 -08:00
committed by GitHub
parent 6ffe1fc02f
commit 861a35fb6e
10 changed files with 76 additions and 17 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/srt/rl/test_fp32_lm_head.py
-111
View File
@@ -1,111 +0,0 @@
import unittest
from types import SimpleNamespace
from unittest.mock import patch
import torch
import torch.nn as nn
import torch.nn.functional as F
from sglang.srt.layers.logits_processor import LogitsProcessor
from sglang.srt.server_args import (
ServerArgs,
get_global_server_args,
set_global_server_args_for_scheduler,
)
class LMHeadStub(nn.Module):
def __init__(self, vocab, hidden, dtype, device="cuda"):
super().__init__()
self.weight = nn.Parameter(
torch.randn(vocab, hidden, dtype=dtype, device=device)
)
class DummyMeta:
gathered_buffer = None
next_token_logits_buffer = None
def compute_dp_attention_metadata(self): ...
class TestLMHeadFP32(unittest.TestCase):
@classmethod
def setUpClass(cls):
if not torch.cuda.is_available():
raise unittest.SkipTest("needs CUDA GPU")
def _make_logprocessor(self, vocab_size, enable_fp32):
set_global_server_args_for_scheduler(ServerArgs(model_path="dummy"))
get_global_server_args().enable_dp_lm_head = False
get_global_server_args().enable_fp32_lm_head = enable_fp32
cfg = SimpleNamespace(vocab_size=vocab_size, final_logit_softcapping=None)
return LogitsProcessor(cfg, skip_all_gather=True, logit_scale=None)
def _run_case(
self,
hidden_state_dtype,
enable_fp32,
weights_dtype,
expected_a_dtype,
expected_b_dtype,
):
device = "cuda"
BATCH_SIZE, HIDDEN_SIZE, VOCAB_SIZE = 2, 64, 128
hidden_state = torch.randn(
BATCH_SIZE, HIDDEN_SIZE, dtype=hidden_state_dtype, device=device
)
head = LMHeadStub(VOCAB_SIZE, HIDDEN_SIZE, dtype=weights_dtype, device=device)
meta = DummyMeta()
logprocessor = self._make_logprocessor(VOCAB_SIZE, enable_fp32)
original_matmul = torch.matmul
original_linear = F.linear
state = {
"called": False, # Whether a matmul/linear call has been intercepted yet
"operation": None, # Which operation was captured ("matmul" or "linear")
"a": None, # The dtype of the first input tensor to the operation
"b": None, # The dtype of the second input tensor to the operation
}
def probe_matmul(a, b, *args, **kw):
if not state["called"]:
state.update(called=True, operation="matmul", a=a.dtype, b=b.dtype)
return original_matmul(a, b, *args, **kw)
def probe_linear(x, w, bias=None):
if not state["called"]:
state.update(called=True, ooperationp="linear", a=x.dtype, b=w.dtype)
return original_linear(x, w, bias)
with patch("torch.matmul", new=probe_matmul), patch(
"torch.nn.functional.linear", new=probe_linear
):
logits = logprocessor._get_logits(hidden_state, head, meta)
self.assertEqual(hidden_state.dtype, hidden_state_dtype)
self.assertTrue(state["called"], "no call lm head matlmul/linear")
self.assertEqual(state["a"], expected_a_dtype)
self.assertEqual(state["b"], expected_b_dtype)
def test_flag_true_fp16_activations(self):
self._run_case(torch.float16, True, torch.float16, torch.float32, torch.float32)
def test_flag_true_bf16_activations(self):
self._run_case(
torch.bfloat16, True, torch.bfloat16, torch.float32, torch.float32
)
def test_flag_false_fp16_path(self):
self._run_case(
torch.float16, False, torch.float16, torch.float16, torch.float16
)
def test_flag_false_bf16_path(self):
self._run_case(
torch.bfloat16, False, torch.bfloat16, torch.bfloat16, torch.bfloat16
)
if __name__ == "__main__":
unittest.main(verbosity=2)
test/srt/rl/test_return_routed_experts.py
-187
View File
@@ -1,187 +0,0 @@
import asyncio
import logging
import unittest
from typing import List
import aiohttp
import requests
import torch
from torch.nn.utils.rnn import pad_sequence
from sglang.srt.layers.moe.routed_experts_capturer import (
extract_routed_experts_from_meta_info,
)
from sglang.srt.utils import kill_process_tree
from sglang.test.test_utils import (
DEFAULT_ENABLE_ROUTED_EXPERTS_MODEL_NAME_FOR_TEST,
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
DEFAULT_URL_FOR_TEST,
CustomTestCase,
popen_launch_server,
)
SHAREGPT_URL = (
"https://huggingface.co/datasets/anon8231489123/"
"ShareGPT_Vicuna_unfiltered/resolve/main/ShareGPT_V3_unfiltered_cleaned_split.json"
)
logger = logging.getLogger(__name__)
class TestReturnRoutedExperts(CustomTestCase):
# modified from test_hicache.py
@classmethod
def setUpClass(cls):
cls.baseline_args = [
"--enable-return-routed-experts",
"--enable-deterministic-inference",
"--disable-overlap-schedule",
"--disable-cuda-graph",
"--disable-radix-cache",
"--tp",
4,
"--dp",
4,
"--enable-dp-attention",
]
cls.reference_args = [
"--enable-return-routed-experts",
"--enable-deterministic-inference",
"--tp",
4,
"--dp",
4,
"--enable-dp-attention",
]
cls.sampling_args = {
"temperature": 0,
}
# prepare ShareGPT dataset
try:
response = requests.get(SHAREGPT_URL, timeout=60)
response.raise_for_status()
data = response.json()
print(f"Dataset size: {len(data)}")
except requests.exceptions.RequestException as e:
raise Exception(f"Failed to download ShareGPT dataset: {e}") from e
cls.texts = []
for s in data:
if "conversations" in s and len(s["conversations"]) > 0:
try:
text = s["conversations"][0]["value"]
if isinstance(text, str) and len(text) <= 2000:
cls.texts.append(text)
except (KeyError, IndexError, TypeError) as e:
print(f"Warning: Skipping invalid conversation data: {e}")
continue
if not cls.texts:
raise ValueError("No valid texts found in the dataset")
cls.texts = cls.texts[:100]
@classmethod
def test_return_routed_experts(cls):
captured_baseline_experts = asyncio.run(
cls.fetch_result("baseline", cls.baseline_args)
)
captured_reference_experts = asyncio.run(
cls.fetch_result("reference", cls.reference_args)
)
check_all_experts_id_valid(captured_baseline_experts)
check_all_experts_id_valid(captured_reference_experts)
num_baseline_topks = (
sum([len(seq) for seq in captured_baseline_experts])
* len(captured_baseline_experts[0][0])
* len(captured_baseline_experts[0][0][0])
)
num_mismatches = compare_baseline_w_reference(
captured_baseline_experts, captured_reference_experts
)
logger.info(
f"Total mismatches report: {num_mismatches} out of {num_baseline_topks} ({num_mismatches/num_baseline_topks:.4%})"
)
print(
f"Total mismatches report: {num_mismatches} out of {num_baseline_topks} ({num_mismatches/num_baseline_topks:.4%})"
)
assert (
num_mismatches / num_baseline_topks < 0.05
), f"Too many mismatches: {num_mismatches} out of {num_baseline_topks} ({num_mismatches/num_baseline_topks:.4%})"
@classmethod
async def fetch_result(cls, title, other_args):
try:
process = popen_launch_server(
DEFAULT_ENABLE_ROUTED_EXPERTS_MODEL_NAME_FOR_TEST,
DEFAULT_URL_FOR_TEST,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=other_args,
)
async with aiohttp.ClientSession() as session:
tasks = [
asyncio.create_task(
make_request(
session,
f"{DEFAULT_URL_FOR_TEST}/generate",
{
"text": text,
"sampling_params": cls.sampling_args,
"return_routed_experts": True,
"max_new_tokens": 100,
},
)
)
for text in cls.texts
]
# return value shape: List[[seq_len, num_layers, topk]...]
http_result = await asyncio.gather(*tasks)
except Exception as e:
raise e
finally:
kill_process_tree(process.pid)
result = [
extract_routed_experts_from_meta_info(res).reshape(-1, 48, 8)
for res in http_result
]
return result
async def make_request(session, url, payload):
"""Make a single async HTTP request"""
async with session.post(url=url, json=payload) as response:
return await response.json()
def check_all_experts_id_valid(experts: List[List[List[int]]]):
tensor_list = [torch.tensor(lst) for lst in experts]
padded_tensor = pad_sequence(tensor_list, batch_first=True, padding_value=0)
# temporary hardcode as we only use Qwen3 30BA3B
if not ((padded_tensor >= 0) & (padded_tensor <= 127)).all():
raise ValueError(
f"Some expert indices are out of valid range [0, 127], MAX: {padded_tensor.max()} MIN: {padded_tensor.min()}"
)
def compare_baseline_w_reference(baseline, reference):
num_total_mismatches = 0
for baseline_seq, reference_seq in zip(baseline, reference):
for bsl_token, ref_token in zip(baseline_seq, reference_seq):
for bsl_topk, ref_topk in zip(bsl_token, ref_token):
len_bsl, len_ref = len(bsl_topk), len(ref_topk)
set_bsl, set_ref = set(bsl_topk), set(ref_topk)
if set_bsl != set_ref:
num_total_mismatches += len(set_bsl - set_ref)
if (len_bsl != len_ref) or (len_bsl != len(set_bsl)):
raise ValueError(
f"Duplicates experts ids found: Baseline({len_bsl}): {bsl_topk} vs Reference({len_ref}): {ref_topk}"
)
return num_total_mismatches
if __name__ == "__main__":
unittest.main()
test/srt/rl/test_update_weights_from_disk.py
-439
View File
@@ -1,439 +0,0 @@
import json
import random
import time
import unittest
from concurrent.futures import ThreadPoolExecutor, as_completed
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,
CustomTestCase,
is_in_ci,
popen_launch_server,
)
###############################################################################
# Engine Mode Tests (Single-configuration)
###############################################################################
class TestEngineUpdateWeightsFromDisk(CustomTestCase):
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(CustomTestCase):
@classmethod
def setUpClass(cls):
cls.model = DEFAULT_SMALL_MODEL_NAME_FOR_TEST
cls.base_url = DEFAULT_URL_FOR_TEST
cls.process = popen_launch_server(
cls.model, cls.base_url, timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH
)
@classmethod
def tearDownClass(cls):
kill_process_tree(cls.process.pid)
def run_decode(self):
response = requests.post(
self.base_url + "/generate",
json={
"text": "The capital of France is",
"sampling_params": {"temperature": 0, "max_new_tokens": 32},
},
)
print("=" * 100)
print(f"[Server Mode] Generated text: {response.json()['text']}")
return response.json()["text"]
def run_decode_random(self, max_new_tokens=32):
response = requests.post(
self.base_url + "/generate",
json={
"text": f"Question: {random.randint(0, 100)},The capital of France is",
"sampling_params": {
"temperature": 0,
"max_new_tokens": max_new_tokens,
"ignore_eos": True,
},
},
)
return response.json()
def get_model_info(self):
response = requests.get(self.base_url + "/get_model_info")
model_path = response.json()["model_path"]
print(json.dumps(response.json()))
return model_path
def run_update_weights(self, model_path, flush_cache=True):
response = requests.post(
self.base_url + "/update_weights_from_disk",
json={
"model_path": model_path,
"flush_cache": flush_cache,
},
)
ret = response.json()
return ret
def pause_generation(self, mode):
response = requests.post(
self.base_url + "/pause_generation",
json={"mode": mode},
)
ret = response.json()
return ret
def continue_generation(self):
response = requests.post(
self.base_url + "/continue_generation",
json={},
)
ret = response.json()
return ret
def test_update_weights(self):
origin_model_path = self.get_model_info()
print(f"[Server Mode] origin_model_path: {origin_model_path}")
origin_response = self.run_decode()
new_model_path = DEFAULT_SMALL_MODEL_NAME_FOR_TEST.replace("-Instruct", "")
ret = self.run_update_weights(new_model_path)
self.assertTrue(ret["success"])
updated_model_path = self.get_model_info()
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()
self.assertNotEqual(origin_response[:32], updated_response[:32])
ret = self.run_update_weights(origin_model_path)
self.assertTrue(ret["success"])
updated_model_path = self.get_model_info()
self.assertEqual(updated_model_path, origin_model_path)
updated_response = self.run_decode()
self.assertEqual(origin_response[:32], updated_response[:32])
def test_update_weights_non_blocking(self):
origin_model_path = self.get_model_info()
print(f"[Server Mode] origin_model_path: {origin_model_path}")
pause_generation_modes = ["in_place", "retract"]
for pause_generation_mode in pause_generation_modes:
num_requests = 32
with ThreadPoolExecutor(num_requests) as executor:
futures = [
executor.submit(self.run_decode_random, 1600)
for _ in range(num_requests)
]
# ensure the decode has been started
time.sleep(2)
new_model_path = DEFAULT_SMALL_MODEL_NAME_FOR_TEST.replace(
"-Instruct", ""
)
ret = self.pause_generation(pause_generation_mode)
ret = self.run_update_weights(
new_model_path, flush_cache=pause_generation_mode == "retract"
)
self.assertTrue(ret["success"])
ret = self.continue_generation()
for future in as_completed(futures):
self.assertNotEqual(
future.result()["meta_info"]["finish_reason"]["type"], "abort"
)
updated_model_path = self.get_model_info()
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)
def test_update_weights_unexist_model(self):
origin_model_path = self.get_model_info()
print(f"[Server Mode] origin_model_path: {origin_model_path}")
origin_response = self.run_decode()
new_model_path = DEFAULT_SMALL_MODEL_NAME_FOR_TEST.replace("-Instruct", "wrong")
ret = self.run_update_weights(new_model_path)
self.assertFalse(ret["success"])
updated_model_path = self.get_model_info()
print(f"[Server Mode] updated_model_path: {updated_model_path}")
self.assertEqual(updated_model_path, origin_model_path)
updated_response = self.run_decode()
self.assertEqual(origin_response[:32], updated_response[:32])
class TestServerUpdateWeightsFromDiskAbortAllRequests(CustomTestCase):
@classmethod
def setUpClass(cls):
cls.model = DEFAULT_SMALL_MODEL_NAME_FOR_TEST
cls.base_url = DEFAULT_URL_FOR_TEST
cls.process = popen_launch_server(
cls.model,
cls.base_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=["--max-running-requests", 8],
)
@classmethod
def tearDownClass(cls):
kill_process_tree(cls.process.pid)
def run_decode(self, max_new_tokens=32):
response = requests.post(
self.base_url + "/generate",
json={
"text": "The capital of France is",
"sampling_params": {
"temperature": 0,
"max_new_tokens": max_new_tokens,
"ignore_eos": True,
},
},
)
return response.json()
def get_model_info(self):
response = requests.get(self.base_url + "/get_model_info")
model_path = response.json()["model_path"]
print(json.dumps(response.json()))
return model_path
def run_update_weights(self, model_path, abort_all_requests=False):
response = requests.post(
self.base_url + "/update_weights_from_disk",
json={
"model_path": model_path,
"abort_all_requests": abort_all_requests,
},
)
ret = response.json()
print(json.dumps(ret))
return ret
def test_update_weights_abort_all_requests(self):
origin_model_path = self.get_model_info()
print(f"[Server Mode] origin_model_path: {origin_model_path}")
num_requests = 32
with ThreadPoolExecutor(num_requests) as executor:
futures = [
executor.submit(self.run_decode, 16000) for _ in range(num_requests)
]
# ensure the decode has been started
time.sleep(2)
new_model_path = DEFAULT_SMALL_MODEL_NAME_FOR_TEST.replace("-Instruct", "")
ret = self.run_update_weights(new_model_path, abort_all_requests=True)
self.assertTrue(ret["success"])
for future in as_completed(futures):
self.assertEqual(
future.result()["meta_info"]["finish_reason"]["type"], "abort"
)
updated_model_path = self.get_model_info()
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)
###############################################################################
# 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(CustomTestCase):
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__":
unittest.main()
test/srt/rl/test_update_weights_from_distributed.py
-765
View File
@@ -1,765 +0,0 @@
"""Test distributed weight updates.
This test suite simulates a distributed training environment to ensure
correct weight synchronization. On rank 0, the instruct model represents
pre-training weights, and the base model represents post-training weights.
The base model's weights are broadcasted to other ranks using the online
weight update API.
On other ranks, an engine is initialized with the instruct model, and its
parameters are verified against the Hugging Face model. After updating
weights from the distributed system, post-training weights are loaded
and verified again to ensure consistency and accuracy across the
distributed setup.
"""
import gc
import os
import random
import time
import unittest
from concurrent.futures import ThreadPoolExecutor
import numpy as np
import requests
import torch
import torch.multiprocessing as mp
from transformers import AutoModelForCausalLM
import sglang as sgl
from sglang.srt.utils import init_custom_process_group
from sglang.srt.weight_sync.tensor_bucket import FlattenedTensorBucket
from sglang.test.test_utils import (
DEFAULT_MODEL_NAME_FOR_TEST,
DEFAULT_SMALL_MODEL_NAME_FOR_TEST,
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
DEFAULT_URL_FOR_TEST,
CustomTestCase,
is_in_ci,
popen_launch_server,
)
from sglang.utils import terminate_process
mp.set_start_method("spawn", force=True)
def verify_params_close(params1, params2, error_msg):
"""Verify if two parameter arrays are close enough."""
try:
assert np.allclose(np.array(params1), np.array(params2)), error_msg
except Exception as e:
print(f"Parameters not close for {error_msg}")
print("Params1:", np.array(params1))
print("Params2:", np.array(params2))
raise e
def verify_params_not_close(params1, params2, error_msg):
"""Verify if two parameter arrays are different enough."""
assert not np.allclose(np.array(params1), np.array(params2)), error_msg
def init_process(
rank,
world_size,
param_queue,
truncate_size,
state_dict_key_to_shape,
tp_size,
model_name,
backend,
checking_parameters,
tie_word_embeddings,
load_format,
barrier,
pause_generation_mode,
):
torch.cuda.set_device(rank)
if rank == 0:
init_process_hf(
rank,
world_size,
param_queue,
truncate_size,
model_name,
checking_parameters,
tie_word_embeddings,
state_dict_key_to_shape,
load_format,
barrier,
)
elif rank in [1, 2]:
init_process_sgl(
rank,
world_size,
param_queue,
truncate_size,
model_name,
checking_parameters,
tie_word_embeddings,
state_dict_key_to_shape,
backend,
tp_size,
load_format,
barrier,
pause_generation_mode,
)
def init_process_hf(
rank,
world_size,
param_queue,
truncate_size,
model_name,
checking_parameters,
tie_word_embeddings,
state_dict_key_to_shape,
load_format,
barrier,
):
# These two environment variables are very important
# to avoid unexpected behaviors of CUDA and NCCL.
os.environ["NCCL_CUMEM_ENABLE"] = "0"
os.environ["NCCL_NVLS_ENABLE"] = "0"
# Load model and get parameters
hf_instruct_model = AutoModelForCausalLM.from_pretrained(
model_name,
torch_dtype="bfloat16",
tie_word_embeddings=tie_word_embeddings,
).to("cuda:0")
base_model_name = model_name.replace("-Instruct", "")
hf_base_model = AutoModelForCausalLM.from_pretrained(
base_model_name,
torch_dtype="bfloat16",
tie_word_embeddings=tie_word_embeddings,
).to("cuda:0")
hf_instruct_params = []
hf_base_params = []
print("[hf] get parameter in hf instruct model and base model")
for parameter_name in checking_parameters:
hf_instruct_params.append(
hf_instruct_model.get_parameter(parameter_name)[:truncate_size]
.cpu()
.detach()
.float()
.numpy()
.tolist()
)
hf_base_params.append(
hf_base_model.get_parameter(parameter_name)[:truncate_size]
.cpu()
.detach()
.float()
.numpy()
.tolist()
)
param_queue.put(("hf_instruct_params", hf_instruct_params))
param_queue.put(("hf_base_params", hf_base_params))
# Init weight update group for rank 0 (the training engine in RLHF).
port = 60000 + int(os.environ.get("CUDA_VISIBLE_DEVICES", "0")[0]) * 100
init_method = f"tcp://localhost:{port}"
print(f"[hf] {rank=} {world_size=} init custom process group. {init_method=}")
group = init_custom_process_group(
backend="nccl",
init_method=init_method,
world_size=world_size,
rank=rank,
group_name="test_parameter_update_group",
)
torch.cuda.synchronize()
barrier.wait()
time_begin_broadcast = time.perf_counter()
# The last parameter is lm_head.weight, which is tied
# with embed_tokens.weight. Actually, we only need
# to broadcast embed_tokens.weight once.
broadcast_parameters = list(state_dict_key_to_shape.keys())
if tie_word_embeddings:
broadcast_parameters.remove("lm_head.weight")
if load_format == "flattened_bucket":
named_tensors = [
(parameter_name, hf_base_model.get_parameter(parameter_name))
for parameter_name in broadcast_parameters
]
bucket = FlattenedTensorBucket(named_tensors=named_tensors)
flattened_tensor = bucket.get_flattened_tensor()
torch.distributed.broadcast(flattened_tensor, src=0, group=group)
else:
# Broadcast all the weights from the training
# engine to other ranks (inference engine).
for parameter_name in broadcast_parameters:
torch.distributed.broadcast(
hf_base_model.get_parameter(parameter_name),
src=0,
group=group,
)
torch.cuda.synchronize()
time_end_broadcast = time.perf_counter()
# Measure the latency of broadcasting/weights update.
broadcast_time = time_end_broadcast - time_begin_broadcast
print(f"[hf] {rank=} {broadcast_time=:.3f}s")
param_queue.put(("broadcast_time", broadcast_time))
# Destroy process group and release related resource
torch.distributed.destroy_process_group(group)
# Delete the huggingface models to free up memory.
del hf_instruct_model
del hf_base_model
gc.collect()
torch.cuda.empty_cache()
def init_process_sgl(
rank,
world_size,
param_queue,
truncate_size,
model_name,
checking_parameters,
tie_word_embeddings,
state_dict_key_to_shape,
backend,
tp_size,
load_format,
barrier,
pause_generation_mode,
):
torch.cuda.set_device(rank)
torch.cuda.synchronize()
base_gpu_id = 1 if rank == 1 else 1 + tp_size
if backend == "Engine":
print(f"[sgl] rank {rank} init engine")
engine = sgl.Engine(
model_path=model_name,
base_gpu_id=base_gpu_id,
tp_size=tp_size,
cuda_graph_max_bs=2,
)
else:
if rank == 1:
url = DEFAULT_URL_FOR_TEST
else:
host, _, port = DEFAULT_URL_FOR_TEST.rpartition(":")
url = ":".join([host, str(int(port) + 10000)])
print(f"[sgl] rank {rank} init server on url: {url}")
process = popen_launch_server(
model_name,
url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=(
"--base-gpu-id",
str(base_gpu_id),
"--tp-size",
str(tp_size),
"--cuda-graph-max-bs",
2,
),
)
torch.cuda.synchronize()
# Get weights of instruct model, i.e. pre-training weights.
instruct_params = []
for parameter_name in checking_parameters:
instruct_params.append(
engine.get_weights_by_name(parameter_name, truncate_size)
if backend == "Engine"
else requests.get(
f"{url}/get_weights_by_name",
json={"name": parameter_name, "truncate_size": truncate_size},
).json()
)
param_queue.put((f"sgl_dp_{rank}_instruct_params", instruct_params))
port = 60000 + int(os.environ.get("CUDA_VISIBLE_DEVICES", "0")[0]) * 100
# Init weight update group with the training engine.
if backend == "Engine":
engine.init_weights_update_group(
master_address="localhost",
master_port=str(port),
rank_offset=base_gpu_id,
world_size=world_size,
group_name="test_parameter_update_group",
backend="nccl",
)
else:
requests.post(
f"{url}/init_weights_update_group",
json={
"master_address": "localhost",
"master_port": str(port),
"rank_offset": base_gpu_id,
"world_size": world_size,
"group_name": "test_parameter_update_group",
"backend": "nccl",
},
)
if pause_generation_mode in ["in_place", "retract"]:
def run_decode(max_new_tokens=32):
response = requests.post(
url + "/generate",
json={
"text": f"Question: {random.randint(0, 100)},The capital of France is",
"sampling_params": {
"temperature": 0,
"max_new_tokens": max_new_tokens,
"ignore_eos": True,
},
},
)
return response.json()
with ThreadPoolExecutor(32) as executor:
futures = [executor.submit(run_decode, 1000) for _ in range(32)]
time.sleep(2)
# The last parameter is lm_head.weight, which is tied
# with embed_tokens.weight. Actually, we only need
# to update embed_tokens.weight once.
tie_word_embeddings = (
True if model_name == DEFAULT_SMALL_MODEL_NAME_FOR_TEST else False
)
update_parameters = list(state_dict_key_to_shape.keys())
if tie_word_embeddings:
update_parameters.remove("lm_head.weight")
# Get weights from the training engine and update the inference engine.
names = [parameter_name for parameter_name in update_parameters]
dtypes = [torch.bfloat16 if backend == "Engine" else "bfloat16"] * len(names)
shapes = [state_dict_key_to_shape[parameter_name] for parameter_name in names]
if pause_generation_mode in ["in_place", "retract"]:
requests.post(
url + "/pause_generation",
json={"mode": pause_generation_mode},
)
torch.cuda.synchronize()
barrier.wait()
time_begin_update = time.perf_counter()
if backend == "Engine":
engine.update_weights_from_distributed(
names,
dtypes=dtypes,
shapes=shapes,
group_name="test_parameter_update_group",
load_format=load_format,
)
else:
requests.post(
f"{url}/update_weights_from_distributed",
json={
"names": names,
"dtypes": dtypes,
"shapes": shapes,
"group_name": "test_parameter_update_group",
"load_format": load_format,
"flush_cache": not (pause_generation_mode == "in_place"),
},
)
torch.cuda.synchronize()
time_end_update = time.perf_counter()
if pause_generation_mode in ["in_place", "retract"]:
requests.post(
url + "/continue_generation",
json={},
)
# discard unfinished requests to save test overhead
time.sleep(2)
requests.post(
url + "/pause_generation",
json={"mode": "abort"},
)
# Measure the latency of broadcast/weights update.
update_time = time_end_update - time_begin_update
print(
f"[sgl] fully update model_name {model_name} rank {rank} parameter from distributed time: {update_time:.3f}s"
)
param_queue.put((f"update_sgl_dp_{rank}_time", update_time))
# Get the weights of post-training model after weights update for correctness check.
base_params = []
for parameter_name in checking_parameters:
if backend == "Engine":
base_params.append(
engine.get_weights_by_name(parameter_name, truncate_size)
)
else:
base_params.append(
requests.get(
f"{url}/get_weights_by_name",
json={
"name": parameter_name,
"truncate_size": truncate_size,
},
).json()
)
param_queue.put((f"sgl_dp_{rank}_base_params", base_params))
if backend == "Engine":
success, _ = engine.destroy_weights_update_group(
group_name="test_parameter_update_group",
)
assert success is True
else:
response = requests.post(
f"{url}/destroy_weights_update_group",
json={
"group_name": "test_parameter_update_group",
},
)
assert response.status_code == 200
# Shutdown the engine or terminate the server process.
if backend == "Engine":
engine.shutdown()
else:
terminate_process(process)
def assert_tied_weights(params_list, message, should_be_tied):
for params in params_list:
if should_be_tied:
assert np.allclose(params[0], params[-1]), message
else:
assert not np.allclose(params[0], params[-1]), message
def test_update_weights_from_distributed(
tp_size,
dp_size,
model_name,
backend,
state_dict_key_to_shape,
truncate_size,
checking_parameters,
load_format=None,
pause_generation_mode=None,
):
tie_word_embeddings = (
True if model_name == DEFAULT_SMALL_MODEL_NAME_FOR_TEST else False
)
print(
f"Testing model: {model_name} tp_size: {tp_size}, dp_size: {dp_size} backend: {backend}"
)
param_queue = mp.Queue()
results = {}
barrier = mp.Barrier(1 + dp_size)
context = mp.spawn(
init_process,
args=(
1 + tp_size * dp_size,
param_queue,
truncate_size,
state_dict_key_to_shape,
tp_size,
model_name,
backend,
checking_parameters,
tie_word_embeddings,
load_format,
barrier,
pause_generation_mode,
),
nprocs=1 + dp_size,
join=False,
)
while len(results) < 3 * (1 + dp_size):
try:
key, value = param_queue.get(timeout=5)
results[key] = value
except Exception as e:
if all(not p.is_alive() for p in context.processes):
break
context.join()
if len(results) != 3 * (1 + dp_size):
raise RuntimeError(
f"Expected {3 * (1 + dp_size)} parameters but got {len(results)}"
)
params = {
"hf_instruct": results.get("hf_instruct_params"),
"hf_base": results.get("hf_base_params"),
"sgl_dp_1_instruct": results.get("sgl_dp_1_instruct_params"),
"sgl_dp_1_base": results.get("sgl_dp_1_base_params"),
"broadcast_time": results.get("broadcast_time"),
"update_sgl_dp_1_time": results.get("update_sgl_dp_1_time"),
}
if dp_size == 2:
dp2_params = {
"sgl_dp_2_instruct": results.get("sgl_dp_2_instruct_params"),
"sgl_dp_2_base": results.get("sgl_dp_2_base_params"),
"update_sgl_dp_2_time": results.get("update_sgl_dp_2_time"),
}
assert all(v is not None for v in dp2_params.values())
params.update(dp2_params)
# Check the correctness of weights update by verifying
# the weights of instruct model and base model.
for i in range(len(params["hf_instruct"])):
verify_params_close(
params["hf_instruct"][i],
params["sgl_dp_1_instruct"][i],
f"sgl_dp_1_instruct_params rank {i}",
)
verify_params_close(
params["hf_base"][i],
params["sgl_dp_1_base"][i],
f"sgl_dp_1_base_params rank {i}",
)
verify_params_not_close(
params["hf_instruct"][i],
params["hf_base"][i],
f"hf_instruct_params rank {i}",
)
if dp_size == 2:
verify_params_close(
params["hf_base"][i],
params["sgl_dp_2_base"][i],
f"sgl_dp_2_base_params rank {i}",
)
verify_params_close(
params["hf_instruct"][i],
params["sgl_dp_2_instruct"][i],
f"sgl_dp_2_instruct_params rank {i}",
)
assert len(params["hf_instruct"]) == len(
params["hf_base"]
), "hf_instruct_params and hf_base_params have different lengths"
# Check if the weights of lm_head are tied with embed_tokens.
params_to_check = [
(
params["hf_instruct"],
"lm_head.weight is not tied with embed_tokens.weight",
),
(
params["hf_base"],
"lm_head.weight is not tied with embed_tokens.weight",
),
(
params["sgl_dp_1_instruct"],
"lm_head.weight is not tied with embed_tokens.weight",
),
(
params["sgl_dp_1_base"],
"lm_head.weight is not tied with embed_tokens.weight",
),
]
if dp_size == 2:
params_to_check.extend(
[
(
params["sgl_dp_2_instruct"],
"lm_head.weight is not tied with embed_tokens.weight",
),
(
params["sgl_dp_2_base"],
"lm_head.weight is not tied with embed_tokens.weight",
),
]
)
assert_tied_weights(
[params for params, _ in params_to_check],
(
"lm_head.weight is not tied with embed_tokens.weight"
if tie_word_embeddings
else "lm_head.weight is tied with embed_tokens.weight"
),
tie_word_embeddings,
)
# Time limit for broadcast and update on CI is 3 / 6
# On local H100, it's 1 / 2
time_limit = 3 if model_name == DEFAULT_SMALL_MODEL_NAME_FOR_TEST else 6
assert (
params["broadcast_time"] < time_limit
), f"broadcast_time exceeds time limit {time_limit}s"
assert (
params["update_sgl_dp_1_time"] < time_limit
), f"update_sgl_dp_one_time exceeds time limit {time_limit}s"
if dp_size == 2:
assert (
params["update_sgl_dp_2_time"] < time_limit
), f"update_sgl_dp_two_time exceeds time limit {time_limit}s"
# Delete the context and close the parameter queue.
del context
param_queue.close()
param_queue.join_thread()
gc.collect()
torch.cuda.empty_cache()
class TestUpdateWeightsFromDistributed(CustomTestCase):
def test_update_weights_from_distributed(self):
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"])
if mode == "Server":
pause_generation_mode = random.choice(["in_place", "retract"])
else:
pause_generation_mode = None
load_format = random.choice(["flattened_bucket", None])
test_suits = [
(
1,
1,
DEFAULT_SMALL_MODEL_NAME_FOR_TEST,
mode,
pause_generation_mode,
load_format,
),
]
else:
test_suits = [
(
1,
1,
DEFAULT_SMALL_MODEL_NAME_FOR_TEST,
"Engine",
None,
random.choice(["flattened_bucket", None]),
),
(
1,
1,
DEFAULT_MODEL_NAME_FOR_TEST,
"Sever",
random.choice(["in_place", "retract"]),
random.choice(["flattened_bucket", None]),
),
]
if torch.cuda.device_count() >= 4:
test_suits.extend(
[
(
2,
1,
DEFAULT_SMALL_MODEL_NAME_FOR_TEST,
"Engine",
None,
random.choice(["flattened_bucket", None]),
),
(
1,
2,
DEFAULT_MODEL_NAME_FOR_TEST,
"Server",
random.choice(["in_place", "retract"]),
random.choice(["flattened_bucket", None]),
),
]
)
if torch.cuda.device_count() >= 5:
test_suits.extend(
[
(
2,
2,
DEFAULT_SMALL_MODEL_NAME_FOR_TEST,
"Engine",
None,
random.choice(["flattened_bucket", None]),
),
(
2,
2,
DEFAULT_MODEL_NAME_FOR_TEST,
"Server",
random.choice(["in_place", "retract"]),
random.choice(["flattened_bucket", None]),
),
]
)
model_state_dict_shapes = {}
test_models = [test_suit[2] for test_suit in test_suits]
for model_name in test_models:
model = AutoModelForCausalLM.from_pretrained(
model_name, torch_dtype="bfloat16"
).to("cuda:0")
state_dict = model.state_dict()
state_dict_keys = list(state_dict.keys())
model_state_dict_shapes[model_name] = {
key: state_dict[key].shape for key in state_dict_keys
}
del model
gc.collect()
torch.cuda.empty_cache()
truncate_size = 10
checking_parameters = [
"model.embed_tokens.weight",
"model.layers.0.input_layernorm.weight",
"model.layers.1.self_attn.q_proj.weight",
"model.layers.2.self_attn.k_proj.weight",
"model.layers.3.self_attn.v_proj.weight",
"model.layers.4.self_attn.o_proj.weight",
"model.layers.5.mlp.gate_proj.weight",
"model.layers.6.mlp.up_proj.weight",
"model.layers.7.mlp.down_proj.weight",
"model.layers.8.post_attention_layernorm.weight",
"model.norm.weight",
"lm_head.weight",
]
for (
tp_size,
dp_size,
model_name,
backend,
pause_generation_mode,
load_format,
) in test_suits:
test_update_weights_from_distributed(
tp_size,
dp_size,
model_name,
backend,
model_state_dict_shapes[model_name],
truncate_size,
checking_parameters,
load_format,
pause_generation_mode,
)
if __name__ == "__main__":
unittest.main()
test/srt/rl/test_update_weights_from_tensor.py
-295
View File
@@ -1,295 +0,0 @@
import gc
import json
import random
import time
import unittest
from concurrent.futures import ThreadPoolExecutor, as_completed
import requests
import torch
import sglang as sgl
from sglang.srt.utils import MultiprocessingSerializer, kill_process_tree
from sglang.srt.weight_sync.tensor_bucket import FlattenedTensorBucket
from sglang.test.test_utils import (
DEFAULT_SMALL_MODEL_NAME_FOR_TEST,
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
DEFAULT_URL_FOR_TEST,
CustomTestCase,
popen_launch_server,
)
def test_update_weights_from_tensor(tp_size):
assert torch.cuda.device_count() >= tp_size, f"At least {tp_size} GPUs are required"
torch.cuda.empty_cache()
engine = sgl.Engine(model_path=DEFAULT_SMALL_MODEL_NAME_FOR_TEST, tp_size=tp_size)
param_names = [f"model.layers.{i}.mlp.up_proj.weight" for i in range(6, 16)]
_check_param(engine, param_names[0], [0.0087, -0.0214, -0.0004, 0.0039, 0.0110])
memory_before = torch.cuda.memory_allocated()
new_tensor = torch.full((16384, 2048), 1.5, device="cuda")
time_start = time.perf_counter()
engine.update_weights_from_tensor([(x, new_tensor) for x in param_names])
print(f"Time delta: {time.perf_counter() - time_start:.03f}")
for param_name in param_names[:3]:
_check_param(engine, param_name, [1.5] * 5)
engine.shutdown()
del new_tensor
gc.collect()
torch.cuda.ipc_collect()
torch.cuda.empty_cache()
memory_after = torch.cuda.memory_allocated()
assert (
memory_after <= memory_before + 1024
), f"Memory leak detected: {memory_after - memory_before} bytes"
class TestUpdateWeightsFromTensor(CustomTestCase):
def test_update_weights_from_tensor(self):
tp_sizes = [1, 2]
for tp_size in tp_sizes:
if torch.cuda.device_count() < tp_size:
continue
with self.subTest(tp_size=tp_size):
test_update_weights_from_tensor(tp_size)
def test_update_weights_from_tensor_load_format_direct(self):
engine = sgl.Engine(model_path=DEFAULT_SMALL_MODEL_NAME_FOR_TEST)
write_param_names = [
f"model.layers.{i}.self_attn.qkv_proj.weight" for i in range(6, 16)
]
read_param_names = [
f"model.layers.{i}.self_attn.k_proj.weight" for i in range(6, 16)
]
_check_param(
engine, read_param_names[0], [-0.0198, 0.0227, 0.0168, 0.0232, -0.0178]
)
new_tensor = torch.full((3072, 2048), 1.5)
engine.update_weights_from_tensor(
[
(write_param_name, new_tensor.clone())
for write_param_name in write_param_names
],
load_format="direct",
)
for read_param_name in read_param_names[:3]:
_check_param(engine, read_param_name, [1.5] * 5)
engine.shutdown()
def test_update_weights_from_tensor_load_format_custom(self):
custom_loader_name = (
"sglang.srt.model_executor.model_runner._model_load_weights_direct"
)
engine = sgl.Engine(
model_path=DEFAULT_SMALL_MODEL_NAME_FOR_TEST,
custom_weight_loader=[custom_loader_name],
)
write_param_names = [
f"model.layers.{i}.self_attn.qkv_proj.weight" for i in range(6, 16)
]
read_param_names = [
f"model.layers.{i}.self_attn.k_proj.weight" for i in range(6, 16)
]
_check_param(
engine, read_param_names[0], [-0.0198, 0.0227, 0.0168, 0.0232, -0.0178]
)
new_tensor = torch.full((3072, 2048), 1.5)
engine.update_weights_from_tensor(
[
(write_param_name, new_tensor.clone())
for write_param_name in write_param_names
],
load_format=custom_loader_name,
)
for read_param_name in read_param_names[:3]:
_check_param(engine, read_param_name, [1.5] * 5)
engine.shutdown()
def test_update_weights_from_tensor_load_format_flattened_bucket(self):
"""Test updating weights using flattened_bucket format"""
engine = sgl.Engine(model_path=DEFAULT_SMALL_MODEL_NAME_FOR_TEST)
# Create a small set of parameters for testing
param_names = [f"model.layers.{i}.mlp.up_proj.weight" for i in range(6, 10)]
# Check original values
_check_param(engine, param_names[0], [0.0087, -0.0214, -0.0004, 0.0039, 0.0110])
# Create new tensors with different values
new_tensors = []
for _, name in enumerate(param_names):
# Create tensors with different values for each parameter
value = 2.0 # Different value for each parameter
new_tensor = torch.full((16384, 2048), value, device="cuda")
new_tensors.append((name, new_tensor))
# Create a flattened bucket
flattened_bucket = FlattenedTensorBucket(named_tensors=new_tensors)
# Extract the flattened tensor and metadata in the format expected by model_runner
flattened_tensor = flattened_bucket.get_flattened_tensor()
metadata = flattened_bucket.get_metadata()
# Create the dict format expected by _update_weights_from_flattened_bucket
bucket_dict = {"flattened_tensor": flattened_tensor, "metadata": metadata}
# Serialize the bucket data
from sglang.srt.utils import MultiprocessingSerializer
serialized_bucket = MultiprocessingSerializer.serialize(
bucket_dict, output_str=True
)
# Create a list where each rank contains the same serialized data
# This simulates the distributed environment where each rank has the same data
serialized_bucket_list = [serialized_bucket]
# Update weights using flattened_bucket format
time_start = time.perf_counter()
engine.update_weights_from_tensor(
named_tensors=serialized_bucket_list, load_format="flattened_bucket"
)
update_time = time.perf_counter() - time_start
print(f"Flattened bucket update time: {update_time:.03f}")
# Verify the weights were updated correctly
for i, param_name in enumerate(param_names):
_check_param(engine, param_name, [2.0] * 5)
engine.shutdown()
class TestServerUpdateWeightsFromTensorNonBlocking(CustomTestCase):
@classmethod
def setUpClass(cls):
cls.model = DEFAULT_SMALL_MODEL_NAME_FOR_TEST
cls.base_url = DEFAULT_URL_FOR_TEST
cls.process = popen_launch_server(
cls.model,
cls.base_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=["--max-running-requests", 8],
)
@classmethod
def tearDownClass(cls):
kill_process_tree(cls.process.pid)
def run_decode(self, max_new_tokens=32):
response = requests.post(
self.base_url + "/generate",
json={
"text": f"Question: {random.randint(0, 100)},The capital of France is",
"sampling_params": {
"temperature": 0,
"max_new_tokens": max_new_tokens,
"ignore_eos": True,
},
},
)
return response.json()
def get_model_info(self):
response = requests.get(self.base_url + "/get_model_info")
model_path = response.json()["model_path"]
print(json.dumps(response.json()))
return model_path
def pause_generation(self, mode):
response = requests.post(
self.base_url + "/pause_generation",
json={"mode": mode},
)
ret = response.json()
return ret
def continue_generation(self):
response = requests.post(
self.base_url + "/continue_generation",
json={},
)
ret = response.json()
return ret
def run_update_weights(self, named_tensors, flush_cache=True):
response = requests.post(
self.base_url + "/update_weights_from_tensor",
json={
"serialized_named_tensors": [
MultiprocessingSerializer.serialize(named_tensors, output_str=True)
],
"flush_cache": flush_cache,
},
)
ret = response.json()
return ret
def test_update_weights(self):
pause_generation_modes = ["in_place", "retract"]
for pause_generation_mode in pause_generation_modes:
num_requests = 32
with ThreadPoolExecutor(num_requests) as executor:
futures = [
executor.submit(self.run_decode, 3000) for _ in range(num_requests)
]
# ensure the decode has been started
time.sleep(2)
param_names = [
f"model.layers.{i}.mlp.up_proj.weight" for i in range(6, 16)
]
new_tensor = torch.full((16384, 2048), 1.5, device="cuda")
named_tensors = [(x, new_tensor) for x in param_names]
ret = self.pause_generation(pause_generation_mode)
ret = self.run_update_weights(
named_tensors, flush_cache=pause_generation_mode == "retract"
)
self.assertTrue(ret["success"])
ret = self.continue_generation()
for future in as_completed(futures):
self.assertNotEqual(
future.result()["meta_info"]["finish_reason"]["type"], "abort"
)
for param_name in param_names[:3]:
response = requests.post(
self.base_url + "/get_weights_by_name",
json={"name": param_name},
)
actual_values = torch.tensor(response.json())[0, :5]
assert torch.allclose(
actual_values, torch.tensor([1.5] * 5), atol=0.002
), f"{actual_values=}"
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()
test/srt/run_suite.py
-16
View File
@@ -28,10 +28,6 @@ suites = {
TestFile("openai_server/validation/test_openai_server_ignore_eos.py", 6),
TestFile("openai_server/validation/test_request_length_validation.py", 38),
TestFile("ops/test_repeat_interleave.py", 60),
# quant tests moved to test/registered/quant/
TestFile("rl/test_fp32_lm_head.py", 9),
# TestFile("rl/test_update_weights_from_disk.py", 210), # Temporarily disabled, see https://github.com/sgl-project/sglang/pull/13998
TestFile("rl/test_update_weights_from_tensor.py", 195),
TestFile("dllm/test_llada2_mini.py", 520),
TestFile("test_abort.py", 131),
TestFile("test_chunked_prefill.py", 312),
@@ -84,12 +80,10 @@ suites = {
TestFile("hicache/test_hicache_storage_mooncake_backend.py", 300),
TestFile("models/test_kimi_linear_models.py", 90),
TestFile("models/test_nvidia_nemotron_nano_v2.py", 132),
TestFile("rl/test_update_weights_from_distributed.py", 103),
TestFile("test_data_parallelism.py", 73),
TestFile("test_disaggregation_basic.py", 400),
TestFile("test_dp_attention.py", 350),
TestFile("test_load_weights_from_remote_instance.py", 72),
TestFile("test_patch_torch.py", 19),
],
"per-commit-4-gpu": [
TestFile("models/test_qwen3_next_models.py", 650),
@@ -97,7 +91,6 @@ suites = {
TestFile("test_multi_instance_release_memory_occupation.py", 64),
TestFile("test_pp_single_node.py", 500),
TestFile("test_epd_disaggregation.py", 150),
TestFile("rl/test_return_routed_experts.py", 300),
],
"per-commit-8-gpu-h200": [
TestFile("test_deepseek_v3_basic.py", 275),
@@ -146,9 +139,6 @@ suites = {
"__not_in_ci__": [
TestFile("test_release_memory_occupation.py", 200), # Temporarily disabled
TestFile("models/test_dummy_grok_models.py"),
TestFile(
"rl/test_update_weights_from_disk.py"
), # Temporarily disabled, see https://github.com/sgl-project/sglang/pull/13998
TestFile("test_bench_one_batch.py"),
TestFile("test_bench_serving.py"),
TestFile("test_eval_accuracy_large.py"),
@@ -191,9 +181,6 @@ suite_amd = {
TestFile("openai_server/validation/test_openai_server_ignore_eos.py", 85),
TestFile("openai_server/validation/test_request_length_validation.py", 31),
TestFile("ops/test_repeat_interleave.py", 75),
# quant tests moved to test/registered/quant/
TestFile("rl/test_fp32_lm_head.py", 15),
# TestFile("rl/test_update_weights_from_disk.py", 210), # Temporarily disabled, see https://github.com/sgl-project/sglang/pull/13998
TestFile("rotary_embedding/test_mrope.py", 15),
TestFile("test_abort.py", 51),
TestFile("test_bench_typebaseddispatcher.py", 10),
@@ -242,11 +229,8 @@ suite_amd = {
TestFile("test_gpt_oss_1gpu.py", 750),
],
"per-commit-2-gpu-amd": [
# TestFile("lora/test_lora_tp.py", 116), # Disabled temporarily, see https://github.com/sgl-project/sglang/issues/13107. Moved to test/registered/lora/
TestFile("rl/test_update_weights_from_distributed.py", 103),
TestFile("test_data_parallelism.py", 73),
TestFile("test_load_weights_from_remote_instance.py", 72),
# TestFile("test_patch_torch.py", 19), # Disabled temporarily, see https://github.com/sgl-project/sglang/issues/11127
],
"per-commit-4-gpu-amd": [
TestFile("test_pp_single_node.py", 150),
test/srt/test_patch_torch.py
-133
View File
@@ -1,133 +0,0 @@
import os
import traceback
import unittest
from typing import Dict, List
import torch
import torch.multiprocessing as mp
from sglang.srt.utils.patch_torch import monkey_patch_torch_reductions
class TestReleaseMemoryOccupation(unittest.TestCase):
def test_monkey_patch_torch_reductions(self):
mp.set_start_method("spawn", force=True)
for enable_patch in [False, True]:
for params in [
# Same visible devices
dict(
sender_info=dict(
visible_devices=[0, 1],
tensor_device=1,
),
receiver_info=dict(
visible_devices=[0, 1],
tensor_device=1,
),
),
# Different visible devices
dict(
sender_info=dict(
visible_devices=[0, 1],
tensor_device=1,
),
receiver_info=dict(
visible_devices=[1, 0],
# If enable patch, this should be fixed, and cuda:1 becomes cuda:0
tensor_device=0 if enable_patch else 1,
),
),
]:
with self.subTest(f"{enable_patch=} {params=}"):
self._test_monkey_patch_torch_reductions_core(
enable_patch=enable_patch, **params
)
def _test_monkey_patch_torch_reductions_core(
self,
sender_info: Dict,
receiver_info: Dict,
enable_patch: bool,
):
print(
f'test_monkey_patch_torch_reductions_core {os.environ.get("CUDA_VISIBLE_DEVICES")=}'
)
cuda_visible_devices_list: List[int] = [
int(x)
for x in os.environ.get("CUDA_VISIBLE_DEVICES", "0,1,2,3,4,5,6,7").split(
","
)
]
processes = []
output_reader, output_writer = mp.Pipe(duplex=False)
queue = mp.Queue()
for role, info in [
("sender", sender_info),
("receiver", receiver_info),
]:
os.environ["CUDA_VISIBLE_DEVICES"] = ",".join(
str(cuda_visible_devices_list[device])
for device in info["visible_devices"]
)
p = mp.Process(
target=_run_subprocess,
kwargs=dict(
role=role,
queue=queue,
output_writer=output_writer,
tensor_device=info["tensor_device"],
enable_patch=enable_patch,
),
)
p.start()
processes.append(p)
for _ in range(len(processes)):
self.assertTrue(
output_reader.recv(), f"Subprocess has error, please see logs above."
)
for p in processes:
p.join()
def _run_subprocess(
role: str, queue: mp.Queue, output_writer, tensor_device: int, enable_patch: bool
):
print(
f'subprocess[{role}] start {os.environ.get("CUDA_VISIBLE_DEVICES")=}',
flush=True,
)
if enable_patch:
print(f"subprocess[{role}] execute monkey_patch_torch_reductions", flush=True)
monkey_patch_torch_reductions()
try:
if role == "sender":
tensor = torch.tensor([1.0, 2.0], device=f"cuda:{tensor_device}")
print(f"sender queue.put {tensor=} {tensor.device=}")
queue.put(tensor)
assert queue.get() == "done"
elif role == "receiver":
tensor = queue.get()
print(f"receiver queue.get {tensor=} {tensor.device=}")
assert str(tensor.device) == f"cuda:{tensor_device}"
queue.put("done")
else:
raise NotImplementedError
execution_ok = True
except Exception as e:
print(f"subprocess[{role}] has error: {e}", flush=True)
traceback.print_exc()
execution_ok = False
output_writer.send(execution_ok)
output_writer.close()
if __name__ == "__main__":
unittest.main()