Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / auto_parallel / hyper_offload / _patch_accelerator.py: 40%

1# Copyright 2026 Huawei Technologies Co., Ltd 

2# 

3# Licensed under the Apache License, Version 2.0 (the "License"); 

4# you may not use this file except in compliance with the License. 

5# You may obtain a copy of the License at 

6# 

7# http://www.apache.org/licenses/LICENSE-2.0 

8# 

9# Unless required by applicable law or agreed to in writing, software 

10# distributed under the License is distributed on an "AS IS" BASIS, 

11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 

12# See the License for the specific language governing permissions and 

13# limitations under the License. 

14# ============================================================================ 

15"""Monkey-patch ``torch.accelerator`` on PyTorch < 2.6. 

16 

17PyTorch 2.5.x (and earlier) do **not** include the ``torch.accelerator`` 

18module that was introduced in PyTorch 2.6. This module creates a 

19compatible namespace and injects it into the ``torch`` module so that 

20all ``torch.accelerator.xxx`` calls work transparently regardless of 

21the PyTorch version. 

22 

23Backend detection priority 

24 #1. Ascend NPU — ``torch.npu.is_available()`` 

25 #2. NVIDIA GPU — ``torch.cuda.is_available()`` 

26 #3. No accelerator — stub (``is_available()`` returns ``False``) 

27 

28Usage 

29 Importing this module applies the patch as a side-effect:: 

30 

31 import hyper_parallel.auto_parallel.hyper_offload._patch_accelerator 

32 

33 Or simply import the parent package (``__init__.py`` already does this):: 

34 

35 import hyper_parallel.auto_parallel.hyper_offload # patch applied 

36 

37Notes 

38 For NPU backends the APIs ``reset_peak_memory_stats`` and 

39 ``max_memory_allocated`` are guarded with ``hasattr`` — if the 

40 installed ``torch_npu`` does not provide them, they fall back to 

41 no-ops / zero-returning stubs so that callers do not crash. 

42 ``current_stream`` and ``empty_cache`` are expected on any 

43 accelerator backend; if absent they will raise ``AttributeError`` 

44 at call time (fail-fast). 

45""" 

46 

47from __future__ import annotations 

48 

49import types 

50 

51import torch 

52 

53 

54def _patch() -> None: 

55 """Inject ``torch.accelerator`` if not already present (PyTorch ≥ 2.6).""" 

56 if hasattr(torch, "accelerator"): 

57 return # PyTorch ≥ 2.6 — native support 

58 

59 accelerator = types.ModuleType("accelerator") 

60 accelerator.__doc__ = "Monkey-patched ``torch.accelerator`` for PyTorch < 2.6" 

61 

62 # ── Backend detection ──────────────────────────────────────────── 

63 

64 if hasattr(torch, "npu") and torch.npu.is_available(): 

65 # ---- Ascend NPU backend ---- 

66 def is_available() -> bool: 

67 return torch.npu.is_available() 

68 

69 def current_accelerator() -> torch.device: 

70 return torch.device("npu", index=torch.npu.current_device()) 

71 

72 accelerator.is_available = is_available 

73 accelerator.current_accelerator = current_accelerator 

74 accelerator.current_stream = torch.npu.current_stream 

75 accelerator.empty_cache = torch.npu.empty_cache 

76 

77 # These may not exist in older torch_npu versions 

78 if hasattr(torch.npu, "reset_peak_memory_stats"): 

79 accelerator.reset_peak_memory_stats = torch.npu.reset_peak_memory_stats 

80 else: 

81 def _npu_reset_peak_memory_stats() -> None: 

82 pass # no-op fallback 

83 accelerator.reset_peak_memory_stats = _npu_reset_peak_memory_stats 

84 

85 if hasattr(torch.npu, "max_memory_allocated"): 

86 accelerator.max_memory_allocated = torch.npu.max_memory_allocated 

87 else: 

88 def _npu_max_memory_allocated() -> int: 

89 return 0 # fallback 

90 accelerator.max_memory_allocated = _npu_max_memory_allocated 

91 

92 elif torch.cuda.is_available(): 

93 # ---- NVIDIA GPU (CUDA) backend ---- 

94 def is_available() -> bool: 

95 return torch.cuda.is_available() 

96 

97 def current_accelerator() -> torch.device: 

98 return torch.device("cuda", index=torch.cuda.current_device()) 

99 

100 accelerator.is_available = is_available 

101 accelerator.current_accelerator = current_accelerator 

102 accelerator.current_stream = torch.cuda.current_stream 

103 accelerator.empty_cache = torch.cuda.empty_cache 

104 accelerator.reset_peak_memory_stats = torch.cuda.reset_peak_memory_stats 

105 accelerator.max_memory_allocated = torch.cuda.max_memory_allocated 

106 

107 else: 

108 # ---- No accelerator available (CPU-only) ---- 

109 def is_available() -> bool: 

110 return False 

111 

112 def current_accelerator() -> torch.device: 

113 raise RuntimeError( 

114 "No accelerator available (torch.accelerator not supported on CPU)" 

115 ) 

116 

117 def current_stream() -> torch.Stream: 

118 raise RuntimeError( 

119 "No accelerator available (torch.accelerator not supported on CPU)" 

120 ) 

121 

122 def empty_cache() -> None: 

123 pass 

124 

125 def reset_peak_memory_stats() -> None: 

126 pass 

127 

128 def max_memory_allocated() -> int: 

129 return 0 

130 

131 accelerator.is_available = is_available 

132 accelerator.current_accelerator = current_accelerator 

133 accelerator.current_stream = current_stream 

134 accelerator.empty_cache = empty_cache 

135 accelerator.reset_peak_memory_stats = reset_peak_memory_stats 

136 accelerator.max_memory_allocated = max_memory_allocated 

137 

138 # ── Inject into torch ──────────────────────────────────────────── 

139 torch.accelerator = accelerator 

140 

141 

142_patch()