fix: safe attention mask and auto-dequantization

modeling_cerebellum.py

@@ -117,7 +117,8 @@ def bidirectional_state_branch_mask_batch(segs, device, dtype=torch.bfloat16):
     allow = allow | torch.eye(L, dtype=torch.bool, device=device)[None]
     
     mask = torch.zeros((len(segs), 1, L, L), dtype=dtype, device=device)
-    mask.masked_fill_(~allow[:, None, :, :], torch.finfo(dtype).min)
+    mask_min = -1e4 if dtype == torch.float16 else -1e9
+    mask.masked_fill_(~allow[:, None, :, :], mask_min)
     return mask
 
 class CerebellumModel(nn.Module):