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

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"""Model wrapper for GraphTrainer TP/SP demo. 

16 

17Self-contained pure-PyTorch model that reuses AutoModel's **pipeline** 

18(``ShardingPlanner.plan`` + ``apply_sharding_plan``) while keeping the 

19module implementations local to avoid NPU-only functional dependencies. 

20 

21Key design property: ``rotary_emb`` lives inside ``Attention.forward`` 

22(not at the model level), so cos/sin are computed *after* the SP boundary 

23all-gather — they naturally have the full sequence length. 

24 

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

26``gate_proj``, ``up_proj``, ``down_proj``, ``embed_tokens``, ``lm_head``) 

27match the ShardingPlanner conventions so ``ShardingPlanner.plan()`` 

28derives correct TP placements without any plan_overrides. 

29 

30Public API: 

31 - ``AutoModelAdapterForCausalLM`` -- CausalLM model 

32 - ``build_model`` -- factory: create model from config dict 

33 - ``shard_for_sp`` -- shard input tensor along sequence dim for SP 

34 - ``DataSampler`` -- data sampler that yields full-sequence batches 

35""" 

36 

37import torch 

38from torch import nn 

39import torch.nn.functional as F 

40from transformers import LlamaConfig 

41 

42from hyper_parallel.compile.examples.automodel_tp.mock_modules import ( 

43 GroupQueryAttention, 

44 RMSNorm, 

45 SwiGLUMLP, 

46) 

47 

48 

49# ============================================================================ 

50# Model assembly 

51# ============================================================================ 

52 

53 

54class AutoModelAdapterDecoderLayer(nn.Module): 

55 """Decoder layer: attention + MLP with pre-norm residuals.""" 

56 

57 def __init__(self, config: LlamaConfig): 

58 super().__init__() 

59 self.self_attn = GroupQueryAttention(config) 

60 self.mlp = SwiGLUMLP(config) 

61 self.input_layernorm = RMSNorm(config.hidden_size) 

62 self.post_attention_layernorm = RMSNorm(config.hidden_size) 

63 

64 def forward(self, hidden_states, position_ids=None): 

65 """Standard pre-norm decoder layer forward: attn + MLP with residuals.""" 

66 residual = hidden_states 

67 hidden_states = self.input_layernorm(hidden_states) 

68 hidden_states = self.self_attn(hidden_states, position_ids=position_ids) 

69 hidden_states = residual + hidden_states 

70 

71 residual = hidden_states 

72 hidden_states = self.post_attention_layernorm(hidden_states) 

73 hidden_states = self.mlp(hidden_states) 

74 hidden_states = residual + hidden_states 

75 return hidden_states 

76 

77 

78class AutoModelAdapterModel(nn.Module): 

79 """Inner model: embed_tokens, layers, norm. 

80 

81 No rotary_emb at this level -- it lives inside each 

82 ``GroupQueryAttention`` (SP-safe design). 

83 """ 

84 

85 def __init__(self, config: LlamaConfig): 

86 super().__init__() 

87 self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size) 

88 self.layers = nn.ModuleList( 

89 [ 

90 AutoModelAdapterDecoderLayer(config) 

91 for _ in range(config.num_hidden_layers) 

92 ] 

93 ) 

94 self.norm = RMSNorm(config.hidden_size) 

95 

96 def forward(self, input_ids, position_ids=None): 

97 """Embed, run each decoder layer, then apply the final norm.""" 

98 hidden_states = self.embed_tokens(input_ids) 

99 for layer in self.layers: 

100 hidden_states = layer(hidden_states, position_ids=position_ids) 

101 return self.norm(hidden_states) 

102 

103 

104class AutoModelAdapterForCausalLM(nn.Module): 

105 """CausalLM model for GraphTrainer TP/SP demo. 

106 

107 Reuses AutoModel's ``ShardingPlanner`` + ``apply_sharding_plan`` 

108 pipeline. The model itself uses pure-PyTorch modules to ensure 

109 CPU/NPU compatibility. 

110 

111 Submodule layout (matches ShardingPlanner conventions): 

112 model.embed_tokens.weight 

113 model.layers.{i}.self_attn.{q,k,v,o}_proj 

114 model.layers.{i}.mlp.{gate,up,down}_proj 

115 model.layers.{i}.{input,post_attention}_layernorm 

116 model.norm 

117 lm_head.weight 

118 """ 

119 

120 def __init__(self, config: LlamaConfig): 

121 super().__init__() 

