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

support non disturbing remote instance weight loader v2 (#14997)

Browse Source 
Signed-off-by: Anqi Shen <amy.saq@antgroup.com>
This commit is contained in:
amysaq2023
2025-12-16 14:39:56 -08:00
committed by GitHub
parent a4c762811a
commit ccc8f3b266
11 changed files with 557 additions and 40 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/srt/test_load_weights_from_remote_instance.py
+54 -6
View File
@@ -62,6 +62,7 @@ def init_process(
seed_instance_group_base_port,
event_seed_ready,
event_dst_ready_list,
remote_instance_loader_backend,
):
torch.cuda.set_device(rank)
@@ -90,6 +91,7 @@ def init_process(
tp_size,
event_seed_ready,
event_dst_ready_list,
remote_instance_loader_backend,
)
@@ -122,6 +124,7 @@ def init_process_seed(
str(rank),
"--tp-size",
str(tp_size),
"--remote-instance-weight-loader-start-seed-via-transfer-engine",
),
)
torch.cuda.synchronize()
@@ -159,6 +162,7 @@ def init_process_dst(
tp_size,
event_seed_ready,
event_dst_ready_list,
remote_instance_loader_backend,
):
torch.cuda.set_device(rank * tp_size)
torch.cuda.synchronize()
@@ -186,6 +190,10 @@ def init_process_dst(
remote_instance_weight_loader_seed_instance_service_port=seed_instance_service_port,
remote_instance_weight_loader_send_weights_group_ports=ports,
load_format="remote_instance",
remote_instance_weight_loader_backend=remote_instance_loader_backend,
remote_instance_weight_loader_start_seed_via_transfer_engine=(
remote_instance_loader_backend == "transfer_engine"
),
)
else:
host, _, port = DEFAULT_URL_FOR_TEST.rpartition(":")
@@ -213,6 +221,9 @@ def init_process_dst(
f"[{','.join(str(port) for port in ports)}]",
"--load-format",
"remote_instance",
"--remote-instance-weight-loader-backend",
remote_instance_loader_backend,
"--remote-instance-weight-loader-start-seed-via-transfer-engine",
),
)
torch.cuda.synchronize()
@@ -250,9 +261,10 @@ def test_load_weights_from_remote_instance(
seed_instance_ip,
seed_instance_service_port,
seed_instance_group_base_port,
remote_instance_loader_backend,
):
print(
f"Testing model: {model_name} tp_size: {tp_size}, dp_size: {dp_size} backend: {backends}"
f"Testing model: {model_name} tp_size: {tp_size}, dp_size: {dp_size} backend: {backends} remote_instance_loader_backend: {remote_instance_loader_backend}"
)
param_queue = mp.Queue()
results = {}
@@ -276,6 +288,7 @@ def test_load_weights_from_remote_instance(
seed_instance_group_base_port,
event_seed_ready,
event_dst_ready_list,
remote_instance_loader_backend,
),
nprocs=1 + dp_size,
join=False,
@@ -340,14 +353,42 @@ class TestLoadWeightsFromRemoteInstance(CustomTestCase):
# test_suits : tp, dp, model_name, backend, dst_instance_id
if is_in_ci():
mode = random.choice(["Engine", "Server"])
remote_instance_loader_backend = random.choice(["nccl", "transfer_engine"])
test_suits = [
(1, 1, DEFAULT_SMALL_MODEL_NAME_FOR_TEST, [mode]),
(
1,
1,
DEFAULT_SMALL_MODEL_NAME_FOR_TEST,
[mode],
remote_instance_loader_backend,
),
]
else:
test_suits = [
(1, 1, DEFAULT_SMALL_MODEL_NAME_FOR_TEST, ["Engine"]),
(1, 1, DEFAULT_SMALL_MODEL_NAME_FOR_TEST, ["Sever"]),
(2, 2, DEFAULT_SMALL_MODEL_NAME_FOR_TEST, ["Engine", "Server"]),
(1, 1, DEFAULT_SMALL_MODEL_NAME_FOR_TEST, ["Engine"], "nccl"),
(1, 1, DEFAULT_SMALL_MODEL_NAME_FOR_TEST, ["Server"], "nccl"),
(2, 2, DEFAULT_SMALL_MODEL_NAME_FOR_TEST, ["Engine", "Server"], "nccl"),
(
1,
1,
DEFAULT_SMALL_MODEL_NAME_FOR_TEST,
["Engine"],
"transfer_engine",
),
(
1,
1,
DEFAULT_SMALL_MODEL_NAME_FOR_TEST,
["Server"],
"transfer_engine",
),
(
2,
2,
DEFAULT_SMALL_MODEL_NAME_FOR_TEST,
["Engine", "Server"],
"transfer_engine",
),
]
truncate_size = 10
@@ -365,7 +406,13 @@ class TestLoadWeightsFromRemoteInstance(CustomTestCase):
"model.norm.weight",
]
for tp_size, dp_size, model_name, backends in test_suits:
for (
tp_size,
dp_size,
model_name,
backends,
remote_instance_loader_backend,
) in test_suits:
test_load_weights_from_remote_instance(
tp_size,
dp_size,
@@ -376,6 +423,7 @@ class TestLoadWeightsFromRemoteInstance(CustomTestCase):
"127.0.0.1",
DEFAULT_PORT_FOR_SRT_TEST_RUNNER + 1000,
60000,
remote_instance_loader_backend,
)