From cfabfd2e84da90a666deacf0c5111452b1101a35 Mon Sep 17 00:00:00 2001 From: clc Date: Thu, 5 Nov 2026 17:16:10 +0800 Subject: [PATCH] fix: release device allocator cache after Accelerate offload --- src/twinkle/model/transformers/strategy/accelerate.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/twinkle/model/transformers/strategy/accelerate.py b/src/twinkle/model/transformers/strategy/accelerate.py index 81a9f1e7c..f101967e9 100644 --- a/src/twinkle/model/transformers/strategy/accelerate.py +++ b/src/twinkle/model/transformers/strategy/accelerate.py @@ -218,8 +218,8 @@ def offload_to_cpu(self, model, optimizer=None, *, offload_model: bool = True) - if optimizer is not None: move_optimizer_state(optimizer, torch.device('cpu')) record_offload(self, model=offload_model, optimizer=optimizer is not None) - if torch.cuda.is_available(): - torch.cuda.empty_cache() + # Release reserved blocks on every backend before the colocated sampler wakes. + release_caching_allocator() def reload_to_gpu(self, model, optimizer=None) -> None: """Bring back what :meth:`offload_to_cpu` handed over, to the device it was moved off."""