Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / trainer / runtime / device.py: 37%

1# Copyright 2025-2026 Bytedance Ltd. and/or its affiliates 

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# Following codes are inspired from https://github.com/volcengine/verl/blob/main/verl/utils/device.py 

16 

17"""Trainer-side runtime device API. 

18 

19Split out of the former ``auto_models/components/utils/device.py`` in stage 7 

20(05 §10.4): runtime device selection belongs to the Trainer. The 

21accelerator-discovery primitives (``get_device_type`` / ``get_torch_device`` 

22/ ``get_device_id`` and the availability flags) are single-sourced in 

23``hyper_parallel.models.build_options`` — the Trainer-independent home 

24that model construction uses when no explicit device option is given — and 

25re-exported here so Trainer consumers have one import site. 

26""" 

27 

28__all__ = [ 

29 "IS_CUDA_AVAILABLE", 

30 "IS_NPU_AVAILABLE", 

31 "get_device_type", 

32 "get_device_name", 

33 "get_torch_device", 

34 "get_device_id", 

35 "get_dist_comm_backend", 

36 "synchronize", 

37 "stream_synchronize", 

38 "empty_cache", 

39 "set_device", 

40 "get_device_rng_state", 

41 "set_device_rng_state", 

42 "is_nccl_backend", 

43 "is_hccl_backend", 

44 "get_gpu_compute_capability", 

45 "is_sm90_or_above", 

46 "get_compute_units", 

47] 

48 

49import logging 

50from typing import Any 

51 

52import torch 

53 

54from hyper_parallel.models.build_options import ( 

55 IS_CUDA_AVAILABLE, 

56 IS_NPU_AVAILABLE, 

57 get_device_id, 

58 get_device_type, 

59 get_torch_device, 

60) 

61 

62 

63logger = logging.getLogger(__name__) 

64 

65 

66def get_device_name() -> str: 

67 """Get real device name, e.g. A100, H100""" 

68 return get_torch_device().get_device_name() 

69 

70 

71def get_dist_comm_backend() -> str: 

72 """Return distributed communication backend type based on device type.""" 

73 if IS_CUDA_AVAILABLE: 

74 return "nccl" 

75 if IS_NPU_AVAILABLE: 

76 return "hccl" 

77 raise RuntimeError(f"No available distributed communication backend found on device type {get_device_type()}.") 

78 

79 

80def synchronize() -> None: 

81 """Execute torch synchronize operation.""" 

82 get_torch_device().synchronize() 

83 

84 

85def stream_synchronize() -> None: 

86 """Execute device stream synchronize operation.""" 

87 if IS_CUDA_AVAILABLE: 

88 torch.cuda.current_stream().synchronize() 

89 elif IS_NPU_AVAILABLE: 

90 torch.npu.current_stream().synchronize() 

91 else: 

92 synchronize() 

93 

94 

95def empty_cache() -> None: 

96 """Execute torch empty cache operation.""" 

97 get_torch_device().empty_cache() 

98 

99 

100def set_device(device: torch.types.Device) -> None: 

101 """Execute set device operation.""" 

102 get_torch_device().set_device(device) 

103 

104 

105def get_device_rng_state() -> Any: 

106 """Return the current accelerator RNG state, or ``None`` when unavailable. 

107 

108 Checkpointing needs the device RNG alongside the CPU one: dropout, 

109 ``torch.randn`` initializers and any device-side sampling draw from it, so a 

110 resume that only restores ``torch.get_rng_state()`` diverges from the 

111 original run. 

112 """ 

113 if get_device_type() == "cpu": 

114 return None 

115 get_rng_state = getattr(get_torch_device(), "get_rng_state", None) 

116 if get_rng_state is None: 

117 return None 

118 return get_rng_state() 

119 

120 

121def set_device_rng_state(state: Any) -> None: 

122 """Restore an accelerator RNG state captured by :func:`get_device_rng_state`.""" 

123 if state is None or get_device_type() == "cpu": 

124 return 

125 set_rng_state = getattr(get_torch_device(), "set_rng_state", None) 

126 if set_rng_state is None: 

127 logger.warning("Device namespace has no set_rng_state; device RNG not restored.") 

128 return 

129 set_rng_state(state) 

130 

131 

132def is_nccl_backend(backend: str | None = None) -> bool: 

133 """Check if the distributed communication backend is NCCL.""" 

134 return (backend or get_dist_comm_backend()) == "nccl" 

135 

136 

137def is_hccl_backend() -> bool: 

138 """Check if the distributed communication backend is HCCL.""" 

139 return get_dist_comm_backend() == "hccl" 

140 

141 

142def get_gpu_compute_capability(device: torch.types.Device | int | None = None) -> int: 

143 """Return the compute capability as an integer (e.g. 70, 80, 90), or 0 if no GPU.""" 

144 if not IS_CUDA_AVAILABLE: 

145 return 0 

146 major, minor = torch.cuda.get_device_capability(device) 

147 return major * 10 + minor 

148 

149 

150def is_sm90_or_above() -> bool: 

151 """Check if the current CUDA device has SM90+ capability.""" 

152 return get_gpu_compute_capability() >= 90 

153 

154 

155def get_compute_units() -> int: 

156 """ 

157 Returns the number of streaming multiprocessors (SMs) or equivalent compute units 

158 for the available accelerator. 

159 """ 

160 device_type = getattr(torch.accelerator.current_accelerator(), "type", "cpu") 

161 

162 if device_type == "cuda": 

163 device_properties = torch.cuda.get_device_properties(0) 

164 num_compute_units = device_properties.multi_processor_count 

165 elif device_type == "xpu": 

166 device_properties = torch.xpu.get_device_properties(0) 

167 num_compute_units = device_properties.max_compute_units 

168 else: 

169 logger.warning("No CUDA or XPU device available. Using CPU.") 

170 num_compute_units = torch.get_num_threads() 

171 

172 return num_compute_units