Address review feedbacks

Signed-off-by: Kai Xu <kaix@nvidia.com>

Author: kaix-nv

examples/llm_sparsity/attention_sparsity/README.md

@@ -1,6 +1,6 @@
 # Attention Sparsity for HuggingFace Models
 
-In this tutorial, we demonstrate how to use NVIDIA TensorRT Model Optimizer to apply attention sparsity to HuggingFace models. Attention sparsity reduces computational cost by skipping near-zero attention scores during the softmax computation.
+In this tutorial, we demonstrate how to use NVIDIA Model Optimizer to apply attention sparsity to HuggingFace models. Attention sparsity reduces computational cost by skipping near-zero attention scores during the softmax computation.
 
 ## Getting Started
 
@@ -63,7 +63,7 @@ pip install nvidia-modelopt[hf]
 If using `SKIP_SOFTMAX_CALIB`, you need to download the RULER calibration dataset first:
 
 ```bash
-bash modelopt/torch/sparsity/attention_sparsity/calibration/download_ruler_data.sh
+bash ./download_ruler_data.sh
 ```
 
 This downloads the Paul Graham essays dataset used for generating calibration samples.
@@ -75,7 +75,7 @@ This downloads the Paul Graham essays dataset used for generating calibration sa
 Apply sparse attention with a fixed threshold:
 
 ```bash
-python examples/llm_sparsity/attention_sparsity/hf_sa.py \
+python hf_sa.py \
     --pyt_ckpt_path Qwen/Qwen3-8B \
     --sparse_attn skip_softmax
 ```
@@ -85,7 +85,7 @@ python examples/llm_sparsity/attention_sparsity/hf_sa.py \
 Apply sparse attention with calibrated thresholds for optimal sparsity:
 
 ```bash
-python examples/llm_sparsity/attention_sparsity/hf_sa.py \
+python hf_sa.py \
     --pyt_ckpt_path Qwen/Qwen3-8B \
     --sparse_attn skip_softmax_calib
 ```
