diff --git a/lightllm/common/basemodel/basemodel.py b/lightllm/common/basemodel/basemodel.py index e83de684a7..cc66dcb1c3 100755 --- a/lightllm/common/basemodel/basemodel.py +++ b/lightllm/common/basemodel/basemodel.py @@ -24,6 +24,7 @@ from lightllm.common.quantization import Quantcfg from lightllm.common.basemodel.triton_kernel.gather_token_id import gather_token, gather_token_prefill_decode_mixed from lightllm.utils.log_utils import init_logger +from lightllm.utils.gc_utils import freeze_gc, gc_frozen_and_disabled from lightllm.utils.dist_utils import get_dp_world_size from lightllm.utils.envs_utils import get_env_start_args, get_llm_data_type, get_added_mtp_kv_layer_num from lightllm.distributed.communication_op import dist_group_manager @@ -127,11 +128,13 @@ def __init__(self, kvargs): self._autotune_warmup() self._init_padded_req() - self._init_cudagraph() - self._init_prefill_cuda_graph() + with gc_frozen_and_disabled("cudagraph-capture"): + self._init_cudagraph() + self._init_prefill_cuda_graph() self._check_max_len_infer() torch.cuda.empty_cache() set_model_init_status(True) + freeze_gc("model-infer-worker") return def _init_config(self): @@ -604,7 +607,6 @@ def _decode( @final def _context_forward(self, infer_state: InferStateInfo): - input_embs = self.pre_infer.context_forward(infer_state.input_ids, infer_state, self.pre_post_weight) if self.args.enable_dp_prefill_balance: assert not self.args.enable_prefill_cudagraph, "not support now" diff --git a/lightllm/server/api_http.py b/lightllm/server/api_http.py index 270e2a8cfd..39c6dcb5d6 100755 --- a/lightllm/server/api_http.py +++ b/lightllm/server/api_http.py @@ -47,6 +47,7 @@ from .api_lightllm import lightllm_get_score from lightllm.utils.envs_utils import get_env_start_args, get_lightllm_websocket_max_message_size from lightllm.utils.log_utils import init_logger +from lightllm.utils.gc_utils import freeze_gc from lightllm.utils.error_utils import ClientDisconnected, ServerBusyError from lightllm.server.metrics.manager import MetricClient from lightllm.utils.envs_utils import get_unique_server_name @@ -497,4 +498,5 @@ async def startup_event(): g_objs.set_args(get_env_start_args()) loop.create_task(g_objs.httpserver_manager.handle_loop()) logger.info(f"server start up ok, loop use is {asyncio.get_event_loop()}") + freeze_gc("httpserver") return diff --git a/lightllm/server/detokenization/manager.py b/lightllm/server/detokenization/manager.py index 8c213914c7..3f3ece65ee 100644 --- a/lightllm/server/detokenization/manager.py +++ b/lightllm/server/detokenization/manager.py @@ -16,6 +16,7 @@ import pickle import time from lightllm.utils.log_utils import init_logger +from lightllm.utils.gc_utils import freeze_gc from lightllm.utils.envs_utils import get_unique_server_name logger = init_logger(__name__) @@ -179,5 +180,6 @@ def start_detokenization_process(args, pipe_writer): pipe_writer.send(str(e)) raise pipe_writer.send("init ok") + freeze_gc("detokenization") manager.handle_loop() return diff --git a/lightllm/server/router/manager.py b/lightllm/server/router/manager.py index dfb8866601..2b0261b452 100644 --- a/lightllm/server/router/manager.py +++ b/lightllm/server/router/manager.py @@ -26,6 +26,7 @@ from lightllm.server.multi_level_kv_cache.cpu_cache_client import CpuKvCacheClient from lightllm.server.core.objs.shm_objs_io_buffer import ShmObjsIOBuffer from lightllm.utils.log_utils import init_logger, log_time_ready +from lightllm.utils.gc_utils import freeze_gc from lightllm.utils.profiler import ProfilerCmd from lightllm.server.router.token_load import TokenLoad from lightllm.server.metrics.manager import MetricClient @@ -569,5 +570,6 @@ def handle_exception(loop, context): raise pipe_writer.send("init ok") + freeze_gc("router") loop.run_until_complete(router.loop_for_fwd()) return diff --git a/lightllm/utils/gc_utils.py b/lightllm/utils/gc_utils.py new file mode 100644 index 0000000000..d726ebe791 --- /dev/null +++ b/lightllm/utils/gc_utils.py @@ -0,0 +1,23 @@ +import gc +from contextlib import contextmanager +from lightllm.utils.log_utils import init_logger + +logger = init_logger(__name__) + + +def freeze_gc(tag: str = "") -> None: + gc.collect() + gc.freeze() + logger.info(f"gc.freeze done ({tag}): frozen={gc.get_freeze_count()}") + + +@contextmanager +def gc_frozen_and_disabled(tag: str = ""): + freeze_gc(tag) + was_enabled = gc.isenabled() + gc.disable() + try: + yield + finally: + if was_enabled: + gc.enable()