122 self.config = config 

123 self.model = AutoModelAdapterModel(config) 

124 self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) 

125 if getattr(config, "tie_word_embeddings", False): 

126 self.lm_head.weight = self.model.embed_tokens.weight 

127 

128 def forward(self, input_ids, labels=None, position_ids=None): 

129 """Forward: embeddings -> decoder layers -> norm -> lm_head -> loss. 

130 

131 Args: 

132 input_ids: token IDs of shape (batch, seq_len). 

133 labels: optional target token IDs for CE loss. 

134 position_ids: optional; if None, derived from seq_len. 

135 

136 Returns: 

137 Dict with ``loss`` (or None) and ``logits``. 

138 """ 

139 seq_len = input_ids.shape[1] 

140 if position_ids is None: 

141 position_ids = torch.arange(seq_len, device=input_ids.device).unsqueeze(0) 

142 

143 hidden_states = self.model(input_ids, position_ids=position_ids) 

144 logits = self.lm_head(hidden_states) 

145 

146 loss = None 

147 if labels is not None: 

148 shift_logits = logits[..., :-1, :].contiguous() 

149 shift_labels = labels[..., 1:].contiguous() 

150 loss = F.cross_entropy( 

151 shift_logits.view(-1, shift_logits.size(-1)), 

152 shift_labels.view(-1), 

153 ignore_index=-100, 

154 ) 

155 return {"loss": loss, "logits": logits} 

156 

157 

158# ============================================================================ 

159# SP data utilities 

160# ============================================================================ 

161 

162 

163def shard_for_sp(tensor, tp_size, tp_rank): 

164 """Shard input tensor along sequence dimension (dim 1) for SP. 

165 

166 In SP mode, each TP rank receives a 1/tp_size slice of the sequence. 

167 The boundary forward's in_plan (Shard(1) -> Replicate()) will 

168 all-gather it back to full sequence inside the model. 

169 """ 

170 seq_len = tensor.shape[1] 

171 chunk = seq_len // tp_size 

172 if seq_len % tp_size != 0: 

173 raise ValueError( 

174 f"Sequence length {seq_len} not divisible by tp_size {tp_size}" 

175 ) 

176 return tensor[:, tp_rank * chunk : (tp_rank + 1) * chunk].contiguous() 

177 

178 

179def build_model(cfg: dict, device: torch.device) -> torch.nn.Module: 

180 """Factory: create a CausalLM from config dict.""" 

181 model_cfg = LlamaConfig( 

182 vocab_size=cfg["vocab_size"], 

183 hidden_size=cfg["hidden_size"], 

184 intermediate_size=cfg["intermediate_size"], 

185 num_hidden_layers=cfg["num_hidden_layers"], 

186 num_attention_heads=cfg["num_attention_heads"], 

187 num_key_value_heads=cfg["num_key_value_heads"], 

188 max_position_embeddings=cfg["max_position_embeddings"], 

189 torch_dtype=cfg.get("torch_dtype", "float32"), 

190 tie_word_embeddings=False, 

191 ) 

192 return AutoModelAdapterForCausalLM(model_cfg).to(device) 

193 

194 

195class DataSampler: 

196 """Yield dummy batches for training. 

197 

198 In SP mode, yields full-sequence data; SP sharding is handled by 

199 the embedding boundary's reduce-scatter on hidden_states. 

200 """ 

201 

202 def __init__( 

203 self, 

204 vocab_size: int, 

205 batch_size: int, 

206 seq_len: int, 

207 max_steps: int, 

208 tp_size: int, 

209 tp_rank: int, 

210 sequence_parallel: bool, 

211 device: torch.device, 

212 ): 

213 self.vocab_size = vocab_size 

214 self.batch_size = batch_size 

215 self.seq_len = seq_len 

216 self.max_steps = max_steps 

217 self.tp_size = tp_size 

218 self.tp_rank = tp_rank 

219 self.sequence_parallel = sequence_parallel 

220 self.device = device 

221 

222 def sample(self): 

223 """Return one (input_ids, labels) pair with full sequence length.""" 

224 inp = torch.randint( 

225 0, self.vocab_size, (self.batch_size, self.seq_len), device=self.device 

226 ) 

227 lbl = torch.randint( 

228 0, self.vocab_size, (self.batch_size, self.seq_len), device=self.device 

229 ) 

230 return inp, lbl 

231 

232 def __iter__(self): 

233 for _ in range(self.max_steps): 

234 yield self.sample()