DOCSEARCHbyAlgolia
Submit Docs

pruna

repository·main·Indexed 23 days ago

https://github.com/prunaai/pruna

A model optimization framework designed to make AI models faster, smaller, cheaper, and greener. Pruna provides a suite of compression algorithms—including quantization, pruning, and distillation—to optimize various model types such as LLMs, Diffusion, and Vision Transformers. It features the `smash()` function and `SmashConfig` to apply optimization techniques like deepcache and stable_fast to pre-trained models.

Tokens
37.9K
Snippets
96
Records
170
Agent score
75%

What's inside pruna

  1. Evaluate model quality with the Evaluation Agent

    main

    The Evaluation Agent in Pruna is used to assess how model optimizations (like compression) affect performance across two primary dimensions:

    • Efficiency Metrics: Measures resource usage including speed (total time, latency, throughput), memory (disk, inference, training), and energy (consumption, CO2 emissions).
    • Quality Metrics: Assesses model fidelity and performance using metrics like FID, CMMD (fidelity), Clip Score (alignment), PSNR, SSIM (diversity), and accuracy, precision, or perplexity.

    Custom metrics are also supported to tailor the evaluation to specific requirements.

  3. What is SmashConfig and how to use it

    main

    The SmashConfig class is the central configuration object used to define model optimization strategies in Pruna AI. It provides a dictionary-like interface to manage optimization algorithms, their hyperparameters, and auxiliary components like tokenizers, processors, and datasets.

    Workflow Overview:

    1. Create a SmashConfig instance.
    2. Configure algorithms and hyperparameters.
    3. Add optional components (tokenizer, processor, dataset).
    4. Pass the SmashConfig and a pre-trained model to the smash() function to receive an optimized PrunaModel.
    from pruna import SmashConfig

# Initialize with algorithm settings
smash_config = SmashConfig({"ifw": {"weight_bits": 16}})

# Add components
model_id = 'openai/whisper-tiny'
smash_config.add_tokenizer(model_id)
smash_config.add_processor(model_id)

  4. Configure optimizations with SmashConfig

    main

    The SmashConfig is a dictionary-like object used to customize optimizations. You can specify multiple algorithms from categories such as batching, caching, and quantization.

    For specific model types like Whisper, you may need to call .add_processor(model_id) and .add_tokenizer(model_id) on the SmashConfig instance.

    from pruna import SmashConfig

# Basic configuration
smash_config = SmashConfig(["hqq_diffusers"])

# Configuration for Whisper-style models
smash_config = SmashConfig(["c_whisper", "whisper_s2t"])
smash_config.add_processor(model_id)
smash_config.add_tokenizer(model_id)

  5. Selective Smashing with Target Modules

    main

    Some algorithms allow you to tailor optimization to specific parts of your model's architecture using the target_modules parameter. This parameter is passed within the algorithm's configuration dictionary inside SmashConfig.

    Configuration Format

    The target_modules parameter is a dictionary with two keys:

    • include: A list of glob pattern strings. A module is targeted if its path matches at least one pattern here.
    • exclude: A list of glob pattern strings. A module is not targeted if its path matches any pattern here.

    Logic: A module is targeted if it matches an include pattern AND does not match any exclude pattern.

    Pattern Syntax (Glob Patterns)

    • *: Matches any number of characters (e.g., attention.* matches attention.to_q).
    • ?: Matches exactly one character.
    • [abc]: Matches any single character from the set (e.g., to_[qk] matches to_q and to_k).

    Default Behavior

    • If target_modules is None: Values are inferred automatically from the model, configuration, and algorithm.
    • If include is missing: Defaults to ["*"] (all modules).
    • If exclude is missing: Defaults to [] (no modules).
    from pruna import SmashConfig

