Update modeling_densebackward_olmoe0125.py

modeling_densebackward_olmoe0125.py

@@ -6,15 +6,46 @@ import torch.nn.functional as F
 
 # 导入官方实现（注意根据你的 transformers 版本调整导入路径）
 from transformers.models.olmoe.modeling_olmoe import OlmoeForCausalLM, OlmoeSparseMoeBlock, OlmoeMLP
-from .configuration_densebackward_olmoe0125 import DenseBackwardOLMoEConfig
+from configuration_custom import DenseBackwardOLMoEConfig
 
 
 class DenseBackwardOlmoeSparseMoeBlock(OlmoeSparseMoeBlock):
+    """
+    继承自官方 OlmoeSparseMoeBlock，实现 dense backward 功能：
+    前向输出依旧保持与官方相同（即稀疏计算结果），
+    但在反向传播时，通过直通梯度让 dense 计算的梯度传递回来，
+    dense 输出通过对每个专家在所有 token 上进行计算，并利用全 routing 权重加权获得。
+    
+    输入：
+        hidden_states: Tensor, shape (batch_size, sequence_length, hidden_dim)
+    输出：
+        final_output: Tensor, shape (batch_size, sequence_length, hidden_dim)
+        router_logits: Tensor, shape (batch_size * sequence_length, num_experts)
+    """
     def forward(self, hidden_states: torch.Tensor):
+        """
+        输入:
+        hidden_states: Tensor, shape (batch_size, sequence_length, hidden_dim)
+        输出:
+            final_output: Tensor, shape (batch_size, sequence_length, hidden_dim)
+            router_logits: Tensor, shape (batch_size * sequence_length, num_experts)
+        实现思路：
+          1. 将输入展平为 (B*seq_len, hidden_dim)，通过 self.gate 得到 router_logits，
+             并计算全专家的 routing 权重（softmax 后）。
+          2. 对 routing 权重取 top-k，得到 routing_weights_topk 与 selected_experts；
+             如配置要求，归一化 top-k 概率。
+          3. 稀疏计算部分：仅计算每个 token 对于 top-k 专家的输出，
+             并累加得到 sparse_output（保留原版计算流程，同时记录激活专家的实际输出）。
+          4. Dense 估计部分：先计算所有专家对所有 token 的输出（all_expert_outputs），
+             再逐 token 调用 estimate_dense_output 得到 dense 输出（dense_estimated）。
+          5. 使用直通梯度技巧：前向输出用 sparse_output，但梯度来源于 dense_estimated。
+          6. 最后 reshape 为 (batch_size, sequence_length, hidden_dim) 并返回 final_output 及 router_logits.
+        """
+        #determine the shape of hidden_states
         batch_size, seq_length, hidden_dim = hidden_states.shape
         flat_hidden = hidden_states.view(-1, hidden_dim)  # (B*seq_len, hidden_dim)
 
-        # 计算路由 logits 和路由权重
+        # 计算路由 logits 和全专家 routing 权重
         router_logits = self.gate(flat_hidden)  # (B*seq_len, num_experts)
         routing_weights = F.softmax(router_logits, dim=1, dtype=torch.float)  # (B*seq_len, num_experts)
 
@@ -25,130 +56,90 @@ class DenseBackwardOlmoeSparseMoeBlock(OlmoeSparseMoeBlock):
         routing_weights_topk = routing_weights_topk.to(flat_hidden.dtype)
 
         # ---------- 稀疏计算部分 ----------
-        # 初始化稀疏输出
+        # 初始化稀疏输出，shape: (B*seq_len, hidden_dim)
         sparse_output = torch.zeros((flat_hidden.size(0), hidden_dim), dtype=flat_hidden.dtype, device=flat_hidden.device)
-        
-        # 记录每个专家被激活的token
-        token_indices_per_expert = [[] for _ in range(self.num_experts)]
-        expert_outputs_dict = {i: {} for i in range(self.num_experts)}
-        
-        # one-hot 编码 top-k 专家
+        # 用于记录每个 token 对激活专家的实际输出
+        activated_outputs = [{} for _ in range(flat_hidden.size(0))]
+        # one-hot 编码 top-k 专家，shape: (B*seq_len, top_k, num_experts)
         expert_mask = F.one_hot(selected_experts, num_classes=self.num_experts)  # (B*seq_len, top_k, num_experts)
         expert_mask = expert_mask.permute(2, 1, 0)  # (num_experts, top_k, B*seq_len)
 
