Upload folder using huggingface_hub

README.md

@@ -0,0 +1,13 @@
+# DenseBackwardOLMoE
+
+自定义的OLMoE模型，使用DenseBackwardOlmoeSparseMoeBlock替换原版的MoE模块，实现dense backward功能。
+
+## 用法
+
+```python
+from transformers import AutoConfig, AutoModelForCausalLM
+
+# 使用trust_remote_code=True加载模型
+config = AutoConfig.from_pretrained("autoprogrammer/olmoe_densebackward", trust_remote_code=True)
+model = AutoModelForCausalLM.from_pretrained("autoprogrammer/olmoe_densebackward", config=config, trust_remote_code=True)
+```

__init__.py

@@ -0,0 +1,13 @@
+# 导出自定义配置和模型类
+from .configuration_densebackward_olmoe0125 import DenseBackwardOLMoEConfig
+from .modeling_densebackward_olmoe0125 import DenseBackwardOLMoEForCausalLM, DenseBackwardOlmoeSparseMoeBlock
+
+# 显式注册模型类型
+from transformers.models.auto.configuration_auto import CONFIG_MAPPING
+from transformers.models.auto.modeling_auto import MODEL_FOR_CAUSAL_LM_MAPPING
+
+__all__ = [
+    "DenseBackwardOLMoEConfig",
+    "DenseBackwardOLMoEForCausalLM", 
+    "DenseBackwardOlmoeSparseMoeBlock"
+]

config.json

@@ -0,0 +1,38 @@
+{
+  "_name_or_path": "allenai/OLMoE-1B-7B-0125",
+  "architectures": [
+    "DenseBackwardOLMoEForCausalLM"
+  ],
+  "auto_map": {
+  "AutoConfig": "configuration_densebackward_olmoe0125.DenseBackwardOLMoEConfig",
+  "AutoModel": "modeling_densebackward_olmoe0125.DenseBackwardOLMoEForCausalLM",
+  "AutoModelForCausalLM": "modeling_densebackward_olmoe0125.DenseBackwardOLMoEForCausalLM"
+},
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "clip_qkv": null,
+  "eos_token_id": 50279,
+  "hidden_act": "silu",
+  "hidden_size": 2048,
+  "initializer_range": 0.02,
+  "intermediate_size": 1024,
+  "max_position_embeddings": 4096,
+  "model_type": "olmoe",
+  "norm_topk_prob": false,
+  "num_attention_heads": 16,
+  "num_experts": 64,
+  "num_experts_per_tok": 8,
+  "num_hidden_layers": 16,
+  "num_key_value_heads": 16,
+  "output_router_logits": false,
+  "pad_token_id": 1,
+  "rms_norm_eps": 1e-05,
+  "rope_scaling": null,
+  "rope_theta": 10000.0,
+  "router_aux_loss_coef": 0.01,
+  "tie_word_embeddings": false,
+  "torch_dtype": "float32",
+  "transformers_version": "4.45.2",
+  "use_cache": true,
+  "vocab_size": 50304
+}

configuration_densebackward_olmoe0125.py

@@ -0,0 +1,25 @@
+# my_custom_olmoe/configuration_custom.py
+
+# 注意：根据你的 transformers 版本，导入官方 OLMoE 配置的路径可能需要调整
+from transformers.models.olmoe.configuration_olmoe import OlmoeConfig
+
+class DenseBackwardOLMoEConfig(OlmoeConfig):
+    model_type = "DenseBackward_olmoe"  # 这里覆盖 model_type 字段，便于后续识别
+    
+    # 添加auto_map用于支持AutoClass
+    auto_map = {
+        "AutoConfig": "configuration_custom.DenseBackwardOLMoEConfig",
+        "AutoModelForCausalLM": "modeling_custom.DenseBackwardOLMoEForCausalLM"
+    }
+
+    def __init__(self, model_marker="DenseBackward_olmoe_marker", **kwargs):
+        super().__init__(**kwargs)
+        self.model_marker = model_marker
+        self.intermediate_size= 1024
+#test
+def main():
+    config = DenseBackwardOLMoEConfig(model_marker="DenseBackward_olmoe_marker")
+    print(config)
+
+if __name__ == "__main__":
+    main()

