Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / compile / examples / automodel_tp / mock_modules.py: 94%

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.apache.org/licenses-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"""Pure-PyTorch mock modules for GraphTrainer TP/SP demo. 

16 

17These modules are self-contained implementations of RotaryEmbedding, 

18RMSNorm, GroupQueryAttention and SwiGLUMLP. 

19 

20Submodule names (``q_proj``, ``k_proj``, ``v_proj``, ``o_proj``, 

21``gate_proj``, ``up_proj``, ``down_proj``) match the ShardingPlanner 

22conventions so ``ShardingPlanner.plan()`` derives correct TP placements 

23without any plan_overrides. 

24 

25THESE ARE FOR DEMONSTRATION ONLY. NOT to be used in real trainings. 

26""" 

27 

28import torch 

29from torch import nn 

30import torch.nn.functional as F 

31from transformers import LlamaConfig 

32 

33 

34class RotaryEmbedding(nn.Module): 

35 """Standard rotary position embedding (Transformers-compatible). 

36 

37 Computes ``inv_freq = theta^(-2i/dim)`` and caches cos/sin for the 

38 requested sequence length. 

39 """ 

40 

41 def __init__(self, head_dim: int, max_seq_len: int = 2048, theta: float = 10000.0): 

42 super().__init__() 

43 inv_freq = 1.0 / ( 

44 theta ** (torch.arange(0, head_dim, 2, dtype=torch.float32) / head_dim) 

45 ) 

46 self.register_buffer("inv_freq", inv_freq, persistent=False) 

47 

48 # Pre-compute cos/sin for max_seq_len 

49 t = torch.arange(max_seq_len, dtype=torch.float32) 

50 freqs = torch.outer(t, self.inv_freq) 

51 emb = torch.cat((freqs, freqs), dim=-1) 

52 self.register_buffer("cos_cached", emb.cos(), persistent=False) 

53 self.register_buffer("sin_cached", emb.sin(), persistent=False) 

54 

55 def forward(self, seq_len: int, device=None): 

56 """Return (cos, sin) for the first ``seq_len`` positions.""" 

57 if seq_len > self.cos_cached.shape[0]: 

58 # Extend cache on the fly 

59 t = torch.arange(seq_len, device=self.inv_freq.device, dtype=torch.float32) 

60 freqs = torch.outer(t, self.inv_freq) 

61 emb = torch.cat((freqs, freqs), dim=-1) 

62 cos = emb.cos().to(device) 

63 sin = emb.sin().to(device) 

64 return cos, sin 

65 cos = self.cos_cached[:seq_len] 

66 sin = self.sin_cached[:seq_len] 

67 if device is not None: 

68 cos = cos.to(device) 

69 sin = sin.to(device) 

70 return cos, sin 

71 

72 

73def rotate_half(x: torch.Tensor) -> torch.Tensor: 

74 """Rotate the two halves of the last dimension.""" 

