alihan/LongCodeZip
1
0
Fork 0
mirror of https://github.com/YerbaPage/LongCodeZip.git synced 2025-10-22 23:19:46 +03:00
Code Issues Packages Projects Releases Wiki Activity

Add gpu_memory_utilization to engine arguments

Browse Source 
Prevent shortage of GPU memory.
This commit is contained in:
STEVENTAN100
2025-10-21 10:32:54 +08:00
committed by GitHub
parent 22d21b9b23
commit 5478b117c8
1 changed files with 3 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
experiments/module-summarization/main.py
View File

@@ -19,7 +19,7 @@ from itertools import cycle
class LLMGenerator:
def __init__(self, model_name, device, **model_args):
# Create a vllm LLM instance
engine_args = EngineArgs(model=model_name, device=device, **model_args)
engine_args = EngineArgs(model=model_name, gpu_memory_utilization=0.8, device=device, **model_args)
self.model = LLM(**vars(engine_args))
self.model_name = model_name
self.device = device
@@ -58,7 +58,7 @@ class LLMGenerator:
class LLMScorer:
def __init__(self, model_name, device, **model_args):
# Create a vllm LLM instance
engine_args = EngineArgs(model=model_name, device=device, **model_args)
engine_args = EngineArgs(model=model_name, gpu_memory_utilization=0.8, device=device, **model_args)
self.model = LLM(**vars(engine_args))
self.model_name = model_name
self.device = device
@@ -1315,4 +1315,4 @@ def run_documentation_task(
if __name__ == "__main__":
fire.Fire(run_documentation_task)
fire.Fire(run_documentation_task)