generation_config.json

@@ -0,0 +1,6 @@
+{
+  "_from_model_config": true,
+  "eos_token_id": 50279,
+  "pad_token_id": 1,
+  "transformers_version": "4.45.2"
+}

model-00001-of-00006.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:df5a700fa91fd94e9d1a7ae523c5ae055f5879778a02ea19758edd37089da1ca
+size 4993992240

model-00002-of-00006.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5f299f21e6de71f5334e937fd51a083bab8cffe55682488219e4082d9558fdba
+size 4992966080

model-00003-of-00006.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4491eb552fd917a6777260f75c88e714f3174fbba7dd69d44b97c4124ddcd56b
+size 4992966080

model-00004-of-00006.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7aeea413752cd17a8e7c343aa436594a8158d279817e4993aa5df4b7c7cfdf22
+size 4992966416

model-00005-of-00006.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:18c7f749af753359e10c33a9c5deec688b7155f7c01052937806cd351c52cc38
+size 4992966680

model-00006-of-00006.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:299009ba118f5e918de33573a316ad3d3b9236651986fb1a14dc894f55269e67
+size 2711184968

modeling_densebackward_olmoe0125.py

@@ -0,0 +1,157 @@
+# 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):
+        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的激活mask，只有激活token梯度被保留
+            activation_mask = (selected_experts == expert_idx).any(dim=1).float().unsqueeze(-1).to(dtype)
+            # 注册梯度hook，使得非激活token的梯度被置零
+            if expert_output.requires_grad:
+                expert_output.register_hook(lambda grad, mask=activation_mask: grad * mask)
+            # 确保专家输出与预期类型一致
+            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)
+        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()

special_tokens_map.json

@@ -0,0 +1,16 @@
+{
+  "eos_token": {
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "pad_token": {
+    "content": "<|padding|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  }
+}

tokenizer_config.json

@@ -0,0 +1,239 @@
+{
+  "add_bos_token": false,
+  "add_eos_token": false,
+  "add_prefix_space": false,
+  "added_tokens_decoder": {
+    "0": {
+      "content": "|||IP_ADDRESS|||",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "1": {
+      "content": "<|padding|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "50254": {
+      "content": "                        ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50255": {
+      "content": "                       ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50256": {
+      "content": "                      ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50257": {
+      "content": "                     ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50258": {
+      "content": "                    ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50259": {
+      "content": "                   ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50260": {
+      "content": "                  ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50261": {
+      "content": "                 ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50262": {
+      "content": "                ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50263": {
+      "content": "               ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50264": {
+      "content": "              ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50265": {
+      "content": "             ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50266": {
+      "content": "            ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50267": {
+      "content": "           ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50268": {
+      "content": "          ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50269": {
+      "content": "         ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50270": {
+      "content": "        ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50271": {
+      "content": "       ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50272": {
+      "content": "      ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50273": {
+      "content": "     ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50274": {
+      "content": "    ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50275": {
+      "content": "   ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50276": {
+      "content": "  ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50277": {
+      "content": "|||EMAIL_ADDRESS|||",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50278": {
+      "content": "|||PHONE_NUMBER|||",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50279": {
+      "content": "<|endoftext|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "bos_token": null,
+  "clean_up_tokenization_spaces": true,
+  "eos_token": "<|endoftext|>",
+  "extra_special_tokens": {},
+  "model_max_length": 1000000000000000019884624838656,
+  "pad_token": "<|padding|>",
+  "tokenizer_class": "GPTNeoXTokenizer",
+  "unk_token": null
+}