Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / activation_memory / _backend.py: 68%

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"""PyTorch device primitives used by the activation-swap implementation.""" 

16 

17__all__ = [ 

18 "Tensor", 

19 "get_device_handle", 

20 "new_stream", 

21 "get_stream_context", 

22 "get_current_stream", 

23 "new_event", 

24 "no_grad", 

25 "preserve_version_counter", 

26 "cat", 

27 "empty_like", 

28 "tree_map", 

29 "get_element_size", 

30 "alloc_tensor_buffer", 

31 "register_forward_pre_hook", 

32 "register_full_backward_hook", 

33 "register_full_backward_pre_hook", 

34] 

35 

36import torch 

37from torch import Tensor 

38 

39from hyper_parallel.core.utils.communication import get_device_handle 

40 

41 

42def new_stream(): 

43 """Create a new device stream on the current accelerator.""" 

44 return get_device_handle().Stream() 

45 

46 

47def get_stream_context(): 

48 """Return the stream context manager (``torch.npu.stream`` / ``torch.cuda.stream``).""" 

49 return get_device_handle().stream 

50 

51 

52def get_current_stream(): 

53 """Return the current device stream.""" 

54 return get_device_handle().current_stream() 

55 

56 

57def new_event(): 

58 """Create a new device event on the current accelerator.""" 

59 return get_device_handle().Event() 

60 

61 

62def no_grad(): 

63 """Return ``torch.no_grad()``.""" 

64 return torch.no_grad() 

65 

66 

67def preserve_version_counter(tensor): 

68 """Temporarily keep the tensor version counter unchanged across an in-place write.""" 

69 return torch.autograd._unsafe_preserve_version_counter(tensor) # pylint: disable=W0212 

70 

71 

72def cat(tensors, dim=0): 

73 """Concatenate tensors along *dim*.""" 

74 return torch.cat(tensors, dim=dim) 

75 

76 

77def empty_like(tensor, *, dtype=None, device=None, pin_memory=False): 

78 """Allocate an uninitialized tensor shaped like *tensor*.""" 

79 return torch.empty_like(tensor, dtype=dtype, device=device, pin_memory=pin_memory) 

80 

81 

82def tree_map(fn, tree): 

83 """Apply *fn* to every leaf of *tree* and rebuild the same structure.""" 

84 return torch.utils._pytree.tree_map(fn, tree) # pylint: disable=W0212 

85 

86 

87def get_element_size(tensor) -> int: 

88 """Return the size in bytes of one element of *tensor*.""" 

89 return tensor.element_size() 

90 

91 

92def alloc_tensor_buffer(numel: int, dtype, device="cpu", pin_memory: bool = False): 

93 """Allocate an uninitialized 1-D tensor buffer.""" 

94 if pin_memory: 

95 return torch.empty(numel, dtype=dtype, device="cpu", pin_memory=True) 

96 return torch.empty(numel, dtype=dtype, device=device) 

97 

98 

99def register_forward_pre_hook(module, hook, prepend=False, with_kwargs=False): 

100 """Register a forward pre-hook on *module*, ignoring *prepend*.""" 

101 del prepend 

102 return module.register_forward_pre_hook(hook, with_kwargs=with_kwargs) 

103 

104 

105def register_full_backward_hook(module, hook, prepend=False): 

106 """Register a full backward hook on *module*, ignoring *prepend*.""" 

107 del prepend 

108 return module.register_full_backward_hook(hook) 

109 

110 

111def register_full_backward_pre_hook(module, hook, prepend=False): 

112 """Register a full backward pre-hook on *module*, ignoring *prepend*.""" 

113 del prepend 

114 return module.register_full_backward_pre_hook(hook)