-        # 保存每个token激活的专家列表
-        token_to_experts = [[] for _ in range(flat_hidden.size(0))]
-        
-        # 稀疏计算
         for expert_idx in range(self.num_experts):
             expert_layer = self.experts[expert_idx]
             idx, top_x = torch.where(expert_mask[expert_idx])
-            
             if top_x.numel() > 0:
-                # 记录该专家被哪些token激活
-                top_x_list = top_x.tolist()
-                token_indices_per_expert[expert_idx].extend(top_x_list)
-                
-                # 记录每个token激活了哪些专家
-                for token_idx in top_x_list:
-                    token_to_experts[token_idx].append(expert_idx)
-                
-                # 标准MoE前向计算
                 current_state = flat_hidden[top_x]  # (n, hidden_dim)
                 current_output = expert_layer(current_state)  # (n, hidden_dim)
                 weight = routing_weights_topk[top_x, idx].unsqueeze(-1)  # (n, 1)
                 weighted_output = current_output * weight
                 sparse_output.index_add_(0, top_x, weighted_output.to(flat_hidden.dtype))
-                
-                # 保存未加权的专家输出，用于后续估计
-                for i, token_idx in enumerate(top_x_list):
-                    expert_outputs_dict[expert_idx][token_idx] = current_output[i]
-        
-        # ---------- 初始化密集估计输出 ----------
-        # 创建一个每个token对所有专家输出的张量
-        dense_outputs = torch.zeros((flat_hidden.size(0), self.num_experts, hidden_dim), 
-                                   dtype=flat_hidden.dtype, device=flat_hidden.device)
-        
-        # 首先将已计算的专家输出填入
-        for expert_idx in range(self.num_experts):
-            for token_idx, output in expert_outputs_dict[expert_idx].items():
-                dense_outputs[token_idx, expert_idx] = output
-        
-        # ---------- 添加伪梯度路径 ----------
-        # 对于每个未激活的专家，创建一个小的梯度路径
-        # 首先创建一个全零的mask标记哪些(token,expert)对已被计算
-        computed_mask = torch.zeros((flat_hidden.size(0), self.num_experts), 
-                                   dtype=torch.bool, device=flat_hidden.device)
-        
-        # 标记已计算的(token,expert)对
-        for expert_idx in range(self.num_experts):
-            for token_idx in token_indices_per_expert[expert_idx]:
-                computed_mask[token_idx, expert_idx] = True
-        
-        # 对未计算的(token,expert)对，添加微小的伪梯度路径
-        # 我们使用一个小的缩放参数来确保前向计算几乎不受影响
-        scale = 1e-4
-        
-        # 计算一个残差连接，确保梯度能够流向所有专家
-        for token_idx in range(flat_hidden.size(0)):
-            token_input = flat_hidden[token_idx:token_idx+1]  # 保持维度
-            
-            # 对每个未激活的专家创建伪梯度路径
-            for expert_idx in range(self.num_experts):
-                if not computed_mask[token_idx, expert_idx]:
-                    # 查找是否有token激活了这个专家
-                    if token_indices_per_expert[expert_idx]:
-                        # 从激活该专家的token中选择一个
-                        # 优先选择与当前token共享其他专家的token
-                        similar_tokens = []
-                        for other_idx in token_indices_per_expert[expert_idx]:
-                            # 检查是否有共同激活的专家
-                            common_experts = set(token_to_experts[token_idx]) & set(token_to_experts[other_idx])
-                            if common_experts:
-                                similar_tokens.append(other_idx)
-                        
-                        if similar_tokens:
-                            # 使用相似token的平均输出
-                            similar_outputs = [expert_outputs_dict[expert_idx][t] for t in similar_tokens]
-                            estimated_output = torch.stack(similar_outputs).mean(0)
-                        else:
-                            # 使用所有激活该专家的token的平均输出
-                            all_outputs = [expert_outputs_dict[expert_idx][t] for t in token_indices_per_expert[expert_idx]]
-                            estimated_output = torch.stack(all_outputs).mean(0)
-                            
-                        # 添加微小的直接计算以维持梯度流
-                        direct_output = self.experts[expert_idx](token_input).squeeze(0)
-                        
-                        # 组合估计输出和直接计算
-                        # 前向使用估计输出，反向使用直接计算
-                        combined = estimated_output.detach() + scale * (direct_output - direct_output.detach())
-                        dense_outputs[token_idx, expert_idx] = combined
-                    else:
-                        # 如果没有token激活该专家，直接进行计算但使用小缩放
-                        # 这保证了梯度流而对前向几乎无影响
-                        direct_output = scale * self.experts[expert_idx](token_input).squeeze(0)
-                        dense_outputs[token_idx, expert_idx] = direct_output
-        
-        # ---------- 组合输出 ----------
-        # 使用路由权重对每个token的所有专家输出进行加权
-        dense_combined = torch.zeros_like(sparse_output)
-        
-        for token_idx in range(flat_hidden.size(0)):
-            # 对该token的所有专家输出进行加权
-            token_experts_output = dense_outputs[token_idx]  # (num_experts, hidden_dim)
-            token_routing_weights = routing_weights[token_idx].unsqueeze(-1)  # (num_experts, 1)
-            weighted_experts = token_experts_output * token_routing_weights
-            token_output = weighted_experts.sum(0)  # (hidden_dim)
-            dense_combined[token_idx] = token_output
-        
-        # 使用直通梯度技巧：前向使用sparse_output，反向使用dense_combined
-        # 增大直通梯度系数以增强梯度流
-        straight_through_scale = 1.0  # 可以尝试不同的值
-        final_flat = sparse_output.detach() + straight_through_scale * (dense_combined - dense_combined.detach())
-        
+                # 保存当前 token 对该专家的实际输出
+                for pos, token_idx in enumerate(top_x.tolist()):
+                    activated_outputs[token_idx][expert_idx] = current_output[pos]
+        # ---------- 稀疏计算结束 ----------
+
+        # ---------- Dense估计部分 ----------
+        # 计算所有专家对所有 token 的 dense 输出，shape: (B*seq_len, num_experts, hidden_dim)
+        all_expert_outputs = torch.stack([expert(flat_hidden) for expert in self.experts], dim=1)
+        # 将 selected_experts 转换为 list，每个 token 的激活专家列表
+        all_routing = selected_experts.tolist()  # 长度为 (B*seq_len)
+        
+        dense_outputs = []
+        for i in range(flat_hidden.size(0)):
+            dense_est = self.estimate_dense_output(
+                token_idx=i,
+                activated=all_routing[i],              # 当前 token 激活的专家列表，例如 [a, b]
+                gate_prob=routing_weights[i],            # 当前 token 的完整 routing 权重 (num_experts,)
+                activated_outputs=activated_outputs[i],  # 当前 token 对激活专家的实际输出
+                all_routing=all_routing,                 # 全 batch 每个 token 的激活专家列表（list of lists）
+                all_expert_outputs=all_expert_outputs      # (B*seq_len, num_experts, hidden_dim)
+            )
+            dense_outputs.append(dense_est.unsqueeze(0))
+        dense_outputs = torch.cat(dense_outputs, dim=0)  # (B*seq_len, hidden_dim)
+        # ---------- Dense估计结束 ----------
+
+        # 使用直通梯度：前向输出用稀疏结果，但反向传播时梯度来源于 dense 估计
+        final_flat = sparse_output.detach() + (dense_outputs - dense_outputs.detach())
         final_output = final_flat.view(batch_size, seq_length, hidden_dim)
         return final_output, router_logits
 
