[RL] Call torch.cuda.empty_cache() for in-place pause mode to avoid OOM (#24854)
Co-authored-by: Byron Hsu <byron@periodiclabs.ai> Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
co-authored by
Byron Hsu
Cursor
parent
44efc23a9a
commit
cfd3fd00d0
@@ -1378,7 +1378,12 @@ class PauseGenerationReqInput(BaseReq):
|
|||||||
|
|
||||||
@dataclass
|
@dataclass
|
||||||
class ContinueGenerationReqInput(BaseReq):
|
class ContinueGenerationReqInput(BaseReq):
|
||||||
pass
|
# Call torch.cuda.empty_cache() before un-pausing. Returns blocks
|
||||||
|
# cached by the PyTorch allocator (left over from transient allocs
|
||||||
|
# during post-weight-update processing) back to the driver before
|
||||||
|
# inference resumes, with no race against active streams. Set to
|
||||||
|
# False to skip the empty_cache call.
|
||||||
|
torch_empty_cache: bool = True
|
||||||
|
|
||||||
|
|
||||||
@dataclass
|
@dataclass
|
||||||
|
|||||||
@@ -3646,6 +3646,15 @@ class Scheduler(
|
|||||||
self.chunked_req = None
|
self.chunked_req = None
|
||||||
|
|
||||||
def continue_generation(self, recv_req: ContinueGenerationReqInput):
|
def continue_generation(self, recv_req: ContinueGenerationReqInput):
|
||||||
|
if recv_req.torch_empty_cache:
|
||||||
|
before_mb = torch.cuda.memory_reserved() / (1024 * 1024)
|
||||||
|
torch.cuda.empty_cache()
|
||||||
|
after_mb = torch.cuda.memory_reserved() / (1024 * 1024)
|
||||||
|
logger.info(
|
||||||
|
f"[continue_generation] torch.cuda.empty_cache() called: "
|
||||||
|
f"reserved {before_mb:.1f} MB -> {after_mb:.1f} MB "
|
||||||
|
f"(freed {before_mb - after_mb:.1f} MB)"
|
||||||
|
)
|
||||||
self._engine_paused = False
|
self._engine_paused = False
|
||||||
|
|
||||||
def load_lora_adapter(
|
def load_lora_adapter(
|
||||||
|
|||||||
Reference in New Issue
Block a user