RYDE-WORK/lnp_ml
1
0
Fork 0
mirror of https://github.com/RYDE-WORK/lnp_ml.git synced 2026-07-22 05:49:55 +08:00
Code Issues Packages Projects Releases Wiki Activity

feat: 同步 RAG 融合与训练核心代码(fusion 支持 f_retr、models/trainer/moe/dataset)

Browse Source
This commit is contained in:
DicongLi 2026-07-05 01:07:26 +08:00
parent 3898fb8973
commit 987ecf6386
2 changed files with 383 additions and 379 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
lnp_ml/modeling/layers/fusion.py
View File

@ -128,6 +128,7 @@ class ResidualConcatFusion(nn.Module):
# 零初始化门控（可学习标量），旁路初始不参与 # 零初始化门控（可学习标量），旁路初始不参与
self.g_moe = nn.Parameter(torch.zeros(())) self.g_moe = nn.Parameter(torch.zeros(()))
self.g_llm = nn.Parameter(torch.zeros(())) self.g_llm = nn.Parameter(torch.zeros(()))
self.g_retr = nn.Parameter(torch.zeros(())) # 检索旁路零初始化门控
def forward( def forward(
self, self,
@ -135,6 +136,7 @@ class ResidualConcatFusion(nn.Module):
tab: torch.Tensor, tab: torch.Tensor,
f_moe: Optional[torch.Tensor] = None, f_moe: Optional[torch.Tensor] = None,
f_llm: Optional[torch.Tensor] = None, f_llm: Optional[torch.Tensor] = None,
f_retr: Optional[torch.Tensor] = None,
return_attn_weights: bool = False, return_attn_weights: bool = False,
) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]: ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
# 只对真实 token（chem + tab）做注意力池化，旁路不参与 softmax 竞争 # 只对真实 token（chem + tab）做注意力池化，旁路不参与 softmax 竞争
@ -148,5 +150,7 @@ class ResidualConcatFusion(nn.Module):
out = out + self.g_moe * f_moe # 残差 + 零初始化门 out = out + self.g_moe * f_moe # 残差 + 零初始化门
if f_llm is not None: if f_llm is not None:
out = out + self.g_llm * f_llm out = out + self.g_llm * f_llm
if f_retr is not None:
out = out + self.g_retr * f_retr # 检索旁路，零初始化门保证起点=不开
return (out, attn) if return_attn_weights else out return (out, attn) if return_attn_weights else out