+    def estimate_dense_output(self, token_idx, activated, gate_prob, activated_outputs, all_routing, all_expert_outputs):
+        """
+        对于当前 token，根据 mini-batch 中的信息估计 dense 输出。
+        参数：
+            token_idx: 当前 token 的索引（标量）
+            activated: 当前 token 激活的专家列表，例如 [1, 3]
+            gate_prob: 当前 token 的 routing 权重，形状 (num_experts,)
+            activated_outputs: dict，当前 token 对激活专家的实际输出，形状 (hidden_dim,)
+            all_routing: list，每个 token 的激活专家列表（长度为 N，每个元素为 list）
+            all_expert_outputs: Tensor, (N, num_experts, hidden_dim)
+        返回：
+            estimated_dense: Tensor, (hidden_dim,)
+        """
+        num_experts = gate_prob.size(0)
+        dense_parts = {}
+        # 对于激活的专家，直接使用其实际输出
+        for idx in activated:
+            dense_parts[idx] = activated_outputs[idx]
+        # 对于未激活的专家，使用 mini-batch 中其他 token 的输出估计
+        non_activated = [i for i in range(num_experts) if i not in activated]
+        for i in non_activated:
+            indices = []
+            for idx, r_dec in enumerate(all_routing):
+                if (i in r_dec) and (len(set(r_dec) & set(activated)) > 0):
+                    indices.append(idx)
+            if indices:
+                selected_outputs = all_expert_outputs[indices, i, :]  # (n, hidden_dim)
+                estimated = selected_outputs.mean(dim=0)
+            else:
+                estimated = all_expert_outputs[:, i, :].mean(dim=0)
+            dense_parts[i] = estimated
+        # 按 gate_prob 加权求和各专家输出
+        estimated_dense = 0
+        for i in range(num_experts):
+            estimated_dense += gate_prob[i] * dense_parts[i]
+        return estimated_dense
+
 
 class DenseBackwardOLMoEForCausalLM(OlmoeForCausalLM):
     """
@@ -161,48 +152,21 @@ class DenseBackwardOLMoEForCausalLM(OlmoeForCausalLM):
     base_model_prefix = "olmoe"
 
     def __init__(self, config):
-        # 首先调用父类初始化方法
         super().__init__(config)
-        
-        # 不要尝试重新赋值self，而是从预训练模型加载并更新当前模型
-        pretrained_model = OlmoeForCausalLM.from_pretrained("allenai/OLMoE-1B-7B-0125")
-        
-        # 复制预训练模型的状态到当前模型
-        self.config = pretrained_model.config
-        self.model = pretrained_model.model
-        self.vocab_size = pretrained_model.vocab_size
-        self.router_aux_loss_coef = pretrained_model.router_aux_loss_coef
-        self.num_experts = pretrained_model.num_experts
-        self.lm_head = pretrained_model.lm_head
-        
         # 遍历���型中所有 decoder 层，替换每个 OlmoeSparseMoeBlock 为 DenseBackward 版本
         # 此处假设官方模型在 self.model.layers 中组织 decoder 层，
         # 且每层中 mlp 模块包含属性 sparse_moe_block。
         for layer in self.model.layers:
-            if hasattr(layer.mlp, "gate"):
-                print("111")
-                orig_block = layer.mlp
+            if hasattr(layer.mlp, "sparse_moe_block"):
+                orig_block = layer.mlp.sparse_moe_block
                 # 通过直接复制原版属性创建新的块
                 new_block = DenseBackwardOlmoeSparseMoeBlock(config)  # 或其他适当参数
                 # 然后手动复制需要共享的属性：
                 new_block.gate = orig_block.gate
                 new_block.experts = orig_block.experts
+                new_block.router = orig_block.router
                 new_block.num_experts = orig_block.num_experts
                 new_block.top_k = orig_block.top_k
                 new_block.norm_topk_prob = orig_block.norm_topk_prob
-                layer.mlp = new_block
-                print(type(layer.mlp))
+                layer.mlp.sparse_moe_block = new_block
 
-def main():
-    config = DenseBackwardOLMoEConfig(        # 官方模型参数
-    model_marker="DenseBackward_olmoe_marker",
-)
-# 创建自定义模型实例
-    model = DenseBackwardOLMoEForCausalLM(config)
-    print(type(model))
-    print(type(model.model))
-    print(type(model.model.layers[0]))
-    print(type(model.model.layers[0].mlp))
-    print(type(model.model.layers[0].mlp.experts))
-if __name__ == "__main__":
-    main()