Add sparse attention integration to llm_eval

Signed-off-by: Kai Xu <[email protected]>

Author: kaix-nv

examples/llm_eval/lm_eval_hf.py

@@ -43,9 +43,11 @@
 from lm_eval.api.model import T
 from lm_eval.models.huggingface import HFLM
 from quantization_utils import quantize_model
+from sparse_attention_utils import sparsify_model
 
 import modelopt.torch.opt as mto
 from modelopt.torch.quantization.utils import is_quantized
+from modelopt.torch.sparsity.attention_sparsity.conversion import is_attn_sparsified
 
 
 def create_from_arg_obj(cls: type[T], arg_dict: dict, additional_config: dict | None = None) -> T:
@@ -57,9 +59,20 @@ def create_from_arg_obj(cls: type[T], arg_dict: dict, additional_config: dict |
     calib_size = arg_dict.pop("calib_size", 512)
     compress = arg_dict.pop("compress", False)
 
+    # Sparse attention arguments
+    sparse_cfg = arg_dict.pop("sparse_cfg", None)
+
     additional_config = {} if additional_config is None else additional_config
     additional_config = {k: v for k, v in additional_config.items() if v is not None}
 
+    # Force eager attention if sparse attention is requested
+    if sparse_cfg:
+        additional_config["attn_implementation"] = "eager"
+        warnings.warn(
+            "Sparse attention requires attn_implementation='eager'. "
+            "Forcing eager attention implementation."
+        )
+
     # Enable automatic save/load of modelopt state huggingface checkpointing
     mto.enable_huggingface_checkpointing()
 
@@ -85,6 +98,15 @@ def create_from_arg_obj(cls: type[T], arg_dict: dict, additional_config: dict |
             compress=compress,
         )
 
+    if sparse_cfg:
+        if is_attn_sparsified(model_obj.model):
+            warnings.warn("Skipping sparse attention: model already has sparse attention applied.")
+        else:
+            sparsify_model(
+                model=model_obj,
+                sparse_cfg=sparse_cfg,
+            )
+
     return model_obj
 
 
@@ -120,6 +142,11 @@ def setup_parser_with_modelopt_args():
         action="store_true",
         help="Compress the model after quantization",
     )
+    parser.add_argument(
+        "--sparse_cfg",
+        type=str,
+        help="Sparse attention configuration (e.g., SKIP_SOFTMAX_DEFAULT, SKIP_SOFTMAX_CALIB)",
+    )
     return parser
 
 
@@ -142,6 +169,7 @@ def setup_parser_with_modelopt_args():
             "calib_batch_size": args.calib_batch_size,
             "calib_size": args.calib_size,
             "compress": args.compress,
+            "sparse_cfg": args.sparse_cfg,
         }
     )
 

examples/llm_eval/mmlu.py

@@ -48,6 +48,7 @@
 from fire import Fire
 from modeling import EvalModel, select_model
 from quantization_utils import MAX_SEQ_LEN, get_tokenizer, quantize_model
+from sparse_attention_utils import sparsify_model
 from tqdm import tqdm
 
 try:
@@ -56,6 +57,7 @@
     LLM = None  # type: ignore[misc]
 import modelopt.torch.opt as mto
 from modelopt.torch.quantization.utils import is_quantized
+from modelopt.torch.sparsity.attention_sparsity.conversion import is_attn_sparsified
 
 os.environ["TOKENIZERS_PARALLELISM"] = "false"
 
@@ -227,6 +229,7 @@ def main(
     batch_size: int = 0,
     calib_size: int = 512,
     dtype: str = "bfloat16",
+    sparse_cfg: str | None = None,
     **kwargs,
 ):
     random.seed(RAND_SEED)
@@ -283,6 +286,20 @@ def main(
                     auto_quantize_bits=auto_quantize_bits,
                 )
 
