mirror of
https://github.com/PaddlePaddle/FastDeploy.git
synced 2025-10-04 16:22:57 +08:00
113 lines
4.4 KiB
Python
113 lines
4.4 KiB
Python
"""
|
|
# Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License"
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
"""
|
|
import gc
|
|
import time
|
|
from typing import List, Optional
|
|
|
|
import paddle
|
|
import paddle.nn as nn
|
|
|
|
from fastdeploy.config import FDConfig
|
|
from fastdeploy.engine.request import Request
|
|
from fastdeploy.utils import get_logger
|
|
from fastdeploy.worker.gpu_model_runner import GPUModelRunner
|
|
from fastdeploy.worker.output import ModelRunnerOutput
|
|
from fastdeploy.worker.gpu_worker import GpuWorker
|
|
|
|
logger = get_logger("dcu_worker", "dcu_worker.log")
|
|
|
|
|
|
class DcuWorker(GpuWorker):
|
|
""" """
|
|
|
|
def __init__(
|
|
self,
|
|
fd_config: FDConfig,
|
|
local_rank: int,
|
|
rank: int,
|
|
):
|
|
super().__init__(
|
|
fd_config=fd_config,
|
|
local_rank=local_rank,
|
|
rank=rank,
|
|
)
|
|
pass
|
|
|
|
def determine_available_memory(self) -> int:
|
|
"""
|
|
Profiles the peak memory usage of the model to determine how much
|
|
memory can be used for KV cache without OOMs.
|
|
|
|
The engine will first conduct a profiling of the existing memory usage.
|
|
Then, it calculate the maximum possible number of GPU and CPU blocks
|
|
that can be allocated with the remaining free memory.
|
|
|
|
Tip:
|
|
You may limit the usage of GPU memory
|
|
by adjusting the `gpu_memory_utilization` parameter.
|
|
"""
|
|
# 1. Record memory state before profile run
|
|
Gb = 1024**3
|
|
start_time = time.perf_counter()
|
|
paddle.device.cuda.reset_max_memory_reserved(self.local_rank)
|
|
paddle.device.cuda.reset_max_memory_allocated(self.local_rank)
|
|
paddle_reserved_mem_before_run = paddle.device.cuda.max_memory_reserved(
|
|
self.local_rank)
|
|
paddle_allocated_mem_before_run = paddle.device.cuda.max_memory_allocated(
|
|
self.local_rank) # not reserved
|
|
|
|
total_gpu_memory = paddle.device.cuda.get_device_properties(self.local_rank).total_memory
|
|
before_used_gpu_memory = paddle.device.cuda.memory_allocated(self.local_rank)
|
|
|
|
|
|
logger.info((
|
|
"Before running the profile, the memory usage info is as follows:",
|
|
f"\nDevice Total memory: {total_gpu_memory / Gb}",
|
|
f"\nDevice used memory: {before_used_gpu_memory / Gb}",
|
|
f"\nPaddle reserved memory: {paddle_reserved_mem_before_run / Gb}",
|
|
f"\nPaddle allocated memory: {paddle_allocated_mem_before_run / Gb}"))
|
|
|
|
# 2. Profile run
|
|
self.model_runner.profile_run()
|
|
|
|
# 3. Statistical memory information
|
|
paddle_reserved_mem_after_run = paddle.device.cuda.max_memory_reserved(
|
|
self.local_rank)
|
|
paddle_allocated_mem_after_run = paddle.device.cuda.max_memory_allocated(
|
|
self.local_rank)
|
|
|
|
after_used_gpu_memory = paddle.device.cuda.memory_allocated(self.local_rank)
|
|
|
|
# v0 worker
|
|
model_block_memory_used = self.cal_theortical_kvcache()
|
|
paddle.device.cuda.empty_cache()
|
|
paddle_peak_increase = paddle_reserved_mem_after_run - paddle_allocated_mem_before_run
|
|
available_kv_cache_memory = total_gpu_memory * \
|
|
self.parallel_config.gpu_memory_utilization - after_used_gpu_memory - paddle_peak_increase
|
|
available_kv_cache_memory += model_block_memory_used * self.parallel_config.max_block_num
|
|
|
|
end_time = time.perf_counter()
|
|
logger.info(
|
|
("After running the profile, the memory usage info is as follows:",
|
|
f"\nDevice Total memory: {total_gpu_memory / Gb}",
|
|
f"\nDevice used memory: {after_used_gpu_memory / Gb}",
|
|
f"\nPaddle reserved memory: {paddle_reserved_mem_after_run / Gb}",
|
|
f"\nPaddle allocated memory: {paddle_allocated_mem_after_run / Gb}",
|
|
f"\nAvailable KV Cache meomory: {available_kv_cache_memory / Gb}",
|
|
f"Profile time: {end_time - start_time}"))
|
|
|
|
return available_kv_cache_memory # return to caculate the block num in this device
|