PrunaAI / PrunaAI/pruna

Use `dtype` instead of deprecated `torch_dtype` for transformers >= 4.56

Open Beginner friendly
#734 1 comment 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

stale
Dominant language
Python
Stars
1.3k
Forks
109
Avg merge
3d 9h
Merged PRs (30d)
3

Description

In src/pruna/algorithms/hqq.py line 272, AutoModelForCausalLM.from_pretrained is called with torch_dtype=...:

model = AutoModelForCausalLM.from_pretrained(
    temp_dir, quantization_config=..., torch_dtype=torch.float16 if ... else ...
)

The torch_dtype keyword argument was deprecated in transformers 4.56 (PR #39782) and replaced by dtype. On transformers 4.56+ this call emits a DeprecationWarning, and the argument will be removed in a future release, breaking HQQ quantization.

Suggested fix: choose the keyword based on the installed transformers version with packaging.version:

import transformers
from packaging.version import Version

def _dtype_kwargs(dtype):
    """`dtype` keyword of `from_pretrained` exists since transformers 4.56 (PR #39782);
    older versions use `torch_dtype`."""
    if Version(transformers.__version__) >= Version("4.56"):
        return {"dtype": dtype}
    return {"torch_dtype": dtype}

model = AutoModelForCausalLM.from_pretrained(
    temp_dir, quantization_config=...,
    **_dtype_kwargs(torch.float16 if ... else ...),
)

This keeps compatibility with transformers < 4.56 and stops the deprecation warning on 4.56+.

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start in src/pruna/algorithms/hqq.py at line 272 and inspect the AutoModelForCausalLM.from_pretrained call and its current dtype handling. Check the installed transformers version comparison approach described in the issue, then verify that transformers before 4.56 still receive torch_dtype while 4.56+ receives dtype and no deprecation warning is emitted.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
machine-learning
Issue type
Bug
Difficulty
2/5
Estimated time
1-3 hours
Activity status
Active
Clarity
Clearly specified
Newbie friendliness
78/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.