回滚到逻辑正确版本

modeling_densebackward_olmoe0125.py

@@ -23,10 +23,27 @@ class DenseBackwardOlmoeSparseMoeBlock(OlmoeSparseMoeBlock):
         router_logits: Tensor, shape (batch_size * sequence_length, num_experts)
     """
     def forward(self, hidden_states: torch.Tensor):
-        # determine the shape of hidden_states
+        """
+        输入:
+        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)
-        total_tokens = flat_hidden.size(0)
 
         # 计算路由 logits 和全专家 routing 权重
         router_logits = self.gate(flat_hidden)  # (B*seq_len, num_experts)
@@ -40,18 +57,9 @@ class DenseBackwardOlmoeSparseMoeBlock(OlmoeSparseMoeBlock):
 
         # ---------- 稀疏计算部分 ----------
         # 初始化稀疏输出，shape: (B*seq_len, hidden_dim)
-        sparse_output = torch.zeros((total_tokens, hidden_dim), dtype=flat_hidden.dtype, device=flat_hidden.device)
-        
-        # 创建一个张量存储激活专家的输出，避免使用Python字典
-        # shape: (B*seq_len, num_experts, hidden_dim)
-        all_expert_outputs = torch.zeros((total_tokens, self.num_experts, hidden_dim), 
-                                       dtype=flat_hidden.dtype, device=flat_hidden.device)
-        
-        # 使用张量掩码跟踪哪些专家被激活
-        # shape: (B*seq_len, num_experts)
-        expert_activated = torch.zeros((total_tokens, self.num_experts), 
-                                      dtype=torch.bool, device=flat_hidden.device)
-        
+        sparse_output = torch.zeros((flat_hidden.size(0), hidden_dim), dtype=flat_hidden.dtype, device=flat_hidden.device)
+        # 用于记录每个 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)
@@ -65,30 +73,29 @@ class DenseBackwardOlmoeSparseMoeBlock(OlmoeSparseMoeBlock):
                 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))
-                
-                # 直接为激活的token分配专家输出
-                for i in range(top_x.shape[0]):
-                    token_idx = top_x[i]
-                    all_expert_outputs[token_idx, expert_idx] = current_output[i]
-                    expert_activated[token_idx, expert_idx] = True
-                
+                # 保存当前 token 对该专家的实际输出
+                for pos, token_idx in enumerate(top_x.tolist()):
+                    activated_outputs[token_idx][expert_idx] = current_output[pos]
         # ---------- 稀疏计算结束 ----------
 
         # ---------- Dense估计部分 ----------
-        # 从GPU获取必要信息，避免过多的tensor->list转换
-        selected_experts_gpu = selected_experts  # 保持在GPU上
-
-        # 预分配结果张量，避免在循环中append
-        dense_outputs = torch.zeros_like(sparse_output)
+        # 计算所有专家对所有 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)
         
-        # 使用向量化的estimate_dense_output函数
-        dense_outputs = self.estimate_dense_output_batch(
-            total_tokens=total_tokens,
-            selected_experts=selected_experts_gpu,
-            routing_weights=routing_weights,
-            expert_activated=expert_activated,
-            all_expert_outputs=all_expert_outputs
-        )
+        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 估计
@@ -96,90 +103,6 @@ class DenseBackwardOlmoeSparseMoeBlock(OlmoeSparseMoeBlock):
         final_output = final_flat.view(batch_size, seq_length, hidden_dim)
         return final_output, router_logits
 
