fix(torchao): update imports of quantizer

[ParagEkbote]

Description

After the release of torchao 0.15.0, the config function which are used in the torchao quantizer have been deprecated and removed, instead being updated with a new class config. We can also view this in the release notes. I have updated the imports to reflect the changes. Could you please review?

cc: @minettekaum

Related Issue

Fixes #(issue number)

Type of Change

• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• This change requires a documentation update

How Has This Been Tested?

Checklist

• My code follows the style guidelines of this project
• I have performed a self-review of my own code
• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes

Additional Notes

[ParagEkbote]

@cursor review

[cursor]

✅ Bugbot reviewed your changes and found no new issues!

Comment @cursor review or bugbot run to trigger another review on this PR

[ParagEkbote]

Could you please review the changes?

cc: @minettekaum

[minettekaum]

Hi @ParagEkbote,

Thanks for the updates! Everything looks good to me :D

I made a small change to the PR title so it would pass the check.

Cheers!

[ParagEkbote]

Is the PR ready, or are there any additional maintainer review needed?

[minettekaum]

Hi! The PR is ready 😊 We're just doing a few updates on the Pruna repo right now, so merging is temporarily blocked. You should be able to merge it at the beginning of next week. I can tag you in a comment once merging is available again!

[begumcig]

Hey Parag, thank you so much for handling this issue! Everything looks almost perfect to me, could you please check this with the oldest version of torch + torchao we currently support in pruna to see if it still works? If not we might need to add a dynamic import check.

[minettekaum]

Hi @ParagEkbote, the release is done. Before you merge, could you please check @begumcig's comment 😊

[ParagEkbote]

I have tested the changes with the following script and it seems to work correctly:

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from pruna import SmashConfig, PrunaModel

MODEL_ID = "HuggingFaceTB/SmolLM2-135M-Instruct"

tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)

base_model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    dtype=torch.bfloat16,
    device_map="auto"
)

# ---- CORRECT CONFIG ----
smash_config = SmashConfig.from_list(
    ["torchao"],   # REQUIRED: registers quantizer
    batch_size=1,
    device="cuda"
)

# Add TorchAO params
smash_config.add({
    "torchao_quant_type": "int8wo",
    "torchao_excluded_modules": "norm+embedding",
    "torchao_target_modules": {
        "include": ["model.layers.*"]
    }
})

# ---- Wrap model ----
model = PrunaModel(base_model, smash_config=smash_config)
model.set_to_eval()

# ---- Inference ----
prompt = "Explain quantization trade-offs briefly."
inputs = tokenizer(prompt, return_tensors="pt").to("cuda")

with torch.no_grad():
    outputs = model.model.generate(
        **inputs,
        max_new_tokens=75,
        temperature=0.2,
        top_p=0.9,
        eos_token_id=tokenizer.eos_token_id,
        pad_token_id=tokenizer.eos_token_id,
    )

print(tokenizer.decode(outputs[0], skip_special_tokens=True))

Could you please review?

cc: @begumcig

[begumcig]

Hi @ParagEkbote thank you so much for testing this! The script is super helpful to confirm the PR works in your environment, I am still specifically worried about compatibility with the oldest torch + torchao versions we support in Pruna. Could you please share which versions you tested, and also verify against the minimum supported versions? If the new config-class imports fail there, we’ll probably need a version check or dynamic import fallback. Thank youu!

[ParagEkbote]

Thanks for the clarification, I did test with different versions as shown in the compatibility table but I did not get a torchao and torch API mismatch.

I tested it with the following versions:

torchao==0.15.0 torch==2.10.0 (default)

torchao==0.15.0 torch==2.9.1 torchvision==0.24.1 torchaudio==2.9.1

torchao==0.14.1 torch==2.9.0 torchvision==0.24.0 torchaudio==2.9.0

torchao==0.13.0 torch==2.8.0 torchvision==0.23.0 torchaudio==2.8.0

torchao==0.12.0 torch==2.7.0 torchvision==0.22.0 torchaudio==2.7.0

Could you please review?

[begumcig]

You dropped this 👑👑👑 @ParagEkbote Thank you so much for handling this! Let's merge it