diff --git a/model/LMConfig.py b/model/LMConfig.py index fb0a76d..fedfc1c 100644 --- a/model/LMConfig.py +++ b/model/LMConfig.py @@ -20,6 +20,10 @@ class LMConfig(PretrainedConfig): dropout: float = 0.0, flash_attn: bool = True, #################################################### + # DB related configurations + #################################################### + disable_db: bool = False, # 特殊模式：禁用数据库功能 + #################################################### # Here are the specific configurations of MOE # When use_moe is false, the following is invalid #################################################### @@ -47,6 +51,10 @@ class LMConfig(PretrainedConfig): self.dropout = dropout self.flash_attn = flash_attn #################################################### + # DB related configurations + #################################################### + self.disable_db = disable_db # 设置是否禁用数据库 + #################################################### # Here are the specific configurations of MOE # When use_moe is false, the following is invalid #################################################### diff --git a/model/model.py b/model/model.py index 878c675..31aca85 100644 --- a/model/model.py +++ b/model/model.py @@ -373,7 +373,7 @@ class MiniMindBlock(nn.Module): # self.product_key_topk = min(16, self.num_keys) # 确保不超过num_keys # self.num_experts_per_head_topk = 1 # 最终每个头选取的专家数 - def forward(self, x,db_value, pos_cis, past_key_value=None, use_cache=False): + def forward(self, x, db_value, pos_cis, past_key_value=None, use_cache=False): # import pdb;pdb.set_trace() # db_value = None @@ -426,7 +426,7 @@ class MiniMindBlock(nn.Module): db_value=db_value ) - h_attn = self.cross_att(h_attn,db_value) + h_attn = self.cross_att(h_attn, db_value) # 残差连接 h = x + h_attn @@ -568,10 +568,19 @@ class MiniMindLM(PreTrainedModel): h_list = [] for l, layer in enumerate(self.layers): - index = self.extract_db.q_to_k(h) - db_value = self.extract_db.get_data(index) + # 禁用数据库模式，使用固定值替代数据库查询 + if self.params.disable_db: + # 创建一个形状为[batch_size, n_layers, dim]的tensor，所有元素值为1e-4 + batch_size = h.size(0) + db_value = torch.full((batch_size, self.n_layers, self.params.dim), 1e-4, + dtype=h.dtype, device=h.device) + else: + # 正常模式，使用数据库查询 + index = self.extract_db.q_to_k(h) + db_value = self.extract_db.get_data(index) + h, past_kv = layer( - h,db_value, pos_cis, + h, db_value, pos_cis, past_key_value=past_key_values[l], use_cache=use_cache ) @@ -579,19 +588,22 @@ class MiniMindLM(PreTrainedModel): past_kvs.append(past_kv) h_list.append(h.unsqueeze(0)) - # 使用detach()分离计算图，避免多次反向传播 - h_tensor = torch.cat(h_list,dim=0).permute(1,0,2,3) - h_tensor_detached = h_tensor.detach() - h_tensor_detached = h_tensor_detached.reshape(h_tensor_detached.shape[0],-1,self.params.dim) + h_tensor = torch.cat(h_list, dim=0).permute(1, 0, 2, 3) - # 数据库更新逻辑与主计算图分离 - with torch.no_grad(): - # Compute shared downsampling layer once - shared_features = self.shared_downsample(h_tensor_detached) - z_v = self.downsample_v_specific(shared_features) - z_q = self.downsample_q_specific(shared_features) - z_k = self.extract_db.q_to_k(z_q) - self.extract_db.updata_value(z_k,z_v) + # 只在非禁用数据库模式下执行数据库更新逻辑 + if not self.params.disable_db: + # 使用detach()分离计算图，避免多次反向传播 + h_tensor_detached = h_tensor.detach() + h_tensor_detached = h_tensor_detached.reshape(h_tensor_detached.shape[0], -1, self.params.dim) + + # 数据库更新逻辑与主计算图分离 + with torch.no_grad(): + # Compute shared downsampling layer once + shared_features = self.shared_downsample(h_tensor_detached) + z_v = self.downsample_v_specific(shared_features) + z_q = self.downsample_q_specific(shared_features) + z_k = self.extract_db.q_to_k(z_q) + self.extract_db.updata_value(z_k, z_v) slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep logits = self.output(self.norm(h)[:, slice_indices, :]) diff --git a/train_pretrain.py b/train_pretrain.py index 92311fa..daabc37 100644 --- a/train_pretrain.py +++ b/train_pretrain.py @@ -179,7 +179,7 @@ if __name__ == "__main__": parser.add_argument("--out_dir", type=str, default="out") # 若要以最快速度实现zero则epochs设置为1轮；否则应当利用有限的数据训练2~6个epochs。 parser.add_argument("--epochs", type=int, default=3) - parser.add_argument("--batch_size", type=int, default=1) + parser.add_argument("--batch_size", type=int, default=32) parser.add_argument("--learning_rate", type=float, default=5e-4) parser.add_argument("--device", type=str, default="cuda:0" if torch.cuda.is_available() else "cpu") #如果GPU可用，则使用GPU，否则使用CPU。 parser.add_argument("--dtype", type=str, default="bfloat16") @@ -193,15 +193,22 @@ if __name__ == "__main__": parser.add_argument("--log_interval", type=int, default=100) #日志打印间隔，用于控制日志打印的频率。 parser.add_argument("--save_interval", type=int, default=100) #模型保存间隔，用于控制模型保存的频率。 parser.add_argument('--local_rank', type=int, default=-1) #本地进程编号，用于分布式训练。 - parser.add_argument('--dim', default=592, type=int) #模型维度，用于控制模型的大小。 - parser.add_argument('--n_layers', default=4, type=int) #层数，用于控制模型层数。 + parser.add_argument('--dim', default=768, type=int) #模型维度，用于控制模型的大小。 + parser.add_argument('--n_layers', default=8, type=int) #层数，用于控制模型层数。 parser.add_argument('--max_seq_len', default=512, type=int) #最大序列长度，用于控制输入序列的最大长度。 parser.add_argument('--use_moe', default=False, type=bool) #是否使用MOE，用于控制是否使用MOE。 + parser.add_argument('--disable_db', action='store_true', help="禁用数据库功能，使用固定值1e-4替代") #禁用数据库功能，启用特殊模式 parser.add_argument("--data_path", type=str, default="./dataset/pretrain_hq.jsonl") #数据路径，用于控制数据集的路径。 parser.add_argument("--pretrained_embedding_path", type=str, default=None, help="Path to pretrained token embedding weights (.pth file)") args = parser.parse_args() - lm_config = LMConfig(dim=args.dim, n_layers=args.n_layers, max_seq_len=args.max_seq_len, use_moe=args.use_moe) #创建LMConfig对象，用于控制模型配置。 + lm_config = LMConfig( + dim=args.dim, + n_layers=args.n_layers, + max_seq_len=args.max_seq_len, + use_moe=args.use_moe, + disable_db=args.disable_db # 添加禁用数据库参数 + ) #创建LMConfig对象，用于控制模型配置。 args.save_dir = os.path.join(args.out_dir) #创建保存目录。 os.makedirs(args.save_dir, exist_ok=True) #创建保存目录。 os.makedirs(args.out_dir, exist_ok=True) #创建输出目录。