75 first = x[..., : x.shape[-1] // 2] 

76 second = x[..., x.shape[-1] // 2 :] 

77 return torch.cat((-second, first), dim=-1) 

78 

79 

80def apply_rotary_pos_emb( 

81 q: torch.Tensor, 

82 k: torch.Tensor, 

83 cos: torch.Tensor, 

84 sin: torch.Tensor, 

85) -> tuple[torch.Tensor, torch.Tensor]: 

86 """Apply rotary position embedding to query and key tensors.""" 

87 # cos/sin: (seq_len, head_dim) -> (1, seq_len, 1, head_dim) 

88 cos = cos.unsqueeze(0).unsqueeze(2) 

89 sin = sin.unsqueeze(0).unsqueeze(2) 

90 q_embed = q * cos + rotate_half(q) * sin 

91 k_embed = k * cos + rotate_half(k) * sin 

92 return q_embed, k_embed 

93 

94 

95class RMSNorm(nn.Module): 

96 """Root mean square normalization (Transformers-compatible).""" 

97 

98 def __init__(self, hidden_size: int, eps: float = 1e-6): 

99 super().__init__() 

100 self.weight = nn.Parameter(torch.ones(hidden_size)) 

101 self.eps = eps 

102 

103 def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: 

104 """Normalize in fp32 and rescale with the learned weight.""" 

105 input_dtype = hidden_states.dtype 

106 hidden_states = hidden_states.to(torch.float32) 

107 variance = hidden_states.pow(2).mean(-1, keepdim=True) 

108 hidden_states = hidden_states * torch.rsqrt(variance + self.eps) 

109 return self.weight * hidden_states.to(input_dtype) 

110 

111 

112class GroupQueryAttention(nn.Module): 

113 """GQA with rotary embedding inside forward (SP-safe design). 

114 

115 Submodule names (``q_proj``, ``k_proj``, ``v_proj``, ``o_proj``) match 

116 the ShardingPlanner conventions. 

117 """ 

118 

119 def __init__(self, config: LlamaConfig): 

120 super().__init__() 

121 self.hidden_size = config.hidden_size 

122 self.num_heads = config.num_attention_heads 

123 self.num_kv_heads = config.num_key_value_heads 

124 self.head_dim = self.hidden_size // self.num_heads 

125 self.num_kv_groups = self.num_heads // self.num_kv_heads 

126 

127 self.q_proj = nn.Linear( 

128 self.hidden_size, self.num_heads * self.head_dim, bias=False 

129 ) 

130 self.k_proj = nn.Linear( 

131 self.hidden_size, self.num_kv_heads * self.head_dim, bias=False 

132 ) 

133 self.v_proj = nn.Linear( 

134 self.hidden_size, self.num_kv_heads * self.head_dim, bias=False 

135 ) 

136 self.o_proj = nn.Linear( 

137 self.num_heads * self.head_dim, self.hidden_size, bias=False 

138 ) 

139 

140 self.rotary_emb = RotaryEmbedding( 

141 self.head_dim, 

142 max_seq_len=config.max_position_embeddings, 

143 theta=getattr(config, "rope_theta", 10000.0), 

144 ) 

145 

146 def _expand_kv_heads(self, tensor: torch.Tensor) -> torch.Tensor: 

147 """Repeat KV heads to match the number of query heads. 

148 

149 Input shape: (batch, seq_len, num_kv_heads, head_dim) 

150 Output shape: (batch, seq_len, num_heads, head_dim) 

151 """ 

152 batch, seq_len, num_kv, head_dim = tensor.shape 

153 tensor = tensor.unsqueeze(3).expand( 

154 batch, seq_len, num_kv, self.num_kv_groups, head_dim 

155 ) 

156 return tensor.reshape(batch, seq_len, num_kv * self.num_kv_groups, head_dim) 

157 

158 # pylint: disable=W0613 

159 def forward(self, hidden_states: torch.Tensor, position_ids=None): 

160 """Forward: QKV proj -> RoPE (inside forward) -> attention -> output.""" 

161 batch, seq_len, _ = hidden_states.shape 

162 

163 query = self.q_proj(hidden_states).view( 

164 batch, seq_len, self.num_heads, self.head_dim 

165 ) 

166 key = self.k_proj(hidden_states).view( 

167 batch, seq_len, self.num_kv_heads, self.head_dim 

168 ) 

169 value = self.v_proj(hidden_states).view( 

170 batch, seq_len, self.num_kv_heads, self.head_dim 

171 ) 

172 

173 # RoPE computed inside attention (after SP boundary all-gather) 

174 cos, sin = self.rotary_emb(seq_len, device=hidden_states.device) 

175 query, key = apply_rotary_pos_emb(query, key, cos, sin) 

176 

177 # Expand KV heads to match query heads 

178 key = self._expand_kv_heads(key) 

179 value = self._expand_kv_heads(value) 

180 

181 # (batch, num_heads, seq_len, head_dim) 

182 query = query.transpose(1, 2) 

183 key = key.transpose(1, 2) 

184 value = value.transpose(1, 2) 

185 

186 # Scaled dot-product attention 

187 scale = self.head_dim**-0.5 

188 attn_weights = torch.matmul(query, key.transpose(-2, -1)) * scale 

189 # Causal mask 

190 causal_mask = torch.triu( 

191 torch.ones(seq_len, seq_len, device=hidden_states.device, dtype=torch.bool), 

192 diagonal=1, 

193 ) 

194 attn_weights = attn_weights.masked_fill(causal_mask, float("-inf")) 

195 attn_weights = F.softmax(attn_weights, dim=-1, dtype=torch.float32).to( 

196 query.dtype 

197 ) 

198 attn_output = torch.matmul(attn_weights, value) 

199 

200 attn_output = attn_output.transpose(1, 2).reshape(batch, seq_len, -1) 

201 return self.o_proj(attn_output) 

202 

203 

204class SwiGLUMLP(nn.Module): 

205 """SwiGLU MLP with separate gate/up/down projections. 

206 

207 Submodule names (``gate_proj``, ``up_proj``, ``down_proj``) match 

208 the ShardingPlanner conventions. 

209 """ 

210 

211 def __init__(self, config: LlamaConfig): 

212 super().__init__() 

213 self.gate_proj = nn.Linear( 

214 config.hidden_size, config.intermediate_size, bias=False 

215 ) 

216 self.up_proj = nn.Linear( 

217 config.hidden_size, config.intermediate_size, bias=False 

218 ) 

219 self.down_proj = nn.Linear( 

220 config.intermediate_size, config.hidden_size, bias=False 

221 ) 

222 

223 def forward(self, x: torch.Tensor) -> torch.Tensor: 

224 """SwiGLU: down(silu(gate(x)) * up(x)).""" 

225 return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))