From 32e85eda58c560e5d5a596f71fce3682ac2ef80a Mon Sep 17 00:00:00 2001 From: Stas Bekman Date: Thu, 15 Jul 2021 18:11:17 -0700 Subject: [PATCH] [see_memory_usage] fix deprecation (#1234) Co-authored-by: Olatunji Ruwase --- deepspeed/runtime/utils.py | 14 ++++++++++++-- 1 file changed, 12 insertions(+), 2 deletions(-) diff --git a/deepspeed/runtime/utils.py b/deepspeed/runtime/utils.py index c792d2a6..55674b54 100755 --- a/deepspeed/runtime/utils.py +++ b/deepspeed/runtime/utils.py @@ -21,6 +21,16 @@ import torch.distributed as dist from deepspeed.utils import logger from numpy import prod +# pt-1.9 deprecations +if hasattr(torch.cuda, "memory_reserved"): + torch_memory_reserved = torch.cuda.memory_reserved +else: + torch_memory_reserved = torch.cuda.memory_allocated +if hasattr(torch.cuda, "max_memory_reserved"): + torch_max_memory_reserved = torch.cuda.max_memory_reserved +else: + torch_max_memory_reserved = torch.cuda.memory_cached + def noop_decorator(func): return func @@ -589,8 +599,8 @@ def see_memory_usage(message, force=False): logger.info( f"MA {round(torch.cuda.memory_allocated() / (1024 * 1024 * 1024),2 )} GB \ Max_MA {round(torch.cuda.max_memory_allocated() / (1024 * 1024 * 1024),2)} GB \ - CA {round(torch.cuda.memory_cached() / (1024 * 1024 * 1024),2)} GB \ - Max_CA {round(torch.cuda.max_memory_cached() / (1024 * 1024 * 1024))} GB ") + CA {round(torch_memory_reserved() / (1024 * 1024 * 1024),2)} GB \ + Max_CA {round(torch_max_memory_reserved() / (1024 * 1024 * 1024))} GB ") vm_stats = psutil.virtual_memory() used_GB = round(((vm_stats.total - vm_stats.available) / (1024**3)), 2) -- GitLab