autocast() got an unexpected keyword argument 'cache_enabled when use trainer.torch_jit_model_eval

[Wanguy]

System Info

• transformers version: 4.46.3
• Platform: Linux-4.18.0-147.el8_1.x86_64-x86_64-with-glibc2.10
• Python version: 3.8.10
• Huggingface_hub version: 0.26.5
• Safetensors version: 0.4.5
• Accelerate version: 1.0.1
• Accelerate config: not found
• PyTorch version (GPU?): 2.4.1+cu118 (True)
• Tensorflow version (GPU?): not installed (NA)
• Flax version (CPU?/GPU?/TPU?): not installed (NA)
• Jax version: not installed
• JaxLib version: not installed
• Using distributed or parallel set-up in script?:
• Using GPU in script?:
• GPU type: NVIDIA A100-SXM4-80GB

Who can help?

@muellerzr
@SunMarc

Information

• The official example scripts
• My own modified scripts

Tasks

• An officially supported task in the examples folder (such as GLUE/SQuAD, ...)
• My own task or dataset (give details below)

Reproduction

When using the torch_jit_model_eval() method in trainer, it prompts

"failed to use PyTorch jit mode due to: autocast() got an unexpected keyword argument 'cache_enabled'."

Looking at the details, it was found that the error was caused by the self.accelerator.autocast(cache_enabled=False) method. Its method definition is def autocast(self, autocast_handler: AutocastKwargs = None), and there is no cache_enabled method.

Is this because the code here has not been updated, or because I ignored some settings?

Is there a solution now?

Expected behavior

Work normally.

[Wanguy]

The last accelerator.autocast method that supports the cache_enabled parameter is in version v0.34.2.

[SunMarc]

Thanks for the report @Wanguy ! This is indeed deprecated as you can see here. Would you like to submit a PR to fix that ? Thanks !

[Wanguy]

@SunMarc Of course, I will submit a fix PR ASAP to fix this bug. :)