smash_config = SmashConfig({"quanto": {"target_modules": {
    "include": ["transformer.*"],
    "exclude": ["*embed*"]
}}})

  6. Understanding Metric Call Types

    main

    Metrics in Pruna operate in two modes: Single-model mode (independent scores for the model) and Pairwise mode (comparing a subsequent model against a base model). The call_type determines how data is fed into the metric.

    Call TypeDescription
    y_gtModel's output first, then ground truth
    gt_yGround truth first, then model's output
    x_gtInput data first, then ground truth
    gt_xGround truth first, then input data
    pairwise_y_gtBase model's output first, then subsequent model's output
    pairwise_gt_ySubsequent model's output first, then base model's output
    yOnly the output is used, the metric has an internal dataset

    Usage Recommendations:

    • Model vs Ground Truth (e.g., accuracy, FID): Use y_gt or gt_y.
    • Model vs Input (e.g., CLIP score): Use x_gt or gt_x.
    • Model vs Model (e.g., PSNR, SSIM, LPIPS): Use pairwise_y_gt or pairwise_gt_y.
    • Internal Dataset (e.g., ANIQA): Use y.

  7. Configure Metric Call Types (Single-Model vs Pairwise)

    main

    Stateful metrics can operate in two high-level modes via the call_type parameter:

    • Single-Model mode: Produces independent scores for the model being evaluated (e.g., comparing against ground truth). IQA metrics typically use this mode.
    • Pairwise mode: Compares a subsequent model against a base model to produce a single comparison score.

    Example usage:

    from pruna.evaluation.metrics import CMMD

# Single-mode (default)
metric = CMMD(call_type="single")

# Pairwise mode
metric = CMMD(call_type="pairwise")

  8. Understand what Pruna telemetry tracks

    main

    Pruna telemetry tracks usage metrics to help improve the package. The following data is collected:

    • Number of function executions (specifically smash, loading a model, saving a model, and calling a model).
    • The smash configuration used.
    • Whether the execution was a success or an error.

    Privacy Note: Pruna does not track any information that could identify individual users.

    Example OTLP export format:

    pruna_function_calls_total{function="test_operation",job="unknown_service",session_id="7bb23832-d733-4404-b43e-7eea8c0b872e",smash_config="the best config",status="success"} 15

  9. How EvaluationAgent modes work

    main

    The EvaluationAgent supports two distinct modes for assessing model performance:

    1. Single-Model mode: Each model is evaluated independently. Metrics are computed based solely on the model's own outputs without any reference to other models.
    2. Pairwise mode: Metrics compare the outputs of the current model against the first model evaluated by the agent. The EvaluationAgent caches the outputs of the first model to serve as a reference for all subsequent evaluations. To use this mode, ensure your metrics are configured for pairwise comparison (e.g., CMMD(call_type="pairwise")).

  10. Configure Datasets in SmashConfig

    main

    Datasets can be added to SmashConfig using three methods:

    1. String Identifiers: Use built-in dataset IDs (e.g., 'WikiText').
    2. Custom Collate Functions: Provide a (train, val, test) tuple of datasets along with a collate_fn name.
    3. PrunaDataModule: Create a custom PrunaDataModule for maximum flexibility.

  11. Choosing between BaseMetric and StatefulMetric

    main

    Pruna's evaluation system provides two base classes for implementing custom metrics, located in pruna/evaluation/metrics. Choosing the correct one depends on whether your metric requires data from multiple batches to compute a final value.

    • BaseMetric: Use this when you need to compute values directly for each instance without maintaining state across batches. This is typically used for isolated inference measurements where shared state might distort results (e.g., latency, disk_memory).
    • StatefulMetric: Use this for most quality evaluations. This class allows you to accumulate state across multiple batches (e.g., accuracy, clip_score).

    Recommendation: In most cases, you should implement a StatefulMetric.

  12. Use jaxtyping for tensor annotations

    main
    Pruna uses jaxtyping to annotate PyTorch tensors with shape and dtype information. This is primarily used in collate functions, dataloaders, and model interfaces to improve readability and catch shape/type mismatches early. Contributors should follow this style when defining tensor inputs or outputs.