+        # Apply sparse attention if requested
+        if sparse_cfg:
+            model.load()
+
+            if is_attn_sparsified(model.model):
+                warnings.warn(
+                    "Skipping sparse attention: model already has sparse attention applied."
+                )
+            else:
+                sparsify_model(
+                    model=model,
+                    sparse_cfg=sparse_cfg,
+                )
+
     for subject in tqdm(subjects):
         dev_df = pd.read_csv(os.path.join(data_dir, "dev", subject + "_dev.csv"), header=None)[
             :ntrain

examples/llm_eval/sparse_attention_utils.py

@@ -0,0 +1,111 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Utilities for sparse attention integration with llm_eval."""
+
+import modelopt.torch.sparsity.attention_sparsity as mtsa
+
+# Custom sparse attention configurations
+CUSTOM_SPARSE_CONFIG = {
+    "SPARSE_CONSERVATIVE": {
+        "sparse_cfg": {
+            "*attn*": {
+                "method": "flash_skip_softmax",
+                "threshold": {"prefill": 5e-4, "decode": 1e-5},
+                "br": 128,
+                "bc": 128,
+                "backend": "pytorch",
+                "enable": True,
+            },
+            "default": {"enable": False},
+        },
+    },
+    "SPARSE_AGGRESSIVE": {
+        "sparse_cfg": {
+            "*attn*": {
+                "method": "flash_skip_softmax",
+                "threshold": {"prefill": 5e-3, "decode": 5e-4},
+                "br": 128,
+                "bc": 128,
+                "backend": "pytorch",
+                "enable": True,
+            },
+            "default": {"enable": False},
+        },
+    },
+}
+
+
+def _extract_model(model_obj):
+    """Extract actual model from wrapper (HFLM or EvalModel)."""
+    if hasattr(model_obj, "gpt2"):
+        return model_obj.gpt2
+    elif hasattr(model_obj, "model"):
+        return model_obj.model
+    else:
+        return model_obj
+
+
+def sparsify_model(
+    model,
+    sparse_cfg: str,
+    backend=None,
+):
+    """Apply sparse attention to model with optional RULER calibration.
+
+    Args:
+        model: Model wrapper (HFLM or EvalModel) or raw model
+        sparse_cfg: Sparse attention config name or dict
+        backend: Backend to use (optional, overrides config backend)
+
+    Returns:
+        The model with sparse attention applied
+
+    Note:
+        Calibration is automatically triggered if the config contains a 'calibration' field.
+        The calibration will auto-generate RULER dataset from the model's tokenizer.
+    """
+    # Extract actual model
+    net = _extract_model(model)
+
+    # Resolve config
+    if isinstance(sparse_cfg, str):
+        # Try custom configs first
+        mtsa_cfg = CUSTOM_SPARSE_CONFIG.get(sparse_cfg)
+        if mtsa_cfg is None:
+            # Try predefined configs
+            mtsa_cfg = getattr(mtsa, sparse_cfg, None)
+        if mtsa_cfg is None:
+            raise ValueError(f"Unknown sparse_cfg: {sparse_cfg}")
+    else:
+        mtsa_cfg = sparse_cfg
+
+    # Override backend if specified
+    if backend:
+        if isinstance(mtsa_cfg, dict) and "sparse_cfg" in mtsa_cfg:
+            modified_sparse_cfg = {}
+            for pattern, cfg in mtsa_cfg["sparse_cfg"].items():
+                modified_cfg = cfg.copy() if isinstance(cfg, dict) else cfg
+                if isinstance(modified_cfg, dict):
+                    modified_cfg["backend"] = backend
+                modified_sparse_cfg[pattern] = modified_cfg
+            mtsa_cfg = {"sparse_cfg": modified_sparse_cfg}
+
+    # Apply sparsification
+    print(f"\nApplying sparse attention with config: {sparse_cfg}")
+    mtsa.sparsify(net, mtsa_cfg)
+    print("Sparse attention applied successfully!")
+
+    return model

examples/llm_sparsity/attention_sparsity/hf_sa.py

@@ -196,29 +196,6 @@ def generate_text(model, inputs, args, tokenizer):
         print("\nOutputs differ")
 
 
-def sparsify_model(model, args):
-    """Apply sparse attention to the model with optional calibration."""
-    print(f"\nApplying sparse attention: {args.sparse_attn} with backend: {args.backend}")
-    base_config = SPARSE_ATTN_CFG_CHOICES[args.sparse_attn]
-
-    # Create modified config with selected backend
-    modified_sparse_cfg = {}
-    for pattern, cfg in base_config["sparse_cfg"].items():
-        modified_cfg = cfg.copy()
-        modified_cfg["backend"] = args.backend
-        modified_sparse_cfg[pattern] = modified_cfg
-
-    # Create new config with modified settings
-    sparse_config = SparseAttentionConfig(sparse_cfg=modified_sparse_cfg)
-
-    # Sparsify the model
-    model = mtsa.sparsify(model, config=sparse_config)
-
-    print("Sparse attention applied successfully!")
-
-    return model
-
-
 def main(args):
     """Main function to run the selected mode."""
     if not torch.cuda.is_available():
@@ -249,8 +226,22 @@ def main(args):
         model = model.cuda()
         print("Model moved to CUDA")
 
-    # Apply sparse attention to the model (with calibration if configured)
-    model = sparsify_model(model, args)
+    # Apply sparse attention with optional calibration
+    print(f"\nApplying sparse attention: {args.sparse_attn} with backend: {args.backend}")
+    base_config = SPARSE_ATTN_CFG_CHOICES[args.sparse_attn]
+
+    # Create modified config with selected backend
+    modified_sparse_cfg = {}
+    for pattern, cfg in base_config["sparse_cfg"].items():
+        modified_cfg = cfg.copy()
+        modified_cfg["backend"] = args.backend
+        modified_sparse_cfg[pattern] = modified_cfg
+
+    # Create config and apply sparsification (calibration happens automatically)
+    sparse_config = SparseAttentionConfig(sparse_cfg=modified_sparse_cfg)
+    model = mtsa.sparsify(model, config=sparse_config)
+
+    print("Sparse attention applied successfully!")
 
     # Verify outputs if requested (compares baseline vs calibrated sparse model)
     if args.verify_output:

modelopt/torch/sparsity/attention_sparsity/calibration/calibrate.py

@@ -24,6 +24,7 @@
 from transformers import AutoTokenizer
 
 from ..config import CalibrationConfig
+from ..conversion import print_sparse_attention_summary
 from ..sparse_attention import SparseAttentionModule
 from .calibrator import DynamicThresholdCalibrator
 from .dataset import RulerDatasetBuilder
@@ -162,6 +163,10 @@ def calibrate_sparse_attention(
     )
     calibration_result = calibrator.calibrate(model, forward_loop)
 
+    # Print calibration statistics (regardless of success/failure for debugging)
+    print("\nCalibration complete!")
+    print_sparse_attention_summary(model)
+
     if "scale_factor" not in calibration_result:
         warnings.warn("Calibration did not produce valid results")
         return {}

modelopt/torch/sparsity/attention_sparsity/calibration/dataset.py

@@ -185,7 +185,7 @@ def __init__(
         self,
         samples: int,
         max_seqlen: int,
-        tokenizer_name_or_path: str,
+        tokenizer_name_or_path: str | object,
         seed: int = 42,
         num_length_bins: int = 4,
         max_length_filter: int = 65536,
@@ -195,7 +195,7 @@ def __init__(
         Args:
             samples: Total number of samples to generate (distributed evenly across length bins)
             max_seqlen: Maximum sequence length (length bins auto-generated as powers of 2)
-            tokenizer_name_or_path: HuggingFace tokenizer path
+            tokenizer_name_or_path: HuggingFace tokenizer path or tokenizer object
             seed: Random seed for reproducibility
             num_length_bins: Number of length bins to generate (default: 4)
             max_length_filter: Maximum sequence length to keep (default: 65536)
@@ -229,8 +229,11 @@ def __init__(
         # Distribute samples evenly across lengths
         self.samples_per_length = [samples // len(self.target_lengths)] * len(self.target_lengths)
 
-        # Initialize tokenizer and seed
-        self.tokenizer = AutoTokenizer.from_pretrained(tokenizer_name_or_path)
+        # Initialize tokenizer
+        if isinstance(tokenizer_name_or_path, str):
+            self.tokenizer = AutoTokenizer.from_pretrained(tokenizer_name_or_path)
+        else:
+            self.tokenizer = tokenizer_name_or_path
         random.seed(seed)
 
     def build_calibration_dataset(self) -> list[dict[str, Any]]:
@@ -247,9 +250,7 @@ def build_calibration_dataset(self) -> list[dict[str, Any]]:
             desc="Generating RULER calibration samples",
             total=len(self.target_lengths),
         ):
-            samples_per_task = num_samples // len(self.subtasks)
-            if samples_per_task <= 0:
-                continue
+            samples_per_task = max(num_samples // len(self.subtasks), 1)
 
             # Generate equal samples for each task
             for task_name in self.subtasks:

modelopt/torch/sparsity/attention_sparsity/config.py

@@ -77,6 +77,12 @@ class SparseAttentionAttributeConfig(ModeloptBaseConfig):
         ),
     )
 
+    collect_stats: bool = ModeloptField(
+        default=False,
+        title="Collect statistics.",
+        description="Whether to collect sparsity statistics during forward pass for monitoring.",
+    )
+
     is_causal: bool = ModeloptField(
         default=True,
         title="Causal attention flag.",
@@ -254,6 +260,7 @@ def validate_num_length_bins(cls, v):
             "br": 128,  # Flash Attention block rows
             "bc": 128,  # Flash Attention block columns
             "backend": "pytorch",  # Only pytorch backend supported
+            "collect_stats": True,
             "enable": True,
         },
         "default": {"enable": False},
@@ -298,6 +305,7 @@ class FlashSkipSoftmaxConfig(SparseAttentionConfig):
                 "br": 128,  # Flash Attention block rows
                 "bc": 128,  # Flash Attention block columns
                 "backend": "pytorch",  # Only pytorch backend supported
+                "collect_stats": True,  # Enable statistics collection
                 "enable": True,
             },
             "default": {"enable": False},
@@ -319,6 +327,7 @@ class FlashSkipSoftmaxConfig(SparseAttentionConfig):
             "br": 128,
             "bc": 128,
             "backend": "pytorch",  # Only pytorch backend supported
+            "collect_stats": True,
             "enable": True,
             "calibration": {
                 "target_sparse_ratio": 0.3,