Fix KV cache pool sized far too small when weight-loading memory is still referenced (#36583)
This commit is contained in:
@@ -1,5 +1,6 @@
|
|||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import gc
|
||||||
import logging
|
import logging
|
||||||
import math
|
import math
|
||||||
from dataclasses import dataclass, field
|
from dataclasses import dataclass, field
|
||||||
@@ -1812,6 +1813,12 @@ class KVCacheConfigurator:
|
|||||||
# KV pool budget = currently-free GPU memory minus the non-static runtime
|
# KV pool budget = currently-free GPU memory minus the non-static runtime
|
||||||
# slack (pre_model_load_memory * (1 - mem_fraction_static)). Whatever is
|
# slack (pre_model_load_memory * (1 - mem_fraction_static)). Whatever is
|
||||||
# already resident (model weights, etc.) is thus charged against it.
|
# already resident (model weights, etc.) is thus charged against it.
|
||||||
|
# Weight-loading temporaries can still be referenced at this point, and
|
||||||
|
# empty_cache() (which get_available_gpu_memory already calls) cannot
|
||||||
|
# reclaim referenced blocks. Without collecting first, the KV budget is
|
||||||
|
# measured against an understated free-memory figure and the pool can be
|
||||||
|
# sized orders of magnitude too small while GPU memory sits idle.
|
||||||
|
gc.collect()
|
||||||
available_gpu_memory = get_available_gpu_memory(
|
available_gpu_memory = get_available_gpu_memory(
|
||||||
self.device,
|
self.device,
|
||||||
self.gpu_id,
|
self.gpu_id,
|
||||||
|
|||||||
Reference in New Issue
Block a user