@@ -121,7 +121,7 @@ The script automatically compares outputs before and after applying sparse atten
 Export the sparsified model to a HuggingFace checkpoint:
 
 ```bash
-python examples/llm_sparsity/attention_sparsity/hf_sa.py \
+python hf_sa.py \
     --pyt_ckpt_path Qwen/Qwen3-8B \
     --sparse_attn skip_softmax_calib \
     --export_dir ./exported_sparse_model
@@ -161,5 +161,5 @@ model = mtsa.sparsify(model, config=custom_config)
 
 ## References
 
-- [TensorRT Model Optimizer Documentation](https://nvidia.github.io/TensorRT-Model-Optimizer/)
+- [Model Optimizer Documentation](https://nvidia.github.io/Model-Optimizer/)
 - [RULER: What's the Real Context Size of Your Long-Context Language Models?](https://github.com/NVIDIA/RULER)

examples/llm_sparsity/attention_sparsity/hf_sa.py

@@ -171,7 +171,7 @@ def main(args):
     print(f"\nApplying sparse attention: {args.sparse_attn}")
     sparse_config = SPARSE_ATTN_CFG_CHOICES[args.sparse_attn]
 
-    # Override target_sparse_ratio if provided via CLI
+    # Override calibration options if provided via CLI
     if args.target_sparse_ratio is not None:
         sparse_config = copy.deepcopy(sparse_config)
         sparse_cfg = sparse_config.get("sparse_cfg", {})

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

@@ -17,7 +17,7 @@
 
 from .calibrate import calibrate_sparse_attention
 from .calibrator import DynamicThresholdCalibrator
-from .dataset import RulerDatasetBuilder
+from .ruler_dataset import RulerDatasetBuilder
 
 __all__ = [
     "DynamicThresholdCalibrator",

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

@@ -15,70 +15,29 @@
 
 """Calibration functions for sparse attention."""
 
-import hashlib
-import json
 import warnings
 from collections.abc import Callable
-from pathlib import Path
 from typing import Any
 
 import torch
 import torch.nn as nn
 from transformers import AutoTokenizer
 
+from modelopt.torch.utils import get_module_device
+
 from ..config import CalibrationConfig
 from ..conversion import print_sparse_attention_summary
-from ..sparse_attention import SparseAttentionModule
+from ..utils import get_named_sparse_attention_modules
 from .calibrator import DynamicThresholdCalibrator
-from .dataset import RulerDatasetBuilder
-
+from .ruler_dataset import RulerDatasetBuilder
 
-def _get_cache_path(
-    tokenizer_path: str, samples: int, max_seqlen: int, cache_dir: str | None = None
-) -> Path:
-    """Generate cache file path based on calibration parameters.
 
-    Args:
-        tokenizer_path: Path to tokenizer (used in hash)
-        samples: Number of calibration samples
-        max_seqlen: Maximum sequence length
-        cache_dir: Optional cache directory. If None, uses ~/.cache/modelopt/sparse_attention/
-    """
-    # Create a hash of the parameters for the cache filename
-    key = f"{tokenizer_path}_{samples}_{max_seqlen}"
-    hash_str = hashlib.md5(key.encode(), usedforsecurity=False).hexdigest()[:12]
-    filename = f"ruler_cache_{samples}s_{max_seqlen}l_{hash_str}.json"
-
-    if cache_dir:
-        base_dir = Path(cache_dir)
-    else:
-        base_dir = Path.home() / ".cache" / "modelopt" / "sparse_attention"
-
-    return base_dir / filename
-
-
-def _load_cached_data(cache_path: Path) -> list[dict[str, Any]] | None:
-    """Load calibration data from cache if it exists."""
-    if cache_path.exists():
-        try:
-            with open(cache_path) as f:
-                data = json.load(f)
-            print(f"Loaded {len(data)} cached calibration samples from {cache_path}")
-            return data
-        except Exception as e:
-            print(f"Warning: Failed to load cache: {e}")
-    return None
-
-
-def _save_cached_data(cache_path: Path, data: list[dict[str, Any]]) -> None:
-    """Save calibration data to cache."""
-    try:
-        cache_path.parent.mkdir(parents=True, exist_ok=True)
-        with open(cache_path, "w") as f:
-            json.dump(data, f)
-        print(f"Saved calibration samples to cache: {cache_path}")
-    except Exception as e:
-        print(f"Warning: Failed to save cache: {e}")
+def _load_tokenizer(tokenizer_name_or_path: str) -> "AutoTokenizer":
+    """Load tokenizer and ensure pad_token is set."""
+    tokenizer = AutoTokenizer.from_pretrained(tokenizer_name_or_path)
+    if not tokenizer.pad_token:
+        tokenizer.pad_token = tokenizer.eos_token
+    return tokenizer
 
 
 def _extract_tokenizer_from_model(model: nn.Module) -> str:
@@ -147,12 +106,10 @@ def create_calibration_forward_loop(
     Returns:
         Forward loop function that takes model as argument
     """
-    tokenizer = AutoTokenizer.from_pretrained(tokenizer_name_or_path)
-    if not tokenizer.pad_token:
-        tokenizer.pad_token = tokenizer.eos_token
+    tokenizer = _load_tokenizer(tokenizer_name_or_path)
 
     def forward_loop(model: nn.Module) -> None:
-        device = next(model.parameters()).device
+        device = get_module_device(model)
 
         for sample in calibration_data:
             inputs = tokenizer(
@@ -205,12 +162,10 @@ def create_decode_calibration_forward_loop(
     Returns:
         Forward loop function that takes model as argument
     """
-    tokenizer = AutoTokenizer.from_pretrained(tokenizer_name_or_path)
-    if not tokenizer.pad_token:
-        tokenizer.pad_token = tokenizer.eos_token
+    tokenizer = _load_tokenizer(tokenizer_name_or_path)
 
     def forward_loop(model: nn.Module) -> None:
-        device = next(model.parameters()).device
+        device = get_module_device(model)
 
         for sample in calibration_data:
             inputs = tokenizer(
@@ -291,9 +246,7 @@ def calibrate_sparse_attention(
         return {}
 
     # Get sparse attention modules
-    sparse_modules = [
-        (name, m) for name, m in model.named_modules() if isinstance(m, SparseAttentionModule)
-    ]
+    sparse_modules = get_named_sparse_attention_modules(model)
 
     if not sparse_modules:
         print("No sparse attention modules found for calibration")
@@ -306,29 +259,16 @@ def calibrate_sparse_attention(
     calibration_data = None
 
     if calibrate_prefill or calibrate_decode:
-        # Try to load from cache first
-        cache_path = _get_cache_path(
-            tokenizer,
-            calib_config.samples,
-            calib_config.max_seqlen,
+        builder = RulerDatasetBuilder(
+            samples=calib_config.samples,
+            max_seqlen=calib_config.max_seqlen,
+            tokenizer_name_or_path=tokenizer,
+            num_length_bins=calib_config.num_length_bins,
+            max_length_filter=int(calib_config.max_seqlen * 1.5),
             cache_dir=calib_config.cache_dir,
+            data_dir=calib_config.data_dir,
         )
-        calibration_data = _load_cached_data(cache_path)
-
-        # Generate if not cached
-        if calibration_data is None:
-            builder = RulerDatasetBuilder(
-                samples=calib_config.samples,
-                max_seqlen=calib_config.max_seqlen,
-                tokenizer_name_or_path=tokenizer,
-                num_length_bins=calib_config.num_length_bins,
-                max_length_filter=int(calib_config.max_seqlen * 1.5),
-            )
-            calibration_data = builder.build_calibration_dataset()
-            print(f"Generated {len(calibration_data)} calibration samples")
-
-            # Save to cache for future runs
-            _save_cached_data(cache_path, calibration_data)
+        calibration_data = builder.build_calibration_dataset()
 
     # Initialize results
     calibration_results: dict[str, Any] = {}

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

@@ -26,8 +26,8 @@
 from scipy.optimize import curve_fit
 from tqdm import tqdm
 
-from ..sparse_attention import SparseAttentionModule
 from ..stats_manager import SparseAttentionStatsManager
+from ..utils import get_sparse_attention_modules
 
 
 class DynamicThresholdCalibrator:
@@ -113,7 +113,7 @@ def calibrate(self, model: nn.Module, forward_loop: Callable, phase: str) -> dic
             Dict with calibration results including a, b, r_squared, and num_data_points
         """
         # Extract attention modules
-        attention_modules = [m for m in model.modules() if isinstance(m, SparseAttentionModule)]
+        attention_modules = get_sparse_attention_modules(model)
 
         if not attention_modules:
             raise ValueError("No sparse attention modules found for calibration")