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

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

2# Copyright 2026 Huawei Technologies Co., Ltd 

3# 

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

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

6# You may obtain a copy of the License at 

7# 

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

9# 

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

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

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

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

14# limitations under the License. 

15# ============================================================================ 

16 

17"""Trainer-side host/device memory helpers. 

18 

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

20(05 §10.4); function names and signatures are unchanged. 

21""" 

22 

23__all__ = [ 

24 "empty_cache", 

25 "peak_device_memory", 

26 "print_cpu_memory_info", 

27 "print_device_mem_info", 

28 "reset_peak_device_memory", 

29] 

30 

31import gc 

32from typing import Dict, Optional 

33 

34import psutil 

35 

36from hyper_parallel.models.build_options import ( 

37 IS_CUDA_AVAILABLE, 

38 IS_NPU_AVAILABLE, 

39 get_device_type, 

40 get_torch_device, 

41) 

42from hyper_parallel.trainer.runtime.logging import _info_rank0 

43 

44 

45def print_device_mem_info(prompt: str = "VRAM usage") -> None: 

46 """ 

47 Logs VRAM info. 

48 """ 

49 if get_device_type() == "cpu": 

50 print_cpu_memory_info() 

51 else: 

52 memory_allocated = get_torch_device().memory_allocated() / (1024**3) 

53 max_memory_allocated = get_torch_device().max_memory_allocated() / (1024**3) 

54 _info_rank0(f"{prompt}: cur {memory_allocated:.2f}GB, max {max_memory_allocated:.2f}GB.") 

55 

56 

57def print_cpu_memory_info() -> None: 

58 """Log CPU usage and system memory information on global rank zero.""" 

59 cpu_usage = psutil.cpu_percent(interval=1) # sampling for 1 sec 

60 _info_rank0(f"CPU Usage: {cpu_usage}%") 

61 

62 memory_info = psutil.virtual_memory() 

63 _info_rank0(f"Total Memory: {memory_info.total / (1024**3):.2f} GB") 

64 _info_rank0(f"Available Memory: {memory_info.available / (1024**3):.2f} GB") 

65 _info_rank0(f"Used Memory: {memory_info.used / (1024**3):.2f} GB") 

66 _info_rank0(f"Memory Usage: {memory_info.percent}%") 

67 

68 

69def empty_cache() -> None: 

70 """ 

71 Collects system memory. 

72 """ 

73 gc.collect() 

74 

75 if IS_CUDA_AVAILABLE or IS_NPU_AVAILABLE: 

76 get_torch_device().empty_cache() 

77 

78 

79 

80def reset_peak_device_memory() -> None: 

81 """Reset peak allocator counters for one measured training step.""" 

82 if IS_CUDA_AVAILABLE or IS_NPU_AVAILABLE: 

83 get_torch_device().reset_peak_memory_stats() 

84 

85 

86def peak_device_memory() -> Dict[str, Optional[int]]: 

87 """Return peak allocated/reserved bytes, or nulls for CPU execution.""" 

88 if not (IS_CUDA_AVAILABLE or IS_NPU_AVAILABLE): 

89 return { 

90 "peak_memory_allocated_bytes": None, 

91 "peak_memory_reserved_bytes": None, 

92 } 

93 device = get_torch_device() 

94 return { 

95 "peak_memory_allocated_bytes": int(device.max_memory_allocated()), 

96 "peak_memory_reserved_bytes": int(device.max_memory_reserved()), 

97 }