minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 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
python/sglang/srt/entrypoints/http_server.py
+6 -4
View File
@@ -78,6 +78,7 @@ from sglang.srt.managers.io_struct import (
AbortReq,
CloseSessionReqInput,
ConfigureLoggingReq,
ContinueGenerationReqInput,
DestroyWeightsUpdateGroupReqInput,
EmbeddingReqInput,
GenerateReqInput,
@@ -87,6 +88,7 @@ from sglang.srt.managers.io_struct import (
LoadLoRAAdapterReqInput,
OpenSessionReqInput,
ParseFunctionCallReq,
PauseGenerationReqInput,
ProfileReqInput,
ReleaseMemoryOccupationReqInput,
ResumeMemoryOccupationReqInput,
@@ -1087,9 +1089,9 @@ async def separate_reasoning_request(obj: SeparateReasoningReqInput, request: Re
@app.post("/pause_generation")
async def pause_generation(request: Request):
async def pause_generation(obj: PauseGenerationReqInput, request: Request):
"""Pause generation."""
await _global_state.tokenizer_manager.pause_generation()
await _global_state.tokenizer_manager.pause_generation(obj)
return ORJSONResponse(
content={"message": "Generation paused successfully.", "status": "ok"},
status_code=200,
@@ -1097,9 +1099,9 @@ async def pause_generation(request: Request):
@app.post("/continue_generation")
async def continue_generation(request: Request):
async def continue_generation(obj: ContinueGenerationReqInput, request: Request):
"""Continue generation."""
await _global_state.tokenizer_manager.continue_generation()
await _global_state.tokenizer_manager.continue_generation(obj)
return ORJSONResponse(
content={"message": "Generation continued successfully.", "status": "ok"},
status_code=200,
python/sglang/srt/managers/io_struct.py
+38 -1
View File
@@ -21,7 +21,7 @@ import uuid
from abc import ABC
from dataclasses import dataclass, field
from enum import Enum
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union
from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Union
from sglang.srt.lora.lora_registry import LoRARef
from sglang.srt.managers.schedule_batch import BaseFinishReason
@@ -1064,6 +1064,41 @@ class FlushCacheReqOutput(BaseReq):
success: bool
@dataclass
class PauseGenerationReqInput(BaseReq):
"""
Note that the PauseGenerationRequests is only supported in SGLang Server.
abort: Abort and return all requests currently being processed.
in_place: Pause the scheduler's event_loop from performing inference;
only non-inference requests (e.g., control commands) will be handled.
The requests in the engine will be paused and stay in the event_loop,
then continue generation after continue_generation with the old kv cache.
Note: In 'inplace' mode, flush_cache will fail if there are any requests
in the running_batch.
retract: Pause the scheduler's event loop from performing inference;
only non-inference requests will be handled, and all currently running
requests will be retracted back to the waiting_queue.
Note: The KV cache can be flushed in this mode and will be automatically
recomputed after continue_generation.
"""
mode: Literal["abort", "retract", "in_place"] = "abort"
def __post_init__(self):
allowed = ["abort", "retract", "in_place"]
if self.mode not in allowed:
raise ValueError(
f"Invalid mode: {self.mode!r}. " f"Expected one of {allowed}."
)
@dataclass
class ContinueGenerationReqInput(BaseReq):
pass
@dataclass
class UpdateWeightFromDiskReqInput(BaseReq):
# The model path with the new weights
@@ -1084,6 +1119,8 @@ class UpdateWeightFromDiskReqInput(BaseReq):
recapture_cuda_graph: bool = False
# The trainer step id. Used to know which step's weights are used for sampling.
token_step: int = 0
# Whether to flush the cache after updating weights
flush_cache: bool = True
@dataclass
python/sglang/srt/managers/schedule_batch.py
+10 -1
View File
@@ -1512,8 +1512,17 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
evict_from_tree_cache(self.tree_cache, num_tokens)
return self._is_available_size_sufficient(num_tokens)
def retract_all(self, server_args: ServerArgs):
retracted_reqs = self.reqs
for idx in range(len(self.reqs)):
self.release_req(idx, len(self.reqs) - idx, server_args)
self.filter_batch(retracted_reqs)
return retracted_reqs
def retract_decode(
self, server_args: ServerArgs
self,
server_args: ServerArgs,
) -> Tuple[List[Req], float, List[Req]]:
"""Retract the decoding requests when there is not enough memory."""
sorted_indices = list(range(len(self.reqs)))
python/sglang/srt/managers/scheduler.py
+35 -2
View File
@@ -73,6 +73,7 @@ from sglang.srt.managers.io_struct import (
ClearHiCacheReqInput,
ClearHiCacheReqOutput,
CloseSessionReqInput,
ContinueGenerationReqInput,
DestroyWeightsUpdateGroupReqInput,
ExpertDistributionReq,
ExpertDistributionReqOutput,
@@ -93,6 +94,7 @@ from sglang.srt.managers.io_struct import (
LoadLoRAAdapterReqOutput,
OpenSessionReqInput,
OpenSessionReqOutput,
PauseGenerationReqInput,
ProfileReq,
ReleaseMemoryOccupationReqInput,
ResumeMemoryOccupationReqInput,
@@ -443,6 +445,7 @@ class Scheduler(
if self.device == "cpu":
self.default_stream.synchronize = lambda: None # No-op for CPU
self.forward_sleep_time = None
self._engine_paused = False
# Init chunked prefill
self.chunked_prefill_size = server_args.chunked_prefill_size
@@ -568,6 +571,8 @@ class Scheduler(
(LoadLoRAAdapterReqInput, self.load_lora_adapter),
(UnloadLoRAAdapterReqInput, self.unload_lora_adapter),
(GetLoadReqInput, self.get_load),
(PauseGenerationReqInput, self.pause_generation),
(ContinueGenerationReqInput, self.continue_generation),
]
)
@@ -953,6 +958,9 @@ class Scheduler(
recv_reqs = self.recv_requests()
self.process_input_requests(recv_reqs)
if self._engine_paused:
continue
batch = self.get_next_batch_to_run()
self.cur_batch = batch
@@ -985,6 +993,9 @@ class Scheduler(
recv_reqs = self.recv_requests()
self.process_input_requests(recv_reqs)
if self._engine_paused:
continue
batch = self.get_next_batch_to_run()
self.cur_batch = batch
@@ -2154,8 +2165,7 @@ class Scheduler(
def _is_no_request(self):
no_request = (
len(self.waiting_queue) == 0
and self.running_batch.is_empty()
self.running_batch.is_empty()
and (self.last_batch is None or self.last_batch.is_empty())
and (self.cur_batch is None or self.cur_batch.is_empty())
and (not self.enable_overlap or len(self.result_queue) == 0)
@@ -2428,6 +2438,29 @@ class Scheduler(
def _pause_engine(self) -> Tuple[List[Req], int]:
raise NotImplementedError()
def pause_generation(self, recv_req: PauseGenerationReqInput):
self._engine_paused = True
if self.enable_overlap and self.last_batch:
# Process the results of the last batch
tmp_batch, tmp_result = self.result_queue.popleft()
self.process_batch_result(tmp_batch, tmp_result)
self.last_batch = None
self.cur_batch = None
if recv_req.mode == "retract":
self.running_batch.filter_batch()
if len(self.running_batch.reqs) != 0:
retracted_reqs = self.running_batch.retract_all(self.server_args)
for req in retracted_reqs:
self._add_request_to_queue(req)
self.running_batch.batch_is_full = False
self.chunked_req = None
def continue_generation(self, recv_req: ContinueGenerationReqInput):
self._engine_paused = False
def load_lora_adapter(
self, recv_req: LoadLoRAAdapterReqInput
) -> LoadLoRAAdapterReqOutput:
python/sglang/srt/managers/scheduler_update_weights_mixin.py
+3 -2
View File
@@ -44,8 +44,9 @@ class SchedulerUpdateWeightsMixin:
"""In-place update of the weights from disk."""
success, message = self.tp_worker.update_weights_from_disk(recv_req)
if success:
flush_cache_success = self.flush_cache()
assert flush_cache_success, "Cache flush failed after updating weights"
if recv_req.flush_cache:
flush_cache_success = self.flush_cache()
assert flush_cache_success, "Cache flush failed after updating weights"
else:
logger.error(message)
return UpdateWeightFromDiskReqOutput(success, message, 0)
python/sglang/srt/managers/tokenizer_communicator_mixin.py
+14
View File
@@ -404,6 +404,14 @@ class TokenizerCommunicatorMixin:
if obj.abort_all_requests:
self.abort_request(abort_all=True)
# Immediately update the weights if the engine is in paused state
async with self.is_pause_cond:
if self.is_pause:
result = (await self.update_weights_from_distributed_communicator(obj))[
0
]
return result.success, result.message
# This means that weight sync
# cannot run while requests are in progress.
async with self.model_update_lock.writer_lock:
@@ -457,6 +465,12 @@ class TokenizerCommunicatorMixin:
if obj.abort_all_requests:
self.abort_request(abort_all=True)
# Immediately update the weights if the engine is in paused state
async with self.is_pause_cond:
if self.is_pause:
result = (await self.update_weights_from_tensor_communicator(obj))[0]
return result.success, result.message
# This means that weight sync
# cannot run while requests are in progress.
async with self.model_update_lock.writer_lock:
python/sglang/srt/managers/tokenizer_manager.py
+21 -10
View File
@@ -55,6 +55,7 @@ from sglang.srt.managers.io_struct import (
BatchTokenizedEmbeddingReqInput,
BatchTokenizedGenerateReqInput,
ConfigureLoggingReq,
ContinueGenerationReqInput,
EmbeddingReqInput,
FreezeGCReq,
GenerateReqInput,
@@ -62,6 +63,7 @@ from sglang.srt.managers.io_struct import (
HealthCheckOutput,
LoadLoRAAdapterReqInput,
OpenSessionReqOutput,
PauseGenerationReqInput,
SessionParams,
TokenizedEmbeddingReqInput,
TokenizedGenerateReqInput,
@@ -1246,21 +1248,25 @@ class TokenizerManager(TokenizerCommunicatorMixin):
self.metrics_collector.labels
)
async def pause_generation(self):
async def pause_generation(self, obj: PauseGenerationReqInput):
async with self.is_pause_cond:
self.is_pause = True
# we are using the model_update_lock to check if there is still on-going requests.
while True:
# TODO: maybe make it async instead of fire-and-forget
self.abort_request(abort_all=True)
is_locked = await self.model_update_lock.is_locked()
if not is_locked:
break
await asyncio.sleep(1.0)
if obj.mode != "abort":
await self.send_to_scheduler.send_pyobj(obj)
else:
# we are using the model_update_lock to check if there is still on-going requests.
while True:
# TODO: maybe make it async instead of fire-and-forget
self.abort_request(abort_all=True)
is_locked = await self.model_update_lock.is_locked()
if not is_locked:
break
await asyncio.sleep(1.0)
async def continue_generation(self):
async def continue_generation(self, obj: ContinueGenerationReqInput):
async with self.is_pause_cond:
self.is_pause = False
await self.send_to_scheduler.send_pyobj(obj)
self.is_pause_cond.notify_all()
async def update_weights_from_disk(
@@ -1278,6 +1284,11 @@ class TokenizerManager(TokenizerCommunicatorMixin):
if obj.abort_all_requests:
self.abort_request(abort_all=True)
# Immediately update the weights if the engine is in paused state
async with self.is_pause_cond:
if self.is_pause:
return await self._wait_for_model_update_from_disk(obj)
if True: # Keep this redundant check to simplify some internal code sync
# Hold the lock if it is not async. This means that weight sync
# cannot run while requests are in progress.
test/srt/rl/test_update_weights_from_disk.py
+71 -3
View File
@@ -98,19 +98,51 @@ class TestServerUpdateWeightsFromDisk(CustomTestCase):
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):
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},
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()
print(json.dumps(ret))
return ret
def test_update_weights(self):
@@ -138,6 +170,42 @@ class TestServerUpdateWeightsFromDisk(CustomTestCase):
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}")
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,
)
test/srt/rl/test_update_weights_from_tensor.py
+118 -1
View File
@@ -1,12 +1,23 @@
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, CustomTestCase
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):
@@ -167,6 +178,112 @@ class TestUpdateWeightsFromTensor(CustomTestCase):
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(