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

[feat] support in-flight weight update (#10071)

Browse Source 
Co-authored-by: 赵晨阳 <zhaochen20@outlook.com>
This commit is contained in:
ShawnY112358
2025-11-25 22:03:13 -08:00
committed by GitHub
co-authored by 赵晨阳
parent 7130ad3a29
commit 007c3e234c
10 changed files with 401 additions and 34 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/srt/rl/test_update_weights_from_distributed.py
+85 -10
View File
@@ -18,6 +18,7 @@ import os
import random
import time
import unittest
from concurrent.futures import ThreadPoolExecutor
import numpy as np
import requests
@@ -68,6 +69,8 @@ def init_process(
backend,
checking_parameters,
tie_word_embeddings,
barrier,
pause_generation_mode,
):
torch.cuda.set_device(rank)
@@ -81,6 +84,7 @@ def init_process(
checking_parameters,
tie_word_embeddings,
state_dict_key_to_shape,
barrier,
)
elif rank in [1, 2]:
init_process_sgl(
@@ -94,6 +98,8 @@ def init_process(
state_dict_key_to_shape,
backend,
tp_size,
barrier,
pause_generation_mode,
)
@@ -106,6 +112,7 @@ def init_process_hf(
checking_parameters,
tie_word_embeddings,
state_dict_key_to_shape,
barrier,
):
# These two environment variables are very important
# to avoid unexpected behaviors of CUDA and NCCL.
@@ -162,6 +169,7 @@ def init_process_hf(
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
@@ -208,6 +216,8 @@ def init_process_sgl(
state_dict_key_to_shape,
backend,
tp_size,
barrier,
pause_generation_mode,
):
torch.cuda.set_device(rank)
torch.cuda.synchronize()
@@ -282,8 +292,25 @@ def init_process_sgl(
},
)
torch.cuda.synchronize()
time_begin_update = time.perf_counter()
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
@@ -300,6 +327,14 @@ def init_process_sgl(
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,
@@ -315,10 +350,23 @@ def init_process_sgl(
"dtypes": dtypes,
"shapes": shapes,
"group_name": "test_parameter_update_group",
"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
@@ -383,6 +431,7 @@ def test_update_weights_from_distributed(
state_dict_key_to_shape,
truncate_size,
checking_parameters,
pause_generation_mode=None,
):
tie_word_embeddings = (
True if model_name == DEFAULT_SMALL_MODEL_NAME_FOR_TEST else False
@@ -393,6 +442,7 @@ def test_update_weights_from_distributed(
)
param_queue = mp.Queue()
results = {}
barrier = mp.Barrier(1 + dp_size)
context = mp.spawn(
init_process,
@@ -406,6 +456,8 @@ def test_update_weights_from_distributed(
backend,
checking_parameters,
tie_word_embeddings,
barrier,
pause_generation_mode,
),
nprocs=1 + dp_size,
join=False,
@@ -558,28 +610,50 @@ class TestUpdateWeightsFromDistributed(CustomTestCase):
# 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
test_suits = [
(1, 1, DEFAULT_SMALL_MODEL_NAME_FOR_TEST, mode),
(1, 1, DEFAULT_SMALL_MODEL_NAME_FOR_TEST, mode, pause_generation_mode),
]
else:
test_suits = [
(1, 1, DEFAULT_SMALL_MODEL_NAME_FOR_TEST, "Engine"),
(1, 1, DEFAULT_MODEL_NAME_FOR_TEST, "Sever"),
(1, 1, DEFAULT_SMALL_MODEL_NAME_FOR_TEST, "Engine", None),
(
1,
1,
DEFAULT_MODEL_NAME_FOR_TEST,
"Sever",
random.choice(["in_place", "retract"]),
),
]
if torch.cuda.device_count() >= 4:
test_suits.extend(
[
(2, 1, DEFAULT_SMALL_MODEL_NAME_FOR_TEST, "Engine"),
(1, 2, DEFAULT_MODEL_NAME_FOR_TEST, "Server"),
(2, 1, DEFAULT_SMALL_MODEL_NAME_FOR_TEST, "Engine", None),
(
1,
2,
DEFAULT_MODEL_NAME_FOR_TEST,
"Server",
random.choice(["in_place", "retract"]),
),
]
)
if torch.cuda.device_count() >= 5:
test_suits.extend(
[
(2, 2, DEFAULT_SMALL_MODEL_NAME_FOR_TEST, "Engine"),
(2, 2, DEFAULT_MODEL_NAME_FOR_TEST, "Server"),
(2, 2, DEFAULT_SMALL_MODEL_NAME_FOR_TEST, "Engine", None),
(
2,
2,
DEFAULT_MODEL_NAME_FOR_TEST,
"Server",
random.choice(["in_place", "retract"]),
),
]
)
@@ -615,7 +689,7 @@ class TestUpdateWeightsFromDistributed(CustomTestCase):
"lm_head.weight",
]
for tp_size, dp_size, model_name, backend in test_suits:
for tp_size, dp_size, model_name, backend, pause_generation_mode in test_suits:
test_update_weights_from_distributed(
tp_size,
dp_size,
@@ -624,6 +698,7 @@ class TestUpdateWeightsFromDistributed(CustomTestCase):
model_state_dict_shapes[model_name],
truncate_size,
checking_parameters,
pause_generation_mode,
)