Skip to content

Use dtype instead of deprecated torch_dtype for transformers >= 4.56 #977

Description

@xyf5432

In contrib/rlhflow/reward_modeling.py line 44, AutoModelForSequenceClassification.from_pretrained is called with torch_dtype=torch.bfloat16:

model = AutoModelForSequenceClassification.from_pretrained(
    model_args.model_name_or_path, num_labels=1, torch_dtype=torch.bfloat16
)

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 the script.

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 = AutoModelForSequenceClassification.from_pretrained(
    model_args.model_name_or_path, num_labels=1, **_dtype_kwargs(torch.bfloat16)
)

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

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions