modeling_densebackward_olmoe0125.py

# my_custom_olmoe/modeling_custom.py

import torch
import torch.nn as nn
import torch.nn.functional as F

# 导入官方实现（注意根据你的 transformers 版本调整导入路径）
from transformers.models.olmoe.modeling_olmoe import OlmoeForCausalLM, OlmoeSparseMoeBlock, OlmoeMLP
from .configuration_densebackward_olmoe0125 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):
        # 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))
                
                # 直接为激活的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
                
        # ---------- 稀疏计算结束 ----------

        # ---------- 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
        
        # 预分配结果张量，注意是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 输出。
        参数：
            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)
                 # 只计算非零值的平均值
                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])
            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])
            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):
    """
    自定义的 Olmoe ForCausalLM 模型，使用新的 DenseBackwardOlmoeSparseMoeBlock 替换原版的 MoE 模块，
    以实现 dense backward 功能。
    
    配置类：DenseBackwardOLMoEConfig
    """
    config_class = DenseBackwardOLMoEConfig
    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
                # 通过直接复制原版属性创建新的块
                new_block = DenseBackwardOlmoeSparseMoeBlock(config)  # 或其他适当参数
                # 然后手动复制需要共享的属性：
                new_block.gate = orig_block.gate
                new_block.experts = orig_block.experts
                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()