Add logging of density and local attention error

Author: apd10

benchmark/executor.py

@@ -16,6 +16,7 @@
 from typing import Any, Dict, List, Optional, Tuple
 from queue import Empty
 from contextlib import contextmanager
+from sparse_attention_hub.metric_logging.logger import MicroMetricLogger
 
 # Set multiprocessing start method to 'spawn' for CUDA compatibility
 if multiprocessing.get_start_method(allow_none=True) != 'spawn':
@@ -233,12 +234,15 @@ def _benchmark_worker(
 
                     # Execute benchmark
                     logger.info(f"Worker {worker_id}: Executing benchmark {stub.benchmark_name} on GPU {current_gpu_id}")
+                    metric_logger = MicroMetricLogger()
+                    metric_logger.configure_logging(log_path=stub.result_dir, enabled_metrics=["research_attention_density", "research_attention_output_error"])
                     metrics = benchmark.run_benchmark(
                         adapter=adapter,
                         result_dir=stub.result_dir,
                         generation_kwargs=stub.generation_kwargs,
                         request_kwargs=stub.request_kwargs
                     )
+                    metric_logger.flush()
 
                     execution_time = time.time() - start_time
                     execution_success = True

sparse_attention_hub/metric_logging/logger.py

@@ -36,7 +36,7 @@ def __new__(cls, *args, **kwargs) -> "MicroMetricLogger":
 
     def __init__(self, 
                  log_path: Optional[str] = None,
-                 flush_every: int = 100,  # Flush every N events
+                 flush_every: int = 1000,  # Flush every N events
                  flush_interval: float = 60.0,  # Flush every N seconds
                  enabled_metrics: Union[List[str], str] = None):  # List of string identifiers to enable, or "all"
         if not self._initialized:
@@ -181,8 +181,7 @@ def flush(self) -> None:
             return
 
         # Get current timestamp for filename
-        timestamp = datetime.now().strftime("%Y%m%d")
-        filename = f"metrics_{timestamp}.jsonl"
+        filename = f"micro_metrics.jsonl"
         filepath = os.path.join(self.log_path, filename)
 
         # Write events to file

sparse_attention_hub/sparse_attention/research_attention/base.py

@@ -8,10 +8,13 @@
 
 from ..base import SparseAttention, SparseAttentionConfig
 from ..utils.mask import Mask
-from ..utils.mask_attention_utils import get_masked_attention_output
+from ..utils.mask_attention_utils import get_masked_attention_output, get_true_attention_output
 from .maskers.base import MaskerConfig, ResearchMasker
 from .maskers.sampling.base import SamplingMasker
 
+from sparse_attention_hub.metric_logging.logger import MicroMetricLogger
+MicroMetricLogger.register_metric("research_attention_density", float)
+MicroMetricLogger.register_metric("research_attention_output_error", float)
 
 @dataclass
 class ResearchAttentionConfig(SparseAttentionConfig):
@@ -101,6 +104,9 @@ def custom_attention(
                 previous_mask=sparse_attention_mask,
                 **kwargs,
             )
+            
+        if MicroMetricLogger().is_metric_enabled("research_attention_density"):
+            MicroMetricLogger().log("research_attention_density", sparse_attention_mask.get_density(), metadata={"layer_idx" : kwargs["layer_idx"]})
 
         # Call compute_masked_attention_output on the result of the last mask
         # Always request attention weights to match the expected return signature
@@ -118,6 +124,12 @@ def custom_attention(
             return_attention_weights=True,
             **kwargs,
         )
+
+        if MicroMetricLogger().is_metric_enabled("research_attention_output_error"):
+            true_attention_output, _ = get_true_attention_output(module, queries, keys, values, attention_mask, scaling, dropout, **kwargs)
+            error = torch.norm(true_attention_output - attention_output) / torch.norm(true_attention_output)
+            MicroMetricLogger().log("research_attention_output_error", float(error.item()), metadata={"layer_idx" : kwargs["layer_idx"]})
+
         return attention_output, attention_weights
 
     @classmethod

sparse_attention_hub/sparse_attention/utils/mask_attention_utils.py

@@ -9,6 +9,50 @@
 from .mask import Mask
 
 
+def get_true_attention_output(
+    module: nn.Module,
+    queries: torch.Tensor,
+    keys: torch.Tensor,
+    values: torch.Tensor,
+    attention_mask: Optional[torch.Tensor],
+    scaling: float,
+    dropout: float,
+    **kwargs: Dict[str, Any],
+) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
+    """Get the true (dense) attention output from the module.
+
+    Args:
+        module: The attention module (used for dropout training flag).
+        queries: Query tensor of shape (..., seq_len_q, d_k).
+        keys: Key tensor of shape (..., seq_len_k, d_k).
+        values: Value tensor of shape (..., seq_len_k, d_v).
+        attention_mask: Optional mask tensor to apply to attention weights.
+        scaling: Scaling factor for attention logits.
+        dropout: Dropout probability for attention weights.
+        **kwargs: Additional keyword arguments (unused).
+
+    Returns:
+        Tuple containing:
+            - attention_output: Output tensor after applying attention.
+            - attention_weights: Softmax-normalized attention weights.
+    """
+    num_key_value_groups: int = _get_num_key_value_groups(queries, keys)
+    key_states = repeat_kv(keys, num_key_value_groups)
+    value_states = repeat_kv(values, num_key_value_groups)
+
+    attn_weights = torch.matmul(queries, key_states.transpose(2, 3)) * scaling
+    if attention_mask is not None:
+        causal_mask = attention_mask[:, :, :, : key_states.shape[-2]]
+        attn_weights = attn_weights + causal_mask
+
+    attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(queries.dtype)
+    attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training)
+    attn_output = torch.matmul(attn_weights, value_states)
+    attn_output = attn_output.transpose(1, 2).contiguous()
+
+    return attn_output, attn_weights
+    
+
 def apply_inv_mask_sum(input_tensor: torch.Tensor, mask: Mask) -> torch.Tensor:
     """Apply inverse mask to input tensor and sum along the last dimension.