Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / components / modules / engram.py: 72%
188 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-11-15 07:18 +0800
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"""Reusable DeepSeek Engram hash, sparse lookup, and fusion modules."""
17from __future__ import annotations
19from collections.abc import Mapping
20from typing import Any
22import torch # pylint: disable=forbidden-backend-import
23import torch.distributed as dist # pylint: disable=forbidden-backend-import
24from torch import nn # pylint: disable=forbidden-backend-import
26from hyper_parallel.distributed.expert_parallel.collectives import ep_all_to_all
27from hyper_parallel.models.materialization import register_rebuildable_buffer
28from hyper_parallel.models.replacement import module_replacement
31def local_token_mask(token_mask: torch.Tensor | None, hidden_states: torch.Tensor,
32 *, tp_rank: int, tp_size: int) -> torch.Tensor | None:
33 """Match a CP-local token mask to the current TP/SP hidden interval.
35 Args:
36 token_mask: CP-local mask replicated before the embedding boundary.
37 hidden_states: Current TP-local hidden tensor.
38 tp_rank: Local TP rank.
39 tp_size: Number of TP ranks.
41 Returns:
42 Mask aligned to the hidden sequence, or ``None``.
43 """
44 if token_mask is None or token_mask.shape[1] == hidden_states.shape[1]:
45 return token_mask
46 local_length = hidden_states.shape[1]
47 if token_mask.shape[1] != local_length * tp_size:
48 raise ValueError("Engram token mask is incompatible with TP/SP hidden")
49 return token_mask[:, tp_rank * local_length:(tp_rank + 1) * local_length]
52class NgramHashMapping(nn.Module):
53 """Map token IDs to one Engram layer's disjoint n-gram hash rows."""
55 DEAD_TOKEN = -1
57 def __init__(self, assets: Mapping[str, Any], layer_id: int) -> None:
58 """Create deterministic, non-persistent hash metadata.
60 Args:
61 assets: Prepared tokenizer map and per-layer hash layout.
62 layer_id: Decoder layer owning this Engram table.
63 """
64 super().__init__()
65 layer_ids = tuple(int(value) for value in assets["layer_ids"])
66 if layer_id not in layer_ids:
67 raise ValueError(f"Engram layer {layer_id} is absent from assets layer_ids={layer_ids}")
68 layer_index = layer_ids.index(layer_id)
69 self.layer_id = layer_id
70 self.max_ngram_size = int(assets["max_ngram_size"])
71 self.num_heads = int(assets["num_heads"])
72 token_map = torch.tensor(assets["token_map"], dtype=torch.long)
73 self.pad_id = int(token_map[int(assets["pad_token_id"])])
74 primes = torch.tensor(assets["primes"][layer_index], dtype=torch.long)
75 if tuple(primes.shape) != (self.max_ngram_size - 1, self.num_heads):
76 raise ValueError(
77 "Engram prime layout must be [max_ngram_size - 1, num_heads], "
78 f"got {tuple(primes.shape)}"
79 )
80 flattened_primes = primes.flatten()
81 offsets = torch.cat(
82 [torch.zeros(1, dtype=torch.long), flattened_primes.cumsum(0)[:-1]]
83 ).view_as(primes)
84 multipliers = torch.tensor(assets["multipliers"][layer_index], dtype=torch.long)
85 if multipliers.numel() != self.max_ngram_size:
86 raise ValueError("Engram multiplier count must equal max_ngram_size")
87 self.logical_num_embeddings = int(flattened_primes.sum().item())
88 initial_buffers = {
89 "token_map": token_map,
90 "primes": primes,
91 "offsets": offsets,
92 "multipliers": multipliers,
93 }
94 for name, value in initial_buffers.items():
95 register_rebuildable_buffer(self, name, value=value, persistent=False)
97 def forward(
98 self,
99 input_ids: torch.Tensor,
100 segment_starts: torch.Tensor | None = None,
101 token_mask: torch.Tensor | None = None,
102 ) -> torch.Tensor:
103 """Return global flat-table rows shaped ``[batch, seq, hash_columns]``.
105 Args:
106 input_ids: Token IDs for hash lookup.
107 segment_starts: Optional packed sequence boundaries.
108 token_mask: Optional mask excluding image tokens.
109 """
110 if input_ids.ndim != 2:
111 raise ValueError(f"Engram expects input_ids [B,S], got {tuple(input_ids.shape)}")
112 compressed = self.token_map[input_ids]
113 if token_mask is not None:
114 if token_mask.shape != input_ids.shape:
115 raise ValueError("Engram token_mask must have the same shape as input_ids")
116 compressed = torch.where(
117 token_mask.to(torch.bool),
118 compressed,
119 compressed.new_full((), self.DEAD_TOKEN),
120 )
121 batch_size, sequence_length = compressed.shape
122 positions = torch.arange(sequence_length, device=input_ids.device).expand(batch_size, -1)
123 segment_ids = None
124 if segment_starts is not None:
125 if segment_starts.shape != input_ids.shape:
126 raise ValueError("Engram segment_starts must have the same shape as input_ids")
127 starts = segment_starts.to(torch.bool).clone()
128 starts[:, 0] = True
129 segment_ids = starts.to(torch.long).cumsum(dim=1)
131 tokens = []
132 blocked = torch.zeros_like(positions, dtype=torch.bool)
133 for shift in range(self.max_ngram_size):
134 source_positions = (positions - shift).clamp_min(0)
135 source = compressed.gather(1, source_positions)
136 blocked = blocked | (positions < shift)
137 if segment_ids is not None:
138 source_segments = segment_ids.gather(1, source_positions)
139 blocked = blocked | (source_segments != segment_ids)
140 blocked = blocked | (source == self.DEAD_TOKEN)
141 tokens.append(torch.where(blocked, self.pad_id, source))
143 products = torch.stack(tokens, dim=-1) * self.multipliers
144 rolling = products[..., 0]
145 hashes = []
146 for ngram_index in range(self.max_ngram_size - 1):
147 rolling = torch.bitwise_xor(rolling, products[..., ngram_index + 1])
148 row = rolling.unsqueeze(-1) % self.primes[ngram_index]
149 hashes.append(row + self.offsets[ngram_index])
150 return torch.cat(hashes, dim=-1)
153def _all_gather_sequence(value: torch.Tensor, group: Any, group_size: int) -> torch.Tensor:
154 """Gather equal-size sequence chunks and concatenate along dimension one."""
155 gathered = [torch.empty_like(value) for _ in range(group_size)]
156 dist.all_gather(gathered, value.contiguous(), group=group)
157 return torch.cat(gathered, dim=1)
160@module_replacement
161class EngramModule(nn.Module):
162 """Structure-preserving Engram replacement with sparse EP row lookup."""
164 def __init__(
165 self,
166 *,
167 module: nn.Module,
168 module_fqn: str = "",
169 context: Mapping[str, Any] | None = None,
170 ) -> None:
171 """Reuse a source Engram module's parameters and hash mapping."""
172 super().__init__()
173 del module_fqn, context
174 required = ("hash_mapping", "embed", "wkv", "q_weight", "k_weight")
175 missing = [name for name in required if not hasattr(module, name)]
176 if missing:
177 raise TypeError(f"Engram source is missing required attributes: {missing}")
178 self.hash_mapping = module.hash_mapping
179 self.embed = module.embed
180 self.wkv = module.wkv
181 self.q_weight = module.q_weight
182 self.k_weight = module.k_weight
183 self.layer_id = int(module.layer_id)
184 self.hidden_size = int(module.hidden_size)
185 self.hc_mult = int(module.hc_mult)
186 self.eps = float(module.eps)
187 self.clamp_value = float(module.clamp_value)
188 self.logical_num_embeddings = int(module.logical_num_embeddings)
189 self.padded_num_embeddings = int(module.padded_num_embeddings)
190 self.train(module.training)
192 def checkpoint_logical_shapes(self) -> dict[str, tuple[int, ...]]:
193 """Return the logical embedding shape for checkpoint loading and checking."""
194 return {"embed.weight": (self.logical_num_embeddings, int(self.embed.weight.shape[1]))}
196 def _fuse(self, hidden_states: torch.Tensor, embeddings: torch.Tensor) -> torch.Tensor:
197 """Apply the exact V4.1 WKV projection and signed-root gate."""
198 key_value = self.wkv(embeddings.flatten(start_dim=-2))
199 key, value = key_value.split(
200 [self.hc_mult * self.hidden_size, self.hidden_size], dim=-1
201 )
202 key = key.float().unflatten(-1, (self.hc_mult, self.hidden_size))
203 hidden_fp32 = hidden_states.float()
204 weight = self.q_weight.float() * self.k_weight.float()
205 reciprocal_std = torch.rsqrt(hidden_fp32.square().mean(-1) + self.eps)
206 reciprocal_std = reciprocal_std * torch.rsqrt(key.square().mean(-1) + self.eps)
207 dot = (hidden_fp32 * weight * key).sum(-1)
208 dot = dot * reciprocal_std * self.hidden_size**-0.5
209 root = dot.abs().clamp_min(self.clamp_value).sqrt()
210 signed_root = torch.where(dot < 0, -root, root)
211 gate = torch.sigmoid(signed_root)
212 return (
213 hidden_fp32 + gate.unsqueeze(-1) * value.float().unsqueeze(-2)
214 ).to(hidden_states.dtype)
216 @staticmethod
217 def _sparse_lookup(
218 hash_ids: torch.Tensor,
219 embedding: nn.Embedding,
220 *,
221 ep_group: Any,
222 ep_rank: int,
223 ep_size: int,
224 padded_num_embeddings: int,
225 ) -> torch.Tensor:
226 """Route requested global rows to their EP owners and restore order."""
227 if padded_num_embeddings % ep_size:
228 raise ValueError(
229 f"Engram padded rows {padded_num_embeddings} are not divisible by ep_size={ep_size}"
230 )
231 num_local_rows = padded_num_embeddings // ep_size
232 row_start = ep_rank * num_local_rows
233 flat_ids = hash_ids.reshape(-1)
234 owner = torch.div(flat_ids, num_local_rows, rounding_mode="floor")
235 if owner.numel() and int(owner.max()) >= ep_size:
236 raise IndexError("Engram hash requested a row outside the padded table")
237 local_counts = torch.bincount(owner, minlength=ep_size).to(torch.long)
238 gathered_counts = [torch.empty_like(local_counts) for _ in range(ep_size)]
239 dist.all_gather(gathered_counts, local_counts, group=ep_group)
240 all_counts = torch.stack(gathered_counts)
241 send_counts = local_counts.cpu().tolist()
242 recv_counts = all_counts[:, ep_rank].cpu().tolist()
244 # Ascend ArgSort executes integer inputs on AiCPU. The owner range is
245 # bounded by ep_size, so float32 preserves it exactly and stays on AiCore.
246 sort_indices = torch.argsort(owner.float(), stable=True)
247 sorted_ids = flat_ids.index_select(0, sort_indices)
248 received_ids = ep_all_to_all(sorted_ids, send_counts, recv_counts, ep_group)
249 local_ids = received_ids - row_start
250 if local_ids.numel() and (int(local_ids.min()) < 0 or int(local_ids.max()) >= num_local_rows):
251 raise IndexError("Engram EP routing delivered a row to the wrong owner")
252 # Call the child module instead of reading ``embedding.weight``
253 # directly. This preserves the FSDP pre/post-forward lifecycle when
254 # the EP table is wrapped independently from the dense Engram gates.
255 values = embedding(local_ids)
256 returned = ep_all_to_all(values, recv_counts, send_counts, ep_group)
257 restored = torch.zeros_like(returned)
258 restored.scatter_add_(
259 0,
260 sort_indices.unsqueeze(1).expand(-1, returned.shape[1]),
261 returned,
262 )
263 return restored.view(*hash_ids.shape, embedding.embedding_dim)
265 def _aligned_hash_ids(
266 self,
267 hidden_states: torch.Tensor,
268 input_ids: torch.Tensor,
269 segment_starts: torch.Tensor | None,
270 token_mask: torch.Tensor | None = None,
271 *,
272 cp_group: Any = None,
273 cp_rank: int = 0,
274 cp_size: int = 1,
275 tp_rank: int = 0,
276 tp_size: int = 1,
277 ) -> torch.Tensor:
278 """Hash the local sequence plus the left halo required by n-grams."""
279 cp_local_ids = input_ids
280 if cp_size > 1:
281 input_ids = _all_gather_sequence(input_ids, cp_group, cp_size)
282 if segment_starts is not None:
283 segment_starts = _all_gather_sequence(segment_starts, cp_group, cp_size)
284 if token_mask is not None:
285 token_mask = _all_gather_sequence(token_mask, cp_group, cp_size)
286 local_sequence = hidden_states.shape[1]
287 cp_sequence = cp_local_ids.shape[1]
288 if cp_sequence == local_sequence:
289 sequence_parallel_parts = 1
290 local_tp_rank = 0
291 elif cp_sequence % local_sequence == 0 and cp_sequence // local_sequence == tp_size:
292 sequence_parallel_parts = tp_size
293 local_tp_rank = tp_rank
294 else:
295 raise ValueError(
296 "Engram input_ids length is incompatible with the local hidden sequence: "
297 f"ids={cp_sequence}, hidden={local_sequence}, tp_size={tp_size}"
298 )
299 del sequence_parallel_parts
300 start = cp_rank * cp_sequence + local_tp_rank * local_sequence
301 end = start + local_sequence
302 context_start = max(0, start - self.hash_mapping.max_ngram_size + 1)
303 ids_window = input_ids[:, context_start:end]
304 starts_window = None if segment_starts is None else segment_starts[:, context_start:end]
305 mask_window = None if token_mask is None else token_mask[:, context_start:end]
306 hashes = self.hash_mapping(ids_window, starts_window, mask_window)
307 target_offset = start - context_start
308 return hashes[:, target_offset:target_offset + local_sequence]
310 def forward(
311 self,
312 hidden_states: torch.Tensor,
313 input_ids: torch.Tensor,
314 segment_starts: torch.Tensor | None = None,
315 token_mask: torch.Tensor | None = None,
316 ) -> torch.Tensor:
317 """Run replicated single-card/EP1 Engram.
319 Args:
320 hidden_states: Current decoder hidden rows.
321 input_ids: CP-local token IDs for hash lookup.
322 segment_starts: Optional packed sequence boundaries.
323 token_mask: Optional mask selecting Engram updates.
324 """
325 hash_ids = self._aligned_hash_ids(hidden_states, input_ids, segment_starts, token_mask)
326 fused = self._fuse(hidden_states, self.embed(hash_ids))
327 if token_mask is None:
328 return fused
329 return torch.where(token_mask.to(torch.bool).unsqueeze(-1).unsqueeze(-1), fused, hidden_states)
331 def parallel_forward(
332 self,
333 hidden_states: torch.Tensor,
334 input_ids: torch.Tensor,
335 segment_starts: torch.Tensor | None = None,
336 token_mask: torch.Tensor | None = None,
337 *,
338 ep_group: Any,
339 ep_rank: int,
340 ep_size: int,
341 cp_group: Any = None,
342 cp_rank: int = 0,
343 cp_size: int = 1,
344 tp_rank: int = 0,
345 tp_size: int = 1,
346 ) -> torch.Tensor:
347 """Run CP/TP sequence alignment and sparse EP lookup.
349 Args:
350 hidden_states: Current decoder hidden rows.
351 input_ids: CP-local token IDs for hash construction.
352 segment_starts: Optional packed sequence boundaries.
353 token_mask: CP-local mask selecting Engram updates.
354 ep_group: Expert-parallel process group.
355 ep_rank: Local expert-parallel rank.
356 ep_size: Expert-parallel degree.
357 cp_group: Context-parallel process group.
358 cp_rank: Local context-parallel rank.
359 cp_size: Context-parallel degree.
360 tp_rank: Local tensor-parallel rank.
361 tp_size: Tensor-parallel degree.
362 """
363 hash_ids = self._aligned_hash_ids(
364 hidden_states,
365 input_ids,
366 segment_starts,
367 token_mask,
368 cp_group=cp_group,
369 cp_rank=cp_rank,
370 cp_size=cp_size,
371 tp_rank=tp_rank,
372 tp_size=tp_size,
373 )
374 embeddings = self._sparse_lookup(
375 hash_ids,
376 self.embed,
377 ep_group=ep_group,
378 ep_rank=ep_rank,
379 ep_size=ep_size,
380 padded_num_embeddings=self.padded_num_embeddings,
381 )
382 fused = self._fuse(hidden_states, embeddings)
383 token_mask = local_token_mask(token_mask, hidden_states, tp_rank=tp_rank, tp_size=tp_size)
384 if token_mask is None:
385 return fused
386 return torch.where(token_mask.to(torch.bool).unsqueeze(-1).unsqueeze(-1), fused, hidden_states)
389__all__ = ["EngramModule", "NgramHashMapping", "local_token_mask"]