huggingface/text-generation-inference
5
0
Fork 0
mirror of https://github.com/huggingface/text-generation-inference.git synced 2025-04-19 22:02:06 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix: also show total memory after full warmup

Browse Source
This commit is contained in:
drbh 2024-08-22 17:57:51 +00:00
parent 8b4cd2a9fc
commit e152cb022b
1 changed files with 2 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
server/text_generation_server/models/flash_causal_lm.py
View File

@ -1386,7 +1386,8 @@ class FlashCausalLM(Model):
total_cuda_graph_memory = free_memory_post_alloc - last_available_memory total_cuda_graph_memory = free_memory_post_alloc - last_available_memory
log_master( log_master(
logger.info, logger.info,
f"Total memory used for CUDA graphs: {total_cuda_graph_memory/1024/1024:.2f} MB", f"Total memory used for CUDA graphs: {total_cuda_graph_memory/1024/1024:.2f} MB"
f"\nTotal memory available: {last_available_memory/1024/1024:.2f} MB",
) )
except torch.cuda.OutOfMemoryError: except torch.cuda.OutOfMemoryError:
logger.exception("Decode cuda graph warmup failed") logger.exception("Decode cuda graph warmup failed")