Training checkpoint at step 2000

modeling_armt.py

@@ -440,13 +440,11 @@ def attn_mask_to_4d(attn_mask, upper, query_len):
     return mask
 
 def invert_attn_mask(attn_mask, dtype):
-    if os.environ.get("NOT_INVERT_ATTN_MASK"):
-        return attn_mask
-    min_dtype = torch.finfo(dtype).min
-    # Use the same dtype as attn_mask to avoid dtype conversion
-    one = torch.tensor(1.0, dtype=attn_mask.dtype, device=attn_mask.device)
-    new_mask = (one - attn_mask) * min_dtype
-    return new_mask
+        min_dtype = torch.finfo(dtype).min
+        # Use the same dtype as attn_mask to avoid dtype conversion
+        one = torch.tensor(1.0, dtype=attn_mask.dtype, device=attn_mask.device)
+        new_mask = (one - attn_mask) * min_dtype
+        return new_mask
 
 
 
@@ -1856,6 +1854,8 @@ def is_empty_past_key_values(past_key_values: Optional[DynamicCache], layer_idx:
         return True
     return False
 
+invert_attn_mask = lambda mask, dtype: (_invert_attn_mask(mask, dtype) if not os.environ.get("NOT_INVERT_ATTN_MASK") else mask)
+
 def segment_tensor(t: torch.Tensor, start_idx: int, end_idx: int, seq_len: int) -> torch.Tensor:
     if not isinstance(t, torch.Tensor):
         return t