add available memory check to accelerators (microsoft#4508)

* add available memory check to accelerator * catch case where nvmlInit fails * add pynvml to reqs * fix for cpu systems * Update accelerator/cuda_accelerator.py Co-authored-by: Michael Wyatt <[email protected]> * simplify --------- Co-authored-by: Michael Wyatt <[email protected]>
EleutherAI · Oct 16, 2023 · 12aedac · 12aedac
1 parent 78c518e
commit 12aedac
Show file tree

Hide file tree

Showing 6 changed files with 39 additions and 0 deletions.
diff --git a/accelerator/abstract_accelerator.py b/accelerator/abstract_accelerator.py
@@ -147,6 +147,10 @@ def max_memory_reserved(self, device_index=None):
  def total_memory(self, device_index=None):
  ...
 
+ @abc.abstractmethod
+ def available_memory(self, device_index=None):
+ ...
+
  # Data types
  @abc.abstractmethod
  def is_bf16_supported(self):

diff --git a/accelerator/cpu_accelerator.py b/accelerator/cpu_accelerator.py
@@ -159,6 +159,9 @@ def max_memory_reserved(self, device_index=None):
  def total_memory(self, device_index=None):
  return psutil.virtual_memory().total
 
+ def available_memory(self, device_index=None):
+ return psutil.virtual_memory().available
+
  # Misc
  def amp(self):
  return torch.cpu.amp

diff --git a/accelerator/cuda_accelerator.py b/accelerator/cuda_accelerator.py
@@ -15,12 +15,29 @@
 except ImportError:
  pass
 
+# Delay import pynvml to avoid import error when CUDA is not available
+pynvml = None
+
 
 class CUDA_Accelerator(DeepSpeedAccelerator):
 
  def __init__(self):
  self._name = 'cuda'
  self._communication_backend_name = 'nccl'
+ if pynvml is None:
+ self._init_pynvml()
+
+ def _init_pynvml(self):
+ global pynvml
+ try:
+ import pynvml
+ except ImportError:
+ return
+ try:
+ pynvml.nvmlInit()
+ except pynvml.NVMLError:
+ pynvml = None
+ return
 
  def is_synchronized_device(self):
  return False
@@ -136,6 +153,14 @@ def max_memory_reserved(self, device_index=None):
  def total_memory(self, device_index=None):
  return torch.cuda.get_device_properties(device_index).total_memory
 
+ def available_memory(self, device_index=None):
+ if pynvml:
+ handle = pynvml.nvmlDeviceGetHandleByIndex(device_index)
+ info = pynvml.nvmlDeviceGetMemoryInfo(handle)
+ return info.free
+ else:
+ return self.total_memory(device_index) - self.memory_allocated(device_index)
+
  # Data types
  def is_bf16_supported(self):
  return torch.cuda.is_bf16_supported()

diff --git a/accelerator/mps_accelerator.py b/accelerator/mps_accelerator.py
@@ -131,6 +131,9 @@ def max_memory_reserved(self, device_index=None):
  def total_memory(self, device_index=None):
  return
 
+ def available_memory(self, device_index=None):
+ return
+
  # Data types
  def is_bf16_supported(self):
  return False

diff --git a/accelerator/npu_accelerator.py b/accelerator/npu_accelerator.py
@@ -127,6 +127,9 @@ def max_memory_reserved(self, device_index=None):
  def total_memory(self, device_index=None):
  return torch.npu.get_device_properties(device_index).total_memory
 
+ def available_memory(self, device_index=None):
+ return self.total_memory(device_index) - self.memory_allocated(device_index)
+
  # Data types
  def is_bf16_supported(self):
  return torch.npu.is_bf16_supported()

diff --git a/requirements/requirements.txt b/requirements/requirements.txt
@@ -5,5 +5,6 @@ packaging>=20.0
 psutil
 py-cpuinfo
 pydantic
+pynvml
 torch
 tqdm