| """Granite-1B-A400M 的 Looped × AttnRes 改造(消融实验 G0-G3 共用一份代码)。 |
| |
| 开关(通过 from_pretrained 的 kwargs 直接传给 __init__): |
| - loop_rounds: int = 1 循环轮数; 1 = G0 基线 |
| - attnres: bool = False 轮内 Block AttnRes (G2+) |
| - cross_round: bool = False 源集合加入之前所有轮次的块摘要 (G3) |
| - attnres_blocks: int = 8 AttnRes 块数 (48 子层 / 8 块 = 6 子层每块) |
| |
| 结构要点(依据 DESIGN.md 与两篇论文): |
| - 循环: 24 层 × R 轮, 权重跨轮硬共享; 每轮输出过共享 lm_head 算 CE, 取各轮均值 |
| - AttnRes: 每子层(attn/MoE 各算一个)一个零初始化可学习伪查询, 按轮独立; |
| key=value=子层输出, 核 exp(q·RMSNorm(k)) (RMSNorm 无可学习参数), softmax 在源维度 |
| - residual_multiplier μ=0.22 折进 value (v_i = μ·f_i) |
| - 块内: 部分和按标准残差累加; 块边界: 块摘要加入源集合 |
| - 每轮末尾: 用独立的终读出伪查询对 [b0] + (跨轮时含历史块) + 本轮块摘要 做最后读出 |
| - 跨轮(G3): 源 = [b0(embedding)] + [之前所有轮次的块摘要] + [本轮已完成块] + [本块部分和] |
| |
| 注意: 训练前向 (input_ids/labels) 不支持 KV cache; 推理用 greedy_generate(手动贪心 + KV cache)。 |
| """ |
| import torch |
| import torch.nn.functional as F |
| from torch import nn |
| from torch.utils.checkpoint import checkpoint |
| from types import SimpleNamespace |
|
|
| from transformers.models.granitemoe.configuration_granitemoe import GraniteMoeConfig |
| from transformers.models.granitemoe.modeling_granitemoe import ( |
| GraniteMoeForCausalLM, |
| load_balancing_loss_func, |
| ) |
| from transformers.masking_utils import create_causal_mask |
|
|
|
|
| def _rms_norm(x, eps=1e-6): |
| """无可学习参数的 RMSNorm(AttnRes 核里用在 key 上)。""" |
| return x * torch.rsqrt(x.float().pow(2).mean(-1, keepdim=True) + eps).to(x.dtype) |
|
|
|
|
| class GraniteLoopForCausalLM(GraniteMoeForCausalLM): |
| def __init__(self, config: GraniteMoeConfig, loop_rounds: int = 1, attnres: bool = False, |
| cross_round: bool = False, attnres_blocks: int = 8, **kwargs): |
| super().__init__(config) |
| |
| self.loop_rounds = getattr(config, "loop_rounds", None) or loop_rounds |
| self.use_attnres = getattr(config, "attnres", None) |
| if self.use_attnres is None: |
| self.use_attnres = attnres |
| self.cross_round = getattr(config, "cross_round", None) |
| if self.cross_round is None: |
| self.cross_round = cross_round |
| self.attnres_blocks = getattr(config, "attnres_blocks", None) or attnres_blocks |
| |
| config.loop_rounds = self.loop_rounds |
| config.attnres = self.use_attnres |
| config.cross_round = self.cross_round |
| config.attnres_blocks = self.attnres_blocks |
| L = config.num_hidden_layers |
| assert (2 * L) % self.attnres_blocks == 0, "子层数需被块数整除" |
| self.sub_per_block = (2 * L) // self.attnres_blocks |
| assert self.sub_per_block % 2 == 0 |
| self.layers_per_block = self.sub_per_block // 2 |
|
|
| d = config.hidden_size |
| self.alpha_recorder = None |
| self.round_recorder = None |
| if self.use_attnres: |
| |
| |
| |
| self.pseudo_queries = nn.Parameter( |
| torch.zeros(self.loop_rounds, 2 * L + 1, d) |
| ) |
|
|
| @torch.no_grad() |
| def reset_pseudo_queries(self): |
| """from_pretrained 后必须调用: transformers 5 的加载管线会把 checkpoint 里 |
| 缺失的 pseudo_queries 重置为未初始化内存(可能含 NaN), 这里重新归零, |
| 恢复论文要求的零初始化(初始 = 等权平均)。自己 save_pretrained 的 |
| checkpoint 含训练后的值, 加载后无需调用。""" |
| if self.use_attnres: |
| self.pseudo_queries.zero_() |
|
|
| |
| def _sublayer(self, layer, kind, h, mask, pos_emb): |
| if kind == "attn": |
| return layer.self_attn( |
| hidden_states=layer.input_layernorm(h), |
| attention_mask=mask, |
| position_embeddings=pos_emb, |
| )[0] |
| |
| |
| moe = layer.block_sparse_moe |
| normed = layer.post_attention_layernorm(h) |
| b, t, d = normed.shape |
| flat = normed.reshape(-1, d) |
| top_idx, top_w, rlogits = moe.router(flat) |
| out = moe.experts(flat, top_idx, top_w).view(b, t, d) |
| return out, rlogits |
|
|
| def _aggregate(self, q, sources): |
| """AttnRes 聚合: q (d,), sources list of (b,s,d) -> (b,s,d)""" |
| V = torch.stack(sources) |
| logits = torch.einsum("d,Sbtd->Sbt", q, _rms_norm(V)) |
| alpha = logits.softmax(dim=0) |
| return torch.einsum("Sbt,Sbtd->btd", alpha, V) |
|
|
| |
| |
| |
| def _layer_fwd(self, layer, h, mask, pos_emb, mu): |
| """非 AttnRes 路径: 一个完整层(attn + moe), 返回 (h, router_logits)。""" |
| f = self._sublayer(layer, "attn", h, mask, pos_emb) |
| h = h + mu * f |
| f, rl = self._sublayer(layer, "moe", h, mask, pos_emb) |
| return h + mu * f, rl |
|
|
| def _attnres_step(self, layer, kind, q, cat, inv_rms, block_sum, mask, pos_emb, mu): |
| """AttnRes 路径: 聚合 + 单个子层 + 块内累加, 返回 (新 block_sum, router_logits 或 None)。 |
| |
| 性能优化版聚合(profile 显示旧的"每子层 stack 全部源 + 对全部源做 RMSNorm" |
| 贡献了 ~60% 的 CUDA 时间, 全是 elementwise/cat): |
| - cat: (S,b,t,d) 块内固定源的堆叠, 每块只 stack 一次(块级缓存) |
| - inv_rms: (S,b,t) 各固定源的 RMSNorm 系数, 利用 q·RMSNorm(k) = (q·k)·inv_rms |
| - softmax 拆成 [固定源 | block_sum] 两段, 每子层只需对 block_sum 单张量求范数 |
| 数学上与 _aggregate 完全等价(仅浮点运算次序不同)。""" |
| logits = torch.einsum("d,sbtd->sbt", q, cat) * inv_rms |
| if block_sum is not None: |
| bs_inv = torch.rsqrt(block_sum.float().pow(2).mean(-1) + 1e-6) |
| bs_logit = torch.einsum("d,btd->bt", q, block_sum) * bs_inv |
| alpha = torch.cat([logits, bs_logit.unsqueeze(0)], dim=0).softmax(dim=0) |
| else: |
| alpha = logits.softmax(dim=0) |
| |
| |
| rec = getattr(self, "alpha_recorder", None) |
| if rec is not None: |
| with torch.no_grad(): |
| rec.append((self._rec_ctx, alpha.detach().mean(dim=(1, 2)).cpu())) |
| alpha = alpha.to(cat.dtype) |
| if block_sum is not None: |
| h = torch.einsum("sbt,sbtd->btd", alpha[:-1], cat) + alpha[-1].unsqueeze(-1) * block_sum |
| else: |
| h = torch.einsum("sbt,sbtd->btd", alpha, cat) |
| f = self._sublayer(layer, kind, h, mask, pos_emb) |
| rl = None |
| if kind == "moe": |
| f, rl = f |
| upd = mu * f |
| return (upd if block_sum is None else block_sum + upd), rl |
|
|
| def _run_block(self, blk_layers, q_blk, fixed, mask, pos_emb, mu): |
| """一个 AttnRes 块(3 层 / 6 个子层), 作为一个 checkpoint 段。 |
| |
| 源堆叠 cat 与 inv_rms 在段内构建(而非段外传入): 段内张量在前向结束后即释放, |
| 反向重算时重建, 避免 32 个块的 cat(~0.5GB each)同时常驻显存。 |
| profile 显示旧的"每子层 stack 全部源 + 对全部源 RMSNorm"贡献 ~60% CUDA 时间, |
| 这里每块只 stack/归一化一次, 块内 6 个子层复用, 并用 |
| q·RMSNorm(k) = (q·k)·inv_rms 把归一化折成标量乘法(数学上与 _aggregate 等价)。""" |
| cat = torch.stack(fixed) |
| inv_rms = torch.rsqrt(cat.float().pow(2).mean(-1) + 1e-6) |
| block_sum = None |
| rlogits = [] |
| for j, layer in enumerate(blk_layers): |
| for kind in ("attn", "moe"): |
| q = q_blk[2 * j] if kind == "attn" else q_blk[2 * j + 1] |
| self._rec_ctx = (*self._rec_ctx_base, j, kind) |
| block_sum, rl = self._attnres_step( |
| layer, kind, q, cat, inv_rms, block_sum, mask, pos_emb, mu) |
| if rl is not None: |
| rlogits.append(rl) |
| return block_sum, rlogits |
|
|
| def _ce(self, h, labels): |
| """单轮 CE, 单独成段: 避免 4 轮的 logits(fp32 上转后每轮 ~1.2GB)常驻反传图。""" |
| logits = self.lm_head(self.model.norm(h)) / self.logits_scaling |
| return self.loss_function(logits, labels, vocab_size=self.config.vocab_size) |
|
|
| |
| def _run_round(self, h, round_idx, b0, prev_blocks, mask, pos_emb, mu, use_ckpt): |
| """返回 (本轮最终 hidden, 块摘要列表, 各 MoE 层的 router_logits)。""" |
| layers = self.model.layers |
| rlogits = [] |
| if not self.use_attnres: |
| for layer in layers: |
| if use_ckpt: |
| h, rl = checkpoint( |
| lambda hh, l=layer: self._layer_fwd(l, hh, mask, pos_emb, mu), |
| h, use_reentrant=False) |
| else: |
| h, rl = self._layer_fwd(layer, h, mask, pos_emb, mu) |
| rlogits.append(rl) |
| block_sums = [] |
| else: |
| fixed = [b0] + (list(prev_blocks) if self.cross_round else []) |
| block_sums = [] |
| for blk in range(self.attnres_blocks): |
| blk_layers = layers[blk * self.layers_per_block:(blk + 1) * self.layers_per_block] |
| q_blk = self.pseudo_queries[round_idx, |
| blk * self.sub_per_block:(blk + 1) * self.sub_per_block] |
| fix_snap = tuple(fixed) |
| self._rec_ctx_base = (round_idx, blk) |
| if use_ckpt: |
| block_sum, rls = checkpoint( |
| lambda _q, bl=blk_layers, qb=q_blk, fs=fix_snap: |
| self._run_block(bl, qb, fs, mask, pos_emb, mu), |
| q_blk, use_reentrant=False) |
| else: |
| block_sum, rls = self._run_block(blk_layers, q_blk, fix_snap, mask, pos_emb, mu) |
| rlogits.extend(rls) |
| fixed.append(block_sum) |
| block_sums.append(block_sum) |
| |
| |
| final_sources = [b0] + (list(prev_blocks) if self.cross_round else []) + block_sums |
| q_final = self.pseudo_queries[round_idx, 2 * len(layers)] |
| if use_ckpt: |
| h = checkpoint( |
| lambda _q, qq=q_final, src=tuple(final_sources): self._aggregate(qq, list(src)), |
| q_final, use_reentrant=False) |
| else: |
| h = self._aggregate(q_final, final_sources) |
| return h, block_sums, rlogits |
|
|
| |
| def forward(self, input_ids=None, labels=None, attention_mask=None, **kwargs): |
| model = self.model |
| bsz, seq = input_ids.shape |
| position_ids = torch.arange(seq, device=input_ids.device).unsqueeze(0) |
| h = model.embed_tokens(input_ids) * model.embedding_multiplier |
| mask = create_causal_mask( |
| config=self.config, inputs_embeds=h, attention_mask=attention_mask, |
| past_key_values=None, position_ids=position_ids, |
| ) |
| pos_emb = model.rotary_emb(h, position_ids) |
| mu = self.config.residual_multiplier |
| b0 = h |
|
|
| use_ckpt = self.training and torch.is_grad_enabled() |
| ce_losses = [] |
| all_rlogits = [] |
| prev_blocks = [] |
| for r in range(self.loop_rounds): |
| h, round_blocks, rlogits = self._run_round( |
| h, r, b0, prev_blocks, mask, pos_emb, mu, use_ckpt) |
| if self.round_recorder is not None: |
| self.round_recorder.append(h.detach()) |
| if self.cross_round: |
| prev_blocks = prev_blocks + round_blocks |
| if labels is not None: |
| if use_ckpt: |
| ce = checkpoint(lambda hh=h, lb=labels: self._ce(hh, lb), |
| h, use_reentrant=False) |
| else: |
| ce = self._ce(h, labels) |
| ce_losses.append(ce) |
| all_rlogits.extend(rlogits) |
|
|
| loss = None |
| aux = None |
| round_losses = None |
| logits = None |
| if labels is None: |
| |
| logits = self.lm_head(self.model.norm(h)) / self.logits_scaling |
| if labels is not None: |
| loss = torch.stack(ce_losses).mean() |
| round_losses = [c.detach() for c in ce_losses] |
| if all_rlogits: |
| aux = load_balancing_loss_func( |
| tuple(all_rlogits), self.num_experts, self.num_experts_per_tok) |
| loss = loss + self.router_aux_loss_coef * aux |
| return SimpleNamespace(loss=loss, logits=logits, aux_loss=aux, round_losses=round_losses) |
|
|
| |
| |
| |
| |
| |
| @torch.no_grad() |
| def _prefill_kv(self, input_ids): |
| """跑一遍标准前向, 返回 (最后一位的 logits, cache[r][l]=(K,V))。""" |
| captured = [] |
| hooks = [] |
| for layer in self.model.layers: |
| hooks.append(layer.self_attn.register_forward_pre_hook( |
| lambda m, args, kwargs: captured.append(kwargs["hidden_states"].detach()), |
| with_kwargs=True)) |
| try: |
| out = self.forward(input_ids=input_ids) |
| finally: |
| for hk in hooks: |
| hk.remove() |
| L = len(self.model.layers) |
| assert len(captured) == self.loop_rounds * L, \ |
| f"hook 捕获数 {len(captured)} != 轮数×层数 {self.loop_rounds * L}" |
| from transformers.models.granitemoe.modeling_granitemoe import apply_rotary_pos_emb |
| bsz, seq = input_ids.shape |
| pos_ids = torch.arange(seq, device=input_ids.device).unsqueeze(0) |
| dummy = torch.zeros(1, seq, self.config.hidden_size, device=input_ids.device, |
| dtype=self.model.embed_tokens.weight.dtype) |
| cos, sin = self.model.rotary_emb(dummy, pos_ids) |
| cache = [] |
| for r in range(self.loop_rounds): |
| rc = [] |
| for l, layer in enumerate(self.model.layers): |
| x = captured[r * L + l] |
| sa = layer.self_attn |
| hd, kvh = sa.head_dim, self.config.num_key_value_heads |
| k = sa.k_proj(x).view(bsz, seq, kvh, hd).transpose(1, 2) |
| v = sa.v_proj(x).view(bsz, seq, kvh, hd).transpose(1, 2) |
| _, k = apply_rotary_pos_emb( |
| torch.zeros(bsz, self.config.num_attention_heads, seq, hd, |
| device=x.device, dtype=k.dtype), k, cos, sin) |
| rc.append((k, v)) |
| cache.append(rc) |
| return out.logits[:, -1], cache |
|
|
| def _attn_decode(self, layer, x, k_cache, v_cache, cos, sin): |
| """解码阶段的单 token 注意力: x (b,1,d) 已加 cache 的 K/V (b,kvh,t,hd)。""" |
| from transformers.models.granitemoe.modeling_granitemoe import apply_rotary_pos_emb |
| sa = layer.self_attn |
| b = x.shape[0] |
| nh, kvh, hd = self.config.num_attention_heads, self.config.num_key_value_heads, sa.head_dim |
| q = sa.q_proj(x).view(b, 1, nh, hd).transpose(1, 2) |
| k = sa.k_proj(x).view(b, 1, kvh, hd).transpose(1, 2) |
| v = sa.v_proj(x).view(b, 1, kvh, hd).transpose(1, 2) |
| q, k = apply_rotary_pos_emb(q, k, cos, sin) |
| K = torch.cat([k_cache, k], dim=2) |
| V = torch.cat([v_cache, v], dim=2) |
| g = nh // kvh |
| Kf = K.repeat_interleave(g, dim=1) |
| Vf = V.repeat_interleave(g, dim=1) |
| w = (q @ Kf.transpose(-1, -2)) * sa.scaling |
| o = w.softmax(-1) @ Vf |
| o = o.transpose(1, 2).reshape(b, 1, nh * hd) |
| return sa.o_proj(o), K, V |
|
|
| @torch.no_grad() |
| def greedy_generate(self, input_ids, max_new_tokens=64, eos_token_id=None): |
| """KV cache 贪心生成, 返回生成的 id 列表(不含 prompt)。仅支持 batch=1。""" |
| assert input_ids.shape[0] == 1 |
| dev = input_ids.device |
| layers = self.model.layers |
| mu = self.config.residual_multiplier |
| L = len(layers) |
|
|
| logits, cache = self._prefill_kv(input_ids) |
| nxt = logits[0].argmax().item() |
| out_ids = [] |
| pos = input_ids.shape[1] |
|
|
| for _ in range(max_new_tokens): |
| if eos_token_id is not None and nxt == eos_token_id: |
| break |
| out_ids.append(nxt) |
| pos_ids = torch.tensor([[pos]], device=dev) |
| cos, sin = self.model.rotary_emb( |
| torch.zeros(1, 1, self.config.hidden_size, device=dev, |
| dtype=self.model.embed_tokens.weight.dtype), pos_ids) |
| h = self.model.embed_tokens( |
| torch.tensor([[nxt]], device=dev)) * self.model.embedding_multiplier |
| b0 = h |
| prev_blocks = [] |
| for r in range(self.loop_rounds): |
| if not self.use_attnres: |
| for l, layer in enumerate(layers): |
| f, K, V = self._attn_decode( |
| layer, layer.input_layernorm(h), cache[r][l][0], cache[r][l][1], cos, sin) |
| cache[r][l] = (K, V) |
| h = h + mu * f |
| f, _ = self._sublayer(layer, "moe", h, None, None) |
| h = h + mu * f |
| else: |
| fixed = [b0] + (list(prev_blocks) if self.cross_round else []) |
| block_sums = [] |
| for blk in range(self.attnres_blocks): |
| blk_layers = layers[blk * self.layers_per_block:(blk + 1) * self.layers_per_block] |
| q_blk = self.pseudo_queries[r, |
| blk * self.sub_per_block:(blk + 1) * self.sub_per_block] |
| cat = torch.stack(fixed) |
| inv_rms = torch.rsqrt(cat.float().pow(2).mean(-1) + 1e-6) |
| block_sum = None |
| for j, layer in enumerate(blk_layers): |
| for kind in ("attn", "moe"): |
| q = q_blk[2 * j] if kind == "attn" else q_blk[2 * j + 1] |
| block_sum, _ = self._attnres_step_decode( |
| layer, kind, q, cat, inv_rms, block_sum, cos, sin, mu, |
| cache, r, blk * self.layers_per_block + j) |
| fixed.append(block_sum) |
| block_sums.append(block_sum) |
| final_sources = [b0] + (list(prev_blocks) if self.cross_round else []) + block_sums |
| q_final = self.pseudo_queries[r, 2 * L] |
| h = self._aggregate(q_final, final_sources) |
| if self.cross_round: |
| prev_blocks = prev_blocks + block_sums |
| logits = self.lm_head(self.model.norm(h)) / self.logits_scaling |
| nxt = logits[0, -1].argmax().item() |
| pos += 1 |
| return out_ids |
|
|
| def _attnres_step_decode(self, layer, kind, q, cat, inv_rms, block_sum, cos, sin, mu, |
| cache, r, l): |
| """_attnres_step 的解码版: 聚合 + 单个子层(单 token, 注意力走 cache)。""" |
| logits = torch.einsum("d,sbtd->sbt", q, cat) * inv_rms |
| if block_sum is not None: |
| bs_inv = torch.rsqrt(block_sum.float().pow(2).mean(-1) + 1e-6) |
| bs_logit = torch.einsum("d,btd->bt", q, block_sum) * bs_inv |
| alpha = torch.cat([logits, bs_logit.unsqueeze(0)], dim=0).softmax(dim=0) |
| else: |
| alpha = logits.softmax(dim=0) |
| alpha = alpha.to(cat.dtype) |
| if block_sum is not None: |
| h = torch.einsum("sbt,sbtd->btd", alpha[:-1], cat) + alpha[-1].unsqueeze(-1) * block_sum |
| else: |
| h = torch.einsum("sbt,sbtd->btd", alpha, cat) |
| if kind == "attn": |
| f, K, V = self._attn_decode( |
| layer, layer.input_layernorm(h), cache[r][l][0], cache[r][l][1], cos, sin) |
| cache[r][l] = (K, V) |
| rl = None |
| else: |
| f, rl = self._sublayer(layer, "moe", h, None, None) |
| upd = mu * f |
| return (upd if block_sum is None else block_sum + upd), rl |
|
|
