iomgaa/Minimind
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

添加了忽视数据库模式

Browse Source
This commit is contained in:
iomgaa 2025-05-09 15:19:41 +08:00
parent b6bd97aaaa
commit 0c8c6e5d1a
3 changed files with 48 additions and 21 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
model/LMConfig.py
View File

@ -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
####################################################

24
model/model.py
View File

@ -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):
# 禁用数据库模式，使用固定值替代数据库查询
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,10 +588,13 @@ class MiniMindLM(PreTrainedModel):
past_kvs.append(past_kv)
h_list.append(h.unsqueeze(0))
h_tensor = torch.cat(h_list, dim=0).permute(1, 0, 2, 3)
# 只在非禁用数据库模式下执行数据库更新逻辑
if not self.params.disable_db:
# 使用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_detached = h_tensor_detached.reshape(h_tensor_detached.shape[0], -1, self.params.dim)
# 数据库更新逻辑与主计算图分离
with torch.no_grad():
@ -591,7 +603,7 @@ class MiniMindLM(PreTrainedModel):
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)
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, :])

15
train_pretrain.py
View File

@ -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) #创建输出目录。