fix : 使用 bool 注意力掩码并支持打包 SFT 文档边界阻断

- 简化 process_attention_mask，通过广播返回 bool 掩码
- 新增 make_doc_boundary_mask 生成块对角因果掩码
- SFT strategy 传入文档边界掩码

astrai/model/transformer.py

@@ -26,24 +26,21 @@ def process_attention_mask(
         return input_mask
 
     device = input_tensor.device
-    dtype = input_tensor.dtype
-    B, S = input_tensor.size()[:2]
+    B = input_tensor.size(0)
     T = position_ids.max().item() + 1
 
     if input_mask is None:
         if position_ids.min().item() == 0 and is_causal:
             return None
-        pad = torch.ones(B, T, dtype=torch.bool, device=device)
+        attend = torch.ones(B, 1, T, dtype=torch.bool, device=device)
     else:
-        pad = input_mask[:, :T].to(device=device, dtype=torch.bool)
+        attend = input_mask[:, :T].to(device=device, dtype=torch.bool).unsqueeze(1)
 
-    attend = pad.view(B, 1, T).expand(B, S, T).clone()
     if is_causal:
-        attend &= position_ids.unsqueeze(-1) >= torch.arange(T, device=device)
+        causal = position_ids.unsqueeze(-1) >= torch.arange(T, device=device)
+        attend = attend & causal
 
-    return torch.full(
-        (B, 1, S, T), -torch.finfo(dtype).max / 2, dtype=dtype, device=device
-    ).masked_fill_(attend.unsqueeze(1), 0.0)
+    return attend.unsqueeze(1)
 
 
 @AutoModel.register("autoregressive_lm")