riaz.somc/transformers
1
0
Fork 0
mirror of https://github.com/huggingface/transformers.git synced 2025-08-03 03:31:05 +06:00
Code Issues Actions 2 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
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 15 additions and 1 deletions
Show Stats 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():
self.torch.npu.reset_peak_memory_stats()
self.torch.npu.empty_cache()
elif is_torch_mps_available():
self.torch.mps.empty_cache()
# gpu
if self.torch is not None:
@ -535,6 +537,8 @@ class TrainerMemoryTracker:
self.gpu_mem_used_at_start = self.torch.xpu.memory_allocated()
elif is_torch_npu_available():
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
self.cpu_mem_used_at_start = self.cpu_mem_used()
@ -564,6 +568,8 @@ class TrainerMemoryTracker:
self.torch.xpu.empty_cache()
elif is_torch_npu_available():
self.torch.npu.empty_cache()
elif is_torch_mps_available():
self.torch.mps.empty_cache()
# concepts:
# - alloc_delta: the difference of allocated memory between the end and the start
@ -581,6 +587,11 @@ class TrainerMemoryTracker:
elif is_torch_npu_available():
self.gpu_mem_used_now = self.torch.npu.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:
raise ValueError("No available GPU device found!")
@ -588,8 +599,11 @@ class TrainerMemoryTracker:
"begin": self.gpu_mem_used_at_start,
"end": self.gpu_mem_used_now,
"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
self.cpu_mem_used_now = self.cpu_mem_used()