Upload SpikingLM checkpoint and code

README.md

@@ -16,8 +16,9 @@ SpikingLM is a BERT-base style masked-language model with spiking attention bloc
 This checkpoint uses:
 
 - temporal Spiking BERT with `T=4`
+- learnable Q/K/V scaling parameters initialized from `7`
 - LIF nodes for projection, Q, K, V, attention output, and MLP blocks
-- `self.learnmax(attention_scores)` for attention normalization
+- `FP16OptimizedExp2Softmax` through `self.learnmax(attention_scores)` for attention normalization
 
 ## Files
 

scripts/finetune_glue.py

@@ -1,5 +1,4 @@
 #!/usr/bin/env python
-# 添加从预训练checkpoint加载的功能
 
 import argparse
 import json
@@ -22,7 +21,7 @@ from datasets import load_dataset
 from huggingface_hub import HfApi
 from torch.utils.data import DataLoader
 from tqdm.auto import tqdm
-from safetensors.torch import load_file  # 添加这个导入
+from safetensors.torch import load_file
 
 import transformers
 from transformers import (
@@ -212,41 +211,32 @@ def parse_args():
 
 def load_model_from_checkpoint(checkpoint_path, config, num_labels):
     """
-    从包含model.safetensors的checkpoint目录加载模型
-    
+    Load a sequence classification model from a checkpoint directory containing model.safetensors.
+
     Args:
-        checkpoint_path: checkpoint目录路径（如 step_10000）
-        config: 模型配置
-        num_labels: 分类任务的标签数量
-    
+        checkpoint_path: Checkpoint directory.
+        config: Model configuration.
+        num_labels: Number of labels for the classification task.
+
     Returns:
-        加载了预训练权重的模型
+        Model initialized with matching pretrained weights.
     """
     from spiking_bert.modeling_spiking_bert import BertForSequenceClassification
-    
-    # 检查safetensors文件是否存在
+
     safetensors_path = os.path.join(checkpoint_path, "model.safetensors")
     if not os.path.exists(safetensors_path):
         raise FileNotFoundError(f"Cannot find model.safetensors in {checkpoint_path}")
-    
+
     logger.info(f"Loading pretrained weights from {safetensors_path}")
-    
-    # 修改config以适应分类任务
+
     config.num_labels = num_labels
-    
-    # 创建分类模型（会初始化新的分类头）
     model = BertForSequenceClassification(config)
-    
-    # 加载预训练的权重
+
     pretrained_state_dict = load_file(safetensors_path)
-    
-    # 获取当前模型的state_dict
     model_state_dict = model.state_dict()
-    
-    # 只加载匹配的权重（跳过分类头）
     matched_keys = []
     mismatched_keys = []
-    
+
     for key in pretrained_state_dict.keys():
         if key in model_state_dict:
             if pretrained_state_dict[key].shape == model_state_dict[key].shape:
@@ -257,10 +247,9 @@ def load_model_from_checkpoint(checkpoint_path, config, num_labels):
                 logger.warning(f"Shape mismatch for {key}: pretrained {pretrained_state_dict[key].shape} vs model {model_state_dict[key].shape}")
         else:
             logger.info(f"Key {key} not found in model, skipping...")
-    
-    # 加载权重
+
     model.load_state_dict(model_state_dict)
-    
+
     logger.info(f"Loaded {len(matched_keys)} matching weights from checkpoint")
     logger.info(f"Skipped {len(mismatched_keys)} mismatched weights")
     logger.info(f"Classifier head will be trained from scratch")

spiking_bert/modeling_spiking_bert.py

@@ -197,7 +197,6 @@ class FP16OptimizedExp2Softmax(nn.Module):
             init_value: Initial value for scale parameter
         """
         super().__init__()
-        # 注册为可学习参数，维度 [T, head_size]
         self.scale = nn.Parameter(torch.full((T,), init_value))
     
     def forward(self, tensor, k=1.0):
@@ -211,9 +210,6 @@ class FP16OptimizedExp2Softmax(nn.Module):
         tensor = tensor - tensor.max(dim=-1, keepdim=True)[0] - 1
         tensor.clamp_(min=-30.0, max=30.0)
         tensor = torch.exp2(tensor)
-        # tensor = tensor / (tensor.sum(dim=-1, keepdim=True) + 1e-6)
-        
-        # scale: [T, head_size] -> [T, 1, head_size, 1, 1] for broadcasting
         scale = self.scale[:, None, None, None, None]
         tensor = tensor * scale
         return tensor
@@ -240,13 +236,10 @@ class BertSelfAttention(nn.Module):
 
         self.query = nn.Linear(config.hidden_size, self.all_head_size)
         self.q_lam = nn.Parameter(torch.ones(self.all_head_size)*7)
-        # self.q_norm = nn.RMSNorm(config.hidden_size)
         self.key = nn.Linear(config.hidden_size, self.all_head_size)
         self.k_lam = nn.Parameter(torch.ones(self.all_head_size)*7)
-        # self.k_norm = nn.RMSNorm(config.hidden_size)
         self.value = nn.Linear(config.hidden_size, self.all_head_size)
         self.v_lam = nn.Parameter(torch.ones(self.all_head_size)*7)
-        # self.v_norm = nn.RMSNorm(config.hidden_size)
         self.learnmax = FP16OptimizedExp2Softmax(T=config.T)
         self.dropout = nn.Dropout(config.attention_probs_dropout_prob)
         self.position_embedding_type = position_embedding_type or getattr(
@@ -258,7 +251,6 @@ class BertSelfAttention(nn.Module):
 
         self.is_decoder = config.is_decoder
         self.layer_idx = layer_idx
-        # True 时不做 q_lam/k_lam/v_lam 缩放（等价于恒等，便于与启用 lamb 的实验对比）
         self.disable_qkv_lamb = bool(getattr(config, "disable_qkv_lamb", False))
 
     @deprecate_kwarg("past_key_value", new_name="past_key_values", version="4.58")
@@ -275,7 +267,6 @@ class BertSelfAttention(nn.Module):
         T, batch_size, seq_length, _ = hidden_states.shape
         hidden_states = self.proj_lif(hidden_states)
         query_layer = self.query(hidden_states) 
-        # query_layer = self.q_norm(query_layer.reshape(T*batch_size, seq_length, -1)).reshape(T, batch_size, seq_length, -1)
         if not self.disable_qkv_lamb:
             query_layer = self.q_lam * query_layer
         query_layer = self.q_lif(query_layer)
@@ -303,7 +294,6 @@ class BertSelfAttention(nn.Module):
             value_layer = curr_past_key_value.layers[self.layer_idx].values
         else:
             key_layer = self.key(current_states) 
-            # key_layer = self.k_norm(key_layer.reshape(T*batch_size, seq_length, -1)).reshape(T, batch_size, seq_length, -1)
             if not self.disable_qkv_lamb:
                 key_layer = self.k_lam * key_layer
             key_layer = self.k_lif(key_layer)
@@ -311,7 +301,6 @@ class BertSelfAttention(nn.Module):
                 2, 3
             )
             value_layer = self.value(current_states) 
-            # value_layer = self.v_norm(value_layer.reshape(T*batch_size, seq_length, -1)).reshape(T, batch_size, seq_length, -1)
             if not self.disable_qkv_lamb:
                 value_layer = self.v_lam * value_layer
             value_layer = self.v_lif(value_layer)
@@ -354,27 +343,10 @@ class BertSelfAttention(nn.Module):
                 relative_position_scores_key = torch.einsum("bhrd,lrd->bhlr", key_layer, positional_embedding)
                 attention_scores = attention_scores + relative_position_scores_query.unsqueeze(0) + relative_position_scores_key.unsqueeze(0)
 
-        # attention_scores = attention_scores / math.sqrt(self.attention_head_size)
-        # print(attention_mask)
         if attention_mask is not None:
-            # Apply the attention mask is (precomputed for all layers in BertModel forward() function)
             attention_scores = attention_scores + attention_mask
 
-        # # # Normalize the attention scores to probabilities.
         attention_probs = self.learnmax(attention_scores)
-        # attention_probs = nn.functional.softmax(attention_scores, dim=-1)
-
-        # # This is actually dropping out entire tokens to attend to, which might
-        # # seem a bit unusual, but is taken from the original Transformer paper.
-        # attention_probs = self.dropout(attention_probs)
-        # if attention_mask is not None:
-        #     attention_scores = attention_scores * attention_mask.unsqueeze(0)
-        
-        # attention_probs = attention_scores * self.scale
-
-        # Mask heads if we want to
-        # if head_mask is not None:
-        #     attention_probs = attention_probs * head_mask
 
         context_layer = torch.matmul(attention_probs, value_layer)  # t,b,h,l,d -> t,b,l,h,d
 
@@ -580,7 +552,6 @@ class BertIntermediate(nn.Module):
     def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
         hidden_states = self.mlp1_lif(hidden_states)
         hidden_states = self.dense(hidden_states)
-        # hidden_states = self.intermediate_act_fn(hidden_states)
         return hidden_states
 
 
@@ -1075,8 +1046,6 @@ class BertModel(BertPreTrainedModel):
         # input head_mask has shape [num_heads] or [num_hidden_layers x num_heads]
         # and head_mask is converted to shape [num_hidden_layers x batch x num_heads x seq_length x seq_length]
 
-        ### update extended_attention_mask during training here!!!
-        # extended_attention_mask = attention_mask.unsqueeze(1).unsqueeze(1)
         head_mask = self.get_head_mask(head_mask, self.config.num_hidden_layers)
         encoder_outputs = self.encoder(
             embedding_output,
@@ -1879,4 +1848,4 @@ __all__ = [
     "BertModel",
     "BertPreTrainedModel",
     "load_tf_weights_in_bert",
-]
+]