diff --git a/python/sglang/srt/entrypoints/http_server.py b/python/sglang/srt/entrypoints/http_server.py index 1ebbbc072..a05179ed5 100644 --- a/python/sglang/srt/entrypoints/http_server.py +++ b/python/sglang/srt/entrypoints/http_server.py @@ -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, diff --git a/python/sglang/srt/managers/io_struct.py b/python/sglang/srt/managers/io_struct.py index f26ce2960..f4acf6907 100644 --- a/python/sglang/srt/managers/io_struct.py +++ b/python/sglang/srt/managers/io_struct.py @@ -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 diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index 45451c92e..851e07437 100644 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -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))) diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 7af9b6205..3ca2861e5 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -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: diff --git a/python/sglang/srt/managers/scheduler_update_weights_mixin.py b/python/sglang/srt/managers/scheduler_update_weights_mixin.py index fa0d612e2..1d2965a86 100644 --- a/python/sglang/srt/managers/scheduler_update_weights_mixin.py +++ b/python/sglang/srt/managers/scheduler_update_weights_mixin.py @@ -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) diff --git a/python/sglang/srt/managers/tokenizer_communicator_mixin.py b/python/sglang/srt/managers/tokenizer_communicator_mixin.py index 70129ea8c..558b78756 100644 --- a/python/sglang/srt/managers/tokenizer_communicator_mixin.py +++ b/python/sglang/srt/managers/tokenizer_communicator_mixin.py @@ -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: diff --git a/python/sglang/srt/managers/tokenizer_manager.py b/python/sglang/srt/managers/tokenizer_manager.py index a953702b4..c27a054da 100644 --- a/python/sglang/srt/managers/tokenizer_manager.py +++ b/python/sglang/srt/managers/tokenizer_manager.py @@ -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. diff --git a/test/srt/rl/test_update_weights_from_disk.py b/test/srt/rl/test_update_weights_from_disk.py index 02b283efd..0127b98dc 100644 --- a/test/srt/rl/test_update_weights_from_disk.py +++ b/test/srt/rl/test_update_weights_from_disk.py @@ -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}") diff --git a/test/srt/rl/test_update_weights_from_distributed.py b/test/srt/rl/test_update_weights_from_distributed.py index 45d2baf84..fd45a4e5c 100644 --- a/test/srt/rl/test_update_weights_from_distributed.py +++ b/test/srt/rl/test_update_weights_from_distributed.py @@ -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, ) diff --git a/test/srt/rl/test_update_weights_from_tensor.py b/test/srt/rl/test_update_weights_from_tensor.py index 0dc947b60..04dd01d96 100644 --- a/test/srt/rl/test_update_weights_from_tensor.py +++ b/test/srt/rl/test_update_weights_from_tensor.py @@ -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(