-    def estimate_dense_output_batch(self, total_tokens, selected_experts, routing_weights, 
-                                  expert_activated, all_expert_outputs):
-        """
-        批量估计所有token的dense输出，优化版本。
-        
-        参数：
-            total_tokens: token总数
-            selected_experts: 每个token激活的专家索引，形状 (total_tokens, top_k)
-            routing_weights: 路由权重，形状 (total_tokens, num_experts)
-            expert_activated: 掩码张量，标记每个token激活了哪些专家，形状 (total_tokens, num_experts)
-            all_expert_outputs: 专家输出，形状 (total_tokens, num_experts, hidden_dim)
-            
-        返回：
-            dense_outputs: 形状 (total_tokens, hidden_dim)
-        """
-        hidden_dim = all_expert_outputs.size(-1)
-        num_experts = routing_weights.size(1)
-        device = all_expert_outputs.device
-        
-        # 预分配结果张量，注意是hidden_dim而不是num_experts
-        dense_outputs = torch.zeros((total_tokens, hidden_dim), dtype=all_expert_outputs.dtype, device=device)
-        
-        # 对每个token单独处理（此处仍需循环，但后续可进一步优化）
-        for token_idx in range(total_tokens):
-            # 对于激活的专家，直接使用输出
-            activated_mask = expert_activated[token_idx]  # (num_experts,)
-            
-            # 对于未激活的专家，找到估计值
-            for expert_idx in range(num_experts):
-                if activated_mask[expert_idx]:
-                    # 直接使用激活专家的输出
-                    expert_output = all_expert_outputs[token_idx, expert_idx]
-                else:
-                    # 寻找可以用于估计的输出
-                    # 找出其他激活了当前专家的token
-                    tokens_with_expert = expert_activated[:, expert_idx]
-                    
-                    # 找出同时激活了当前token的某些专家和当前专家的其他token
-                    # 首先获取当前token激活的专家
-                    current_activated = selected_experts[token_idx]
-                    
-                    # 在其他token中寻找同时激活了current_activated中专家和expert_idx的token
-                    valid_tokens = torch.zeros(total_tokens, dtype=torch.bool, device=device)
-                    
-                    # 对于每个其他token，检查它是否同时激活了当前token的某个专家和当前专家
-                    for other_token in range(total_tokens):
-                        if other_token == token_idx:
-                            continue
-                        
-                        # 检查其他token是否激活了当前专家
-                        if expert_activated[other_token, expert_idx]:
-                            # 检查是否有共同激活的专家
-                            other_experts = selected_experts[other_token]
-                            common = torch.any(torch.isin(other_experts, current_activated))
-                            if common:
-                                valid_tokens[other_token] = True
-                    
-                    # 如果找到了有效token
-                    if valid_tokens.any():
-                        # 获取有效token对当前专家的输出
-                        valid_outputs = all_expert_outputs[valid_tokens, expert_idx]
-                        # 只计算非零值的平均值
-                        mask = (valid_outputs.sum(dim=-1) != 0).to(valid_outputs.dtype).unsqueeze(-1)
-                        if mask.sum() > 0:
-                            expert_output = (valid_outputs * mask).sum(dim=0) / mask.sum()
-                        else:
-                            expert_output = torch.zeros(hidden_dim, dtype=all_expert_outputs.dtype, device=device)
-                    else:
-                        # 如果没有找到有效token，使用所有激活了当前专家的token的输出
-                        if tokens_with_expert.any():
-                            all_valid_outputs = all_expert_outputs[tokens_with_expert, expert_idx]
-                            mask = (all_valid_outputs.sum(dim=-1) != 0).to(all_valid_outputs.dtype).unsqueeze(-1)
-                            if mask.sum() > 0:
-                                expert_output = (all_valid_outputs * mask).sum(dim=0) / mask.sum()
-                            else:
-                                expert_output = torch.zeros(hidden_dim, dtype=all_expert_outputs.dtype, device=device)
-                        else:
-                            expert_output = torch.zeros(hidden_dim, dtype=all_expert_outputs.dtype, device=device)
-                
-                # 根据routing权重加权
-                dense_outputs[token_idx] += routing_weights[token_idx, expert_idx] * expert_output
-        
-        return dense_outputs
-
     def estimate_dense_output(self, token_idx, activated, gate_prob, activated_outputs, all_routing, all_expert_outputs):
         """
         对于当前 token，根据 mini-batch 中的信息估计 dense 输出。
@@ -207,21 +130,9 @@ class DenseBackwardOlmoeSparseMoeBlock(OlmoeSparseMoeBlock):
                     indices.append(idx)
             if indices:
                 selected_outputs = all_expert_outputs[indices, i, :]  # (n, hidden_dim)
-                 # 只计算非零值的平均值
-                mask = (selected_outputs.sum(dim=-1) != 0).to(selected_outputs.dtype).unsqueeze(-1)
-                if mask.sum() > 0:
-                    estimated = (selected_outputs * mask).sum(dim=0) / mask.sum()
-                else:
-                   # 如果全是零，返回零向量
-                    estimated = torch.zeros_like(selected_outputs[0])
+                estimated = selected_outputs.mean(dim=0)
             else:
-                all_outputs = all_expert_outputs[:, i, :]
-                mask = (all_outputs.sum(dim=-1) != 0).to(all_outputs.dtype).unsqueeze(-1)
-                if mask.sum() > 0:
-                    estimated = (all_outputs * mask).sum(dim=0) / mask.sum()
-                else:
-                    # 如果全是零，返回零向量
-                    estimated = torch.zeros_like(all_outputs[0])
+                estimated = all_expert_outputs[:, i, :].mean(dim=0)
             dense_parts[i] = estimated
         # 按 gate_prob 加权求和各专家输出
         estimated_dense = 0
@@ -241,49 +152,20 @@ 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))
-
-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()
+                layer.mlp.sparse_moe_block = new_block