# 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_v1 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): batch_size, seq_length, hidden_dim = hidden_states.shape # 记录输入张量的数据类型，确保所有计算保持一致 dtype = hidden_states.dtype device = hidden_states.device flat_hidden = hidden_states.view(-1, hidden_dim) # (B*seq_len, hidden_dim) N_tokens = flat_hidden.size(0) # 计算路由逻辑 router_logits = self.gate(flat_hidden) # (B*seq_len, num_experts) # 确保router_logits和flat_hidden数据类型一致 router_logits = router_logits.to(dtype=dtype) routing_weights = F.softmax(router_logits, dim=1, dtype=dtype) # (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(dtype=dtype) # ---------- 真实计算所有专家输出（密集计算）---------- all_expert_outputs = torch.zeros((N_tokens, self.num_experts, hidden_dim), dtype=dtype, device=device) for expert_idx in range(self.num_experts): expert_layer = self.experts[expert_idx] # 对所有token都计算当前专家的输出 expert_output = expert_layer(flat_hidden) # (N_tokens, hidden_dim) # 确保专家输出与预期类型一致 expert_output = expert_output.to(dtype=dtype) all_expert_outputs[:, expert_idx, :] = expert_output # ---------- 提取激活专家输出（稀疏前向）- 使用张量批处理 ---------- # 创建索引张量，第一维是token索引，第二维是专家索引 token_indices = torch.arange(N_tokens, device=device).unsqueeze(1).expand(-1, self.top_k) batch_indices = token_indices.reshape(-1) expert_indices = selected_experts.reshape(-1) # 批量提取激活专家的输出 selected_outputs = all_expert_outputs[batch_indices, expert_indices].view(N_tokens, self.top_k, hidden_dim) # 扩展权重以便批量相乘 expanded_weights = routing_weights_topk.unsqueeze(-1) # (N_tokens, top_k, 1) expanded_weights = expanded_weights.to(dtype=dtype) # 权重乘以专家输出并求和 sparse_output = (selected_outputs * expanded_weights).sum(dim=1) # (N_tokens, hidden_dim) # ---------- 密集计算聚合（用于反向传播）---------- # 使用所有专家的输出和路由权重计算密集输出 routing_weights_expanded = routing_weights.unsqueeze(-1) # (N_tokens, num_experts, 1) routing_weights_expanded = routing_weights_expanded.to(dtype=dtype) print(expanded_weights.shape) print("sparse",expanded_weights) print(routing_weights_expanded.shape) print("dense",routing_weights_expanded) dense_outputs = (all_expert_outputs * routing_weights_expanded).sum(dim=1) # (N_tokens, hidden_dim) # ---------- 组合稀疏前向和密集反向 ---------- # sparse_output.detach()保留稀疏前向计算图 # (dense_outputs - dense_outputs.detach())只保留密集反向梯度 final_flat = sparse_output.detach() + (dense_outputs - dense_outputs.detach()) final_flat = final_flat.to(dtype=dtype) # 确保最终输出类型一致 final_output = final_flat.view(batch_size, seq_length, hidden_dim) return final_output, router_logits 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)) # 释放预训练模型内存 del pretrained_model import gc gc.collect() torch.cuda.empty_cache() print("原始预训练模型已释放") 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()