fix(trainer): 修复多轮对话中的因果注意力掩码计算逻辑等

2025-09-28 15:15:19 +08:00
parent 0b96b11a6e
commit 1169cfad82
3 changed files with 15 additions and 3 deletions
@@ -99,6 +99,8 @@ def process_attention_mask(
        return None
    
    if seq_mask.dim() > 2:
+        # shape (bsz, seq_len) or (bsz,n_heads, seq_len, seq_len + start_pos)
+        # if ndim > 2, it's 4D tensor
        return seq_mask
    
    batch_size = seq_mask.size(0)