1from __future__ import annotations

2

3import ctypes

4import os

5from typing import Dict, List, Optional

6

7import torch

8from transformers import TrainerCallback

9

10

11def _format_bytes_gb(num_bytes: int) -> str:

12 return f"{num_bytes / (1024 ** 3):.2f}GB"

13

14

15def get_ram_usage() -> Optional[Dict[str, float | int]]:

16 if os.name == "nt":

17 class MEMORYSTATUSEX(ctypes.Structure):

18 _fields_ = [

19 ("dwLength", ctypes.c_ulong),

20 ("dwMemoryLoad", ctypes.c_ulong),

21 ("ullTotalPhys", ctypes.c_ulonglong),

22 ("ullAvailPhys", ctypes.c_ulonglong),

23 ("ullTotalPageFile", ctypes.c_ulonglong),

24 ("ullAvailPageFile", ctypes.c_ulonglong),

25 ("ullTotalVirtual", ctypes.c_ulonglong),

26 ("ullAvailVirtual", ctypes.c_ulonglong),

27 ("ullAvailExtendedVirtual", ctypes.c_ulonglong),

28 ]

29

30 stat = MEMORYSTATUSEX()

31 stat.dwLength = ctypes.sizeof(MEMORYSTATUSEX)

32 ok = ctypes.windll.kernel32.GlobalMemoryStatusEx(ctypes.byref(stat))

33 if not ok:

34 return None

35 total = int(stat.ullTotalPhys)

36 avail = int(stat.ullAvailPhys)

37 if total <= 0:

38 return None

39 used = total - avail

40 return {"ratio": float(used / total), "used": used, "total": total}

41

42 try:

43 page_size = int(os.sysconf("SC_PAGE_SIZE"))

44 total_pages = int(os.sysconf("SC_PHYS_PAGES"))

45 avail_pages = int(os.sysconf("SC_AVPHYS_PAGES"))

46 except Exception:

47 return None

48 total = page_size * total_pages

49 avail = page_size * avail_pages

50 if total <= 0:

51 return None

52 used = total - avail

53 return {"ratio": float(used / total), "used": used, "total": total}

54

55

56def get_vram_usage() -> List[Dict[str, float | int]]:

57 if not torch.cuda.is_available():

58 return []

59

60 out: List[Dict[str, float | int]] = []

61 for idx in range(torch.cuda.device_count()):

62 try:

63 free, total = torch.cuda.mem_get_info(idx)

64 used = int(total - free)

65 ratio = float(used / total) if total > 0 else 0.0

66 out.append({"device": idx, "ratio": ratio, "used": used, "total": int(total)})

67 except Exception:

68 total = int(torch.cuda.get_device_properties(idx).total_memory)

69 used = int(torch.cuda.memory_reserved(idx))

70 ratio = float(used / total) if total > 0 else 0.0

71 out.append({"device": idx, "ratio": ratio, "used": used, "total": total})

72 return out

73

74

75class MemoryWarningCallback(TrainerCallback):

76 def __init__(

77 self,

78 ram_threshold: float = 0.90,

79 vram_threshold: float = 0.90,

80 check_every_steps: int = 10,

81 repeat_warning_steps: int = 200,

82 ) -> None:

83 if not (0.0 < ram_threshold <= 1.0):

84 raise ValueError(f"ram_threshold must be in (0,1], got {ram_threshold}")

85 if not (0.0 < vram_threshold <= 1.0):

86 raise ValueError(f"vram_threshold must be in (0,1], got {vram_threshold}")

87 self.ram_threshold = float(ram_threshold)

88 self.vram_threshold = float(vram_threshold)

89 self.check_every_steps = max(1, int(check_every_steps))

90 self.repeat_warning_steps = max(1, int(repeat_warning_steps))

91 self._last_check_step = -1

92 self._last_ram_warn_step = -10**9

93 self._last_vram_warn_step: Dict[int, int] = {}

94

95 def on_train_begin(self, args, state, control, **kwargs):

96 print(

97 "[MemWarn] enabled: "

98 f"ram>={self.ram_threshold:.2f}, "

99 f"vram>={self.vram_threshold:.2f}, "

100 f"check_every_steps={self.check_every_steps}, "

101 f"repeat_warning_steps={self.repeat_warning_steps}"

102 )

103 return control

104

105 def _should_check(self, step: int) -> bool:

106 if step <= 0:

107 return False

108 if self._last_check_step < 0:

109 self._last_check_step = step

110 return True

111 if (step - self._last_check_step) >= self.check_every_steps:

112 self._last_check_step = step

113 return True

114 return False

115

116 def on_step_end(self, args, state, control, **kwargs):

117 step = int(getattr(state, "global_step", 0))

118 if not self._should_check(step):

119 return control

120

121 ram = get_ram_usage()

122 if ram is not None and float(ram["ratio"]) >= self.ram_threshold:

123 if (step - self._last_ram_warn_step) >= self.repeat_warning_steps:

124 self._last_ram_warn_step = step

125 print(

126 "[MemWarn][RAM] "

127 f"step={step}, usage={float(ram['ratio']) * 100:.1f}%, "

128 f"used={_format_bytes_gb(int(ram['used']))}, "

129 f"total={_format_bytes_gb(int(ram['total']))}"

130 )

131

132 for item in get_vram_usage():

133 device = int(item["device"])

134 ratio = float(item["ratio"])

135 if ratio < self.vram_threshold:

136 continue

137 last_step = self._last_vram_warn_step.get(device, -10**9)

138 if (step - last_step) < self.repeat_warning_steps:

139 continue

140 self._last_vram_warn_step[device] = step

141 print(

142 "[MemWarn][VRAM] "

143 f"step={step}, cuda:{device}, usage={ratio * 100:.1f}%, "

144 f"used={_format_bytes_gb(int(item['used']))}, "

145 f"total={_format_bytes_gb(int(item['total']))}"

146 )

147 return control