SUMIN/HuggingFace_transformer
1
0
Fork 0
Code Issues Pull Requests Actions 7 Packages Projects Releases Wiki Activity

Add compatibility with skip_memory_metrics for mps device (#29264)

Browse Source 
* Add compatibility with mps device

* fix

* typo and style
This commit is contained in:
Marc Sun
2024-02-27 09:58:43 -05:00
committed by GitHub
parent 5c341d4555
commit 8a1faf2803
1 changed files with 15 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

16
src/transformers/trainer_utils.py
View File

@@ -526,6 +526,8 @@ class TrainerMemoryTracker:
elif is_torch_npu_available(): elif is_torch_npu_available():
self.torch.npu.reset_peak_memory_stats() self.torch.npu.reset_peak_memory_stats()
self.torch.npu.empty_cache() self.torch.npu.empty_cache()
elif is_torch_mps_available():
self.torch.mps.empty_cache()
# gpu # gpu
if self.torch is not None: if self.torch is not None:
@@ -535,6 +537,8 @@ class TrainerMemoryTracker:
self.gpu_mem_used_at_start = self.torch.xpu.memory_allocated() self.gpu_mem_used_at_start = self.torch.xpu.memory_allocated()
elif is_torch_npu_available(): elif is_torch_npu_available():
self.gpu_mem_used_at_start = self.torch.npu.memory_allocated() self.gpu_mem_used_at_start = self.torch.npu.memory_allocated()
elif is_torch_mps_available():
self.gpu_mem_used_at_start = self.torch.mps.current_allocated_memory()
# cpu # cpu
self.cpu_mem_used_at_start = self.cpu_mem_used() self.cpu_mem_used_at_start = self.cpu_mem_used()
@@ -564,6 +568,8 @@ class TrainerMemoryTracker:
self.torch.xpu.empty_cache() self.torch.xpu.empty_cache()
elif is_torch_npu_available(): elif is_torch_npu_available():
self.torch.npu.empty_cache() self.torch.npu.empty_cache()
elif is_torch_mps_available():
self.torch.mps.empty_cache()
# concepts: # concepts:
# - alloc_delta: the difference of allocated memory between the end and the start # - alloc_delta: the difference of allocated memory between the end and the start
@@ -581,6 +587,11 @@ class TrainerMemoryTracker:
elif is_torch_npu_available(): elif is_torch_npu_available():
self.gpu_mem_used_now = self.torch.npu.memory_allocated() self.gpu_mem_used_now = self.torch.npu.memory_allocated()
self.gpu_mem_used_peak = self.torch.npu.max_memory_allocated() self.gpu_mem_used_peak = self.torch.npu.max_memory_allocated()
elif is_torch_mps_available():
self.gpu_mem_used_now = self.torch.mps.current_allocated_memory()
# self.torch.mps.max_memory_allocated() does not exist yet
self.gpu_mem_used_peak = None
else: else:
raise ValueError("No available GPU device found!") raise ValueError("No available GPU device found!")
@@ -588,8 +599,11 @@ class TrainerMemoryTracker:
"begin": self.gpu_mem_used_at_start, "begin": self.gpu_mem_used_at_start,
"end": self.gpu_mem_used_now, "end": self.gpu_mem_used_now,
"alloc": (self.gpu_mem_used_now - self.gpu_mem_used_at_start), "alloc": (self.gpu_mem_used_now - self.gpu_mem_used_at_start),
"peaked": max(0, self.gpu_mem_used_peak - self.gpu_mem_used_now),
} }
if self.gpu_mem_used_peak is not None:
self.gpu[self.cur_stage]["peaked"] = max(0, self.gpu_mem_used_peak - self.gpu_mem_used_now)
else:
self.gpu[self.cur_stage]["peaked"] = "Not available"
# cpu # cpu
self.cpu_mem_used_now = self.cpu_mem_used() self.cpu_mem_used_now = self.cpu_mem_used()