Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / distributed / attention_swap.py: 28%

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"""Attention activation swap support for Hugging Face models.""" 

16 

17__all__ = [ 

18 "apply_attention_swap", 

19 "attention_swap_policy", 

20 "validate_attention_swap", 

21] 

22 

23import logging 

24import weakref 

25from dataclasses import dataclass, field 

26from typing import Optional 

27 

28import torch 

29from torch import nn 

30 

31from hyper_parallel.core.activation_memory import CheckpointPolicy, SwapManager, swap_wrapper 

32from hyper_parallel.core.activation_memory.swap import _teardown_wired_swap_layers 

33 

34logger = logging.getLogger(__name__) 

35 

36_MIN_SWAP_TENSOR_BYTES = 1024 * 1024 

37 

38 

39@dataclass 

40class _AttentionTarget: 

41 """One attention module and every parent attribute that references it.""" 

42 

43 module: nn.Module 

44 references: list[tuple[nn.Module, str]] = field(default_factory=list) 

45 

46 

47def attention_swap_policy(tensor: torch.Tensor) -> CheckpointPolicy: 

48 """Select large, independently-owned attention activations for swapping. 

49 

50 Args: 

51 tensor: Tensor saved by autograd inside an attention forward pass. 

52 

53 Returns: 

54 ``MUST_SWAP`` for tensors that are safe and worthwhile to transfer; 

55 otherwise ``MUST_SAVE``. 

56 """ 

57 if not tensor.requires_grad or tensor.dim() < 2: 

58 return CheckpointPolicy.MUST_SAVE 

59 storage_bytes = tensor.untyped_storage().size() 

60 tensor_bytes = tensor.numel() * tensor.element_size() 

61 if storage_bytes != tensor_bytes or tensor_bytes < _MIN_SWAP_TENSOR_BYTES: 

62 return CheckpointPolicy.MUST_SAVE 

63 return CheckpointPolicy.MUST_SWAP 

64 

65 

66def validate_attention_swap( 

67 activation_swap: str, 

68 *, 

69 activation_checkpoint: Optional[str] = None, 

70 enable_compile: bool = False, 

71 pp_size: int = 1, 

72) -> None: 

73 """Validate attention swap mode and incompatible model-build features. 

74 

75 Args: 

76 activation_swap: Requested activation swap mode. 

77 activation_checkpoint: Activation recomputation mode. 

78 enable_compile: Whether graph compilation is enabled. 

79 pp_size: Configured pipeline-parallel world size. 

80 

81 Raises: 

82 ValueError: If the mode is invalid or an unsupported combination is enabled. 

83 """ 

84 if activation_swap not in ("none", "attention"): 

85 raise ValueError( 

86 "activation_swap must be one of ('none', 'attention'), " 

87 f"got {activation_swap!r}" 

88 ) 

89 if activation_swap == "none": 

90 return 

91 if enable_compile: 

92 raise ValueError("activation_swap='attention' is incompatible with torch.compile") 

93 if activation_checkpoint not in (None, "off"): 

94 raise ValueError("activation_swap='attention' is incompatible with activation checkpointing") 

95 if pp_size != 1: 

96 raise ValueError("activation_swap='attention' does not support pipeline parallelism") 

97 

98 

99def _record_attention_target( 

100 targets_by_id: dict[int, _AttentionTarget], 

101 parent: nn.Module, 

102 child_name: str, 

103 child: nn.Module, 

104) -> None: 

105 """Record a replaceable reference while deduplicating shared modules.""" 

106 target = targets_by_id.setdefault(id(child), _AttentionTarget(child)) 

107 target.references.append((parent, child_name)) 

108 

109 

110def _find_attention_targets(model: nn.Module) -> list[_AttentionTarget]: 

111 """Find replaceable ``self_attn`` modules without assuming a layer path.""" 

112 targets_by_id: dict[int, _AttentionTarget] = {} 

113 visited_modules: set[int] = set() 

114 

115 def visit(parent: nn.Module) -> None: 

116 """Recursively record ``self_attn`` children, deduplicating shared modules. 

117 

118 Args: 

119 parent: Module whose children are inspected. 

120 """ 

121 parent_id = id(parent) 

122 if parent_id in visited_modules: 

123 return 

124 visited_modules.add(parent_id) 

125 

126 for child_name, child in parent.named_children(): 

127 if child_name == "self_attn": 

128 _record_attention_target(targets_by_id, parent, child_name, child) 

129 continue 

130 visit(child) 

131 

132 visit(model) 

133 targets = list(targets_by_id.values()) 

134 if not targets: 

135 raise ValueError( 

136 f"activation_swap='attention' found no self_attn modules in {type(model).__name__}; " 

137 "expected attention modules registered as self_attn" 

138 ) 

139 return targets 

140 

141 

142def apply_attention_swap(model: nn.Module, activation_swap: str) -> nn.Module: 

143 """Wrap attention modules and install model-order swap scheduling. 

144 

145 Args: 

146 model: Model being prepared, before sharding and FSDP wrapping. 

147 activation_swap: Requested activation swap mode. 

148 

149 Returns: 

150 The input model, patched in place when attention swapping is enabled. 

151 """ 

152 if activation_swap != "attention": 

153 raise ValueError( 

154 "activation_swap must be one of ('none', 'attention'), " 

155 f"got {activation_swap!r}" 

156 ) 

157 attention_targets = _find_attention_targets(model) 

158 wrapped_attentions: list[nn.Module] = [] 

159 saved_refs: list[tuple[nn.Module, str, nn.Module]] = [] 

160 try: 

161 for target in attention_targets: 

162 attention = target.module 

163 wrapped_attention = swap_wrapper( 

164 attention, 

165 policy_fn=attention_swap_policy, 

166 group_swap=True, 

167 ) 

168 for parent, child_name in target.references: 

169 saved_refs.append((parent, child_name, getattr(parent, child_name))) 

170 setattr(parent, child_name, wrapped_attention) 

171 wrapped_attentions.append(wrapped_attention) 

172 

173 swap_manager = SwapManager() 

174 for current_attention, next_attention in zip(wrapped_attentions, wrapped_attentions[1:]): 

175 swap_manager.set_forward_prefetch_layer(current_attention, next_attention) 

176 if wrapped_attentions: 

177 weakref.finalize( 

178 model, 

179 _teardown_wired_swap_layers, 

180 wrapped_attentions, 

181 ) 

182 except Exception: 

183 # Partial failure: restore the original attention modules on every 

184 # reference already re-pointed, and release any swap wiring registered 

185 # by the chain loop before it raised (idempotent when none was). 

186 _teardown_wired_swap_layers(wrapped_attentions) 

187 for parent, child_name, original in saved_refs: 

188 setattr(parent, child_name, original) 

189 raise 

190 

191 logger.info( 

192 "Enabled attention activation swap for %d module(s) in %s", 

193 len(wrapped_attentions), 

194 type(model).__name__, 

195 ) 

196 return model