From 40cba20e8781e2ac5936fca081a88493c3ce8b43 Mon Sep 17 00:00:00 2001 From: Cyril Vallez Date: Thu, 17 Apr 2025 16:11:54 +0200 Subject: [PATCH] Ensure positive warm-up size (#37581) ensure > 0 --- src/transformers/modeling_utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/transformers/modeling_utils.py b/src/transformers/modeling_utils.py index da2f07c0a6..3be400b93d 100644 --- a/src/transformers/modeling_utils.py +++ b/src/transformers/modeling_utils.py @@ -5898,7 +5898,7 @@ def caching_allocator_warmup(model: PreTrainedModel, expanded_device_map: Dict, # to OOM. See https://github.com/huggingface/transformers/issues/37436#issuecomment-2808982161 for more details. # Note that we use an absolute value instead of device proportion here, as a 8GiB device could still allocate too much # if using e.g. 90% of device size, while a 140GiB device would allocate too little - byte_count = min(byte_count, int(device_memory - 1.2 * 1024**3)) + byte_count = min(byte_count, max(0, int(device_memory - 1.2 * 1024**3))) # Allocate memory _ = torch.empty(byte_count // factor, dtype=torch.float16, device=device, requires_grad=False)