modeling_densebackward_olmoe.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_olmoe import DenseBackwardOLMoEConfig


class DenseBackwardOlmoeSparseMoeBlock(OlmoeSparseMoeBlock):
    def forward(self, hidden_states: torch.Tensor):
        batch_size, seq_length, hidden_dim = hidden_states.shape
        flat_hidden = hidden_states.view(-1, hidden_dim)  # (B*seq_len, hidden_dim)

        # 计算路由 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)

        # ---------- 稀疏计算部分 ----------
        # 初始化稀疏输出
        sparse_output = torch.zeros((flat_hidden.size(0), hidden_dim), dtype=flat_hidden.dtype, device=flat_hidden.device)
        
        # 存储所有激活信息的数据结构
        num_tokens = flat_hidden.size(0)
        all_activated_outputs = {}  # {expert_idx: {token_idx: output_tensor}}
        all_routing_indices = {}    # {expert_idx: [token_indices]}
        token_activated_experts = {}  # {token_idx: [activated_expert_indices]}
        
        # one-hot 编码 top-k 专家
        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和对应输出
                all_activated_outputs[expert_idx] = {}
                all_routing_indices[expert_idx] = top_x.tolist()
                
                for pos, token_idx in enumerate(top_x.tolist()):
                    # 记录该专家对该token的输出
                    all_activated_outputs[expert_idx][token_idx] = current_output[pos]
                    
                    # 记录该token激活的专家
                    if token_idx not in token_activated_experts:
                        token_activated_experts[token_idx] = []
                    token_activated_experts[token_idx].append(expert_idx)
        # ---------- 稀疏计算结束 ----------

        # ---------- Dense估计部分 ----------
        # 将activated_experts 转换为list格式，与路由权重匹配
        all_routing = selected_experts.tolist()  # 长度为 (B*seq_len)
        
        # 使用已激活信息估计dense输出
        dense_outputs = []
        for token_idx in range(num_tokens):
            # 获取当前token的激活专家列表
            activated = all_routing[token_idx] if token_idx in token_activated_experts else []
            
            # 估计dense输出（只使用已经计算过的专家输出）
            dense_est = self.estimate_dense_output_efficient(
                token_idx=token_idx,
                activated=activated,
                gate_prob=routing_weights[token_idx],
                all_activated_outputs=all_activated_outputs,
                all_routing_indices=all_routing_indices,
                token_activated_experts=token_activated_experts
            )
            dense_outputs.append(dense_est.unsqueeze(0))
            
        dense_outputs = torch.cat(dense_outputs, dim=0)  # (B*seq_len, hidden_dim)
        # ---------- 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_efficient(self, token_idx, activated, gate_prob, 
                                       all_activated_outputs, all_routing_indices, token_activated_experts):
        """
        优化版本的dense输出估计，只使用已计算的专家输出
        """
        num_experts = gate_prob.size(0)
        dense_parts = {}
        
        # 对于激活的专家，直接使用其输出
        for expert_idx in activated:
            if expert_idx in all_activated_outputs and token_idx in all_activated_outputs[expert_idx]:
                dense_parts[expert_idx] = all_activated_outputs[expert_idx][token_idx]
        
        # 对于未激活的专家，使用其他token的激活输出估计
        non_activated = [i for i in range(num_experts) if i not in activated]
        for expert_idx in non_activated:
            # 如果该专家没有被任何token激活，跳过
            if expert_idx not in all_routing_indices or not all_routing_indices[expert_idx]:
                # 使用零向量或平均值作为估计
                dense_parts[expert_idx] = torch.zeros_like(next(iter(dense_parts.values()))) if dense_parts else 0
                continue
                
            # 找出激活了该专家的token，并且这些token也激活了当前token激活的某些专家
            candidate_tokens = []
            for other_token in all_routing_indices[expert_idx]:
                # 检查other_token是否与当前token共享某些激活专家
                if other_token in token_activated_experts:
                    common_experts = set(activated) & set(token_activated_experts[other_token])
                    if common_experts:
                        candidate_tokens.append(other_token)
            
            # 如果找到了候选token，使用它们的输出平均值
            if candidate_tokens:
                expert_outputs = [all_activated_outputs[expert_idx][t] for t in candidate_tokens]
                estimated = torch.stack(expert_outputs).mean(dim=0)
            else:
                # 找不到合适的候选，使用所有激活了该专家的token
                expert_outputs = [all_activated_outputs[expert_idx][t] for t in all_routing_indices[expert_idx]]
                estimated = torch.stack(expert_outputs).mean(dim=0)
                
            dense_parts[expert_idx] = estimated
            
        # 按路由权重加权求和
        estimated_dense = 0
        for expert_idx in range(num_experts):
            if expert_idx in dense_parts:
                estimated_dense += gate_prob[expert_idx] * dense_parts[expert_idx]
                
        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-0924", torch_dtype=torch.bfloat16)
        
        # 复制预训练模型的状态到当前模型
        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))
        # 在调用post_init()前
        test_param = self.model.layers[0].mlp.experts[0].up_proj.weight.data[0, 0].item()
        print(f"权重示例值（前）: {test_param}")
        self.post_init()
        # 在调用post_init()后
        test_param_after = self.model.layers[0].mlp.experts[0].up_proj.weight.data[0, 0].item()
        print(f"权重示例值（后）: {test_param_after}")

def main():
    config = DenseBackwardOLMoEConfig(        # 官方模型参数
    model_marker="DenseBackward_olmoe_marker",
    torch_dtype="bfloat16"
)
# 创建自定义模型实例
    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()