Update modeling_densebackward_olmoe0125.py

modeling_densebackward_olmoe0125.py

@@ -23,173 +23,173 @@ 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
-    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)
-    routing_weights = F.softmax(router_logits, dim=1, dtype=torch.float)  # (B*seq_len, num_experts)
-
-    # Top-k 选择
-    routing_weights_topk, selected_experts = torch.topk(routing_weights, self.top_k, dim=-1)
-    if self.norm_topk_prob:
-        routing_weights_topk = routing_weights_topk / routing_weights_topk.sum(dim=-1, keepdim=True)
-    routing_weights_topk = routing_weights_topk.to(flat_hidden.dtype)
-
-    # ---------- 稀疏计算部分 ----------
-    # 初始化稀疏输出，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)
-    
-    # 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)
-
-    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:
-            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))
-            
-            # 保存专家输出到张量中，而不是使用字典
-            all_expert_outputs.index_copy_(0, top_x, 
-                                         torch.zeros_like(all_expert_outputs[0:top_x.size(0)]).scatter_(
-                                             1, expert_idx * torch.ones((top_x.size(0), 1), 
-                                                                     dtype=torch.long, 
-                                                                     device=flat_hidden.device), 
-                                             current_output.unsqueeze(1)))
-            
-            # 标记哪些专家被激活
-            expert_activated.index_copy_(0, top_x,
-                                       torch.zeros_like(expert_activated[0:top_x.size(0)]).scatter_(
-                                           1, expert_idx * torch.ones((top_x.size(0), 1), 
-                                                                   dtype=torch.long, 
-                                                                   device=flat_hidden.device), 
-                                           torch.ones((top_x.size(0), 1), 
-                                                    dtype=torch.bool, 
-                                                    device=flat_hidden.device)))
-    # ---------- 稀疏计算结束 ----------
-
-    # ---------- Dense估计部分 ----------
-    # 从GPU获取必要信息，避免过多的tensor->list转换
-    selected_experts_gpu = selected_experts  # 保持在GPU上
-
-    # 预分配结果张量，避免在循环中append
-    dense_outputs = torch.zeros_like(sparse_output)
-    
-    # 使用向量化的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估计结束 ----------
-
-    # 使用直通梯度：前向输出用稀疏结果，但反向传播时梯度来源于 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_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)
+        # 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)
+        routing_weights = F.softmax(router_logits, dim=1, dtype=torch.float)  # (B*seq_len, num_experts)
+
+        # Top-k 选择
+        routing_weights_topk, selected_experts = torch.topk(routing_weights, self.top_k, dim=-1)
+        if self.norm_topk_prob:
+            routing_weights_topk = routing_weights_topk / routing_weights_topk.sum(dim=-1, keepdim=True)
+        routing_weights_topk = routing_weights_topk.to(flat_hidden.dtype)
+
+        # ---------- 稀疏计算部分 ----------
+        # 初始化稀疏输出，shape: (B*seq_len, hidden_dim)
+        sparse_output = torch.zeros((total_tokens, hidden_dim), dtype=flat_hidden.dtype, device=flat_hidden.device)
         
-    返回：
-        dense_outputs: 形状 (total_tokens, hidden_dim)
-    """
-    hidden_dim = all_expert_outputs.size(-1)
-    num_experts = routing_weights.size(1)
-    device = all_expert_outputs.device
-    
-    # 预分配结果张量
-    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,)
+        # 创建一个张量存储激活专家的输出，避免使用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)
         
-        # 对于未激活的专家，找到估计值
-        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)
+        # 使用张量掩码跟踪哪些专家被激活
+        # shape: (B*seq_len, num_experts)
+        expert_activated = torch.zeros((total_tokens, self.num_experts), 
+                                      dtype=torch.bool, device=flat_hidden.device)
+        
+        # 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)
+
+        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:
+                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))
                 
-                # 对于每个其他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
+                # 保存专家输出到张量中，而不是使用字典
+                all_expert_outputs.index_copy_(0, top_x, 
+                                             torch.zeros_like(all_expert_outputs[0:top_x.size(0)]).scatter_(
+                                                 1, expert_idx * torch.ones((top_x.size(0), 1), 
+                                                                         dtype=torch.long, 
+                                                                         device=flat_hidden.device), 
+                                                 current_output.unsqueeze(1)))
                 
-                # 如果找到了有效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)
+                # 标记哪些专家被激活
+                expert_activated.index_copy_(0, top_x,
+                                           torch.zeros_like(expert_activated[0:top_x.size(0)]).scatter_(
+                                               1, expert_idx * torch.ones((top_x.size(0), 1), 
+                                                                       dtype=torch.long, 
+                                                                       device=flat_hidden.device), 
+                                               torch.ones((top_x.size(0), 1), 
+                                                        dtype=torch.bool, 
+                                                        device=flat_hidden.device)))
+        # ---------- 稀疏计算结束 ----------
+
+        # ---------- Dense估计部分 ----------
+        # 从GPU获取必要信息，避免过多的tensor->list转换
+        selected_experts_gpu = selected_experts  # 保持在GPU上
+
+        # 预分配结果张量，避免在循环中append
+        dense_outputs = torch.zeros_like(sparse_output)
+        
+        # 使用向量化的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估计结束 ----------
+
+        # 使用直通梯度：前向输出用稀疏结果，但反向传播时梯度来源于 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_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
+        
+        # 预分配结果张量
+        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，使用所有激活了当前专家的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)
+                    # 寻找可以用于估计的输出
+                    # 找出其他激活了当前专家的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 = (all_valid_outputs * mask).sum(dim=0) / mask.sum()
+                            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:
-                        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
+                        # 如果没有找到有效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
 
 
 class DenseBackwardOLMoEForCausalLM(OlmoeForCausalLM):
@@ -238,13 +238,14 @@ class DenseBackwardOLMoEForCausalLM(OlmoeForCausalLM):
 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()