ModuleNotFoundError: No module named 'torch.distributed.checkpoint.format_utils'
Author: johnmendez2
Pull request: no

### System Info

Transformers version: 4.42.0.dev0
Platform: Linux-6.1.85+-x86_64-with-glibc2.35
Python version: 3.10.12
Huggingface_hub version: 0.23.0
Safetensors version: 0.4.3
Accelerate version: 0.31.0.dev0
PyTorch version (True): 2.1.2+cu121+12.1

### Who can help?

_No response_

### Information

- [ ] The official example scripts
- [X] 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

from transformers import AutoModelForCausalLM, AutoTokenizer
device = "cuda" # the device to load the model onto

model = AutoModelForCausalLM.from_pretrained(
    "/content/drive/MyDrive/LLM_MODELS/CodeQwen1.5-7B-Chat-model",
    torch_dtype="auto",
    device_map="auto"
)
tokenizer = AutoTokenizer.from_pretrained("/content/drive/MyDrive/LLM_MODELS/CodeQwen1.5-7B-Chat-tokenizer")


### Expected behavior

---------------------------------------------------------------------------
ModuleNotFoundError                       Traceback (most recent call last)
[/usr/local/lib/python3.10/dist-packages/transformers/utils/import_utils.py](https://localhost:8080/#) in _get_module(self, module_name)
   1534         try:
-> 1535             return importlib.import_module("." + module_name, self.__name__)
   1536         except Exception as e:

33 frames
[/usr/lib/python3.10/importlib/__init__.py](https://localhost:8080/#) in import_module(name, package)
    125             level += 1
--> 126     return _bootstrap._gcd_import(name[level:], package, level)
    127 

/usr/lib/python3.10/importlib/_bootstrap.py in _gcd_import(name, package, level)

/usr/lib/python3.10/importlib/_bootstrap.py in _find_and_load(name, import_)

/usr/lib/python3.10/importlib/_bootstrap.py in _find_and_load_unlocked(name, import_)

/usr/lib/python3.10/importlib/_bootstrap.py in _load_unlocked(spec)

/usr/lib/python3.10/importlib/_bootstrap_external.py in exec_module(self, module)

/usr/lib/python3.10/importlib/_bootstrap.py in _call_with_frames_removed(f, *args, **kwds)

[/usr/local/lib/python3.10/dist-packages/transformers/generation/utils.py](https://localhost:8080/#) in <module>
     96 if is_accelerate_available():
---> 97     from accelerate.hooks import AlignDevicesHook, add_hook_to_module
     98 

[/usr/local/lib/python3.10/dist-packages/accelerate/__init__.py](https://localhost:8080/#) in <module>
     15 
---> 16 from .accelerator import Accelerator
     17 from .big_modeling import (

[/usr/local/lib/python3.10/dist-packages/accelerate/accelerator.py](https://localhost:8080/#) in <module>
     34 
---> 35 from .checkpointing import load_accelerator_state, load_custom_state, save_accelerator_state, save_custom_state
     36 from .data_loader import DataLoaderDispatcher, prepare_data_loader, skip_first_batches

[/usr/local/lib/python3.10/dist-packages/accelerate/checkpointing.py](https://localhost:8080/#) in <module>
     23 
---> 24 from .utils import (
     25     MODEL_NAME,

[/usr/local/lib/python3.10/dist-packages/accelerate/utils/__init__.py](https://localhost:8080/#) in <module>
    181 from .bnb import has_4bit_bnb_layers, load_and_quantize_model
--> 182 from .fsdp_utils import load_fsdp_model, load_fsdp_optimizer, merge_fsdp_weights, save_fsdp_model, save_fsdp_optimizer
    183 from .launch import (

[/usr/local/lib/python3.10/dist-packages/accelerate/utils/fsdp_utils.py](https://localhost:8080/#) in <module>
     29     import torch.distributed.checkpoint as dist_cp
---> 30     import torch.distributed.checkpoint.format_utils as dist_cp_format_utils
     31     from torch.distributed.checkpoint.default_planner import DefaultLoadPlanner, DefaultSavePlanner

ModuleNotFoundError: No module named 'torch.distributed.checkpoint.format_utils'

The above exception was the direct cause of the following exception:

RuntimeError                              Traceback (most recent call last)
[/usr/local/lib/python3.10/dist-packages/transformers/utils/import_utils.py](https://localhost:8080/#) in _get_module(self, module_name)
   1534         try:
-> 1535             return importlib.import_module("." + module_name, self.__name__)
   1536         except Exception as e:

[/usr/lib/python3.10/importlib/__init__.py](https://localhost:8080/#) in import_module(name, package)
    125             level += 1
--> 126     return _bootstrap._gcd_import(name[level:], package, level)
    127 

/usr/lib/python3.10/importlib/_bootstrap.py in _gcd_import(name, package, level)

/usr/lib/python3.10/importlib/_bootstrap.py in _find_and_load(name, import_)

/usr/lib/python3.10/importlib/_bootstrap.py in _find_and_load_unlocked(name, import_)

/usr/lib/python3.10/importlib/_bootstrap.py in _load_unlocked(spec)

/usr/lib/python3.10/importlib/_bootstrap_external.py in exec_module(self, module)

/usr/lib/python3.10/importlib/_bootstrap.py in _call_with_frames_removed(f, *args, **kwds)

[/usr/local/lib/python3.10/dist-packages/transformers/models/qwen2/modeling_qwen2.py](https://localhost:8080/#) in <module>
     39 )
---> 40 from ...modeling_utils import PreTrainedModel
     41 from ...utils import (

[/usr/local/lib/python3.10/dist-packages/transformers/modeling_utils.py](https://localhost:8080/#) in <module>
     44 from .dynamic_module_utils import custom_object_save
---> 45 from .generation import GenerationConfig, GenerationMixin
     46 from .integrations import PeftAdapterMixin, deepspeed_config, is_deepspeed_zero3_enabled

/usr/lib/python3.10/importlib/_bootstrap.py in _handle_fromlist(module, fromlist, import_, recursive)

[/usr/local/lib/python3.10/dist-packages/transformers/utils/import_utils.py](https://localhost:8080/#) in __getattr__(self, name)
   1524         elif name in self._class_to_module.keys():
-> 1525             module = self._get_module(self._class_to_module[name])
   1526             value = getattr(module, name)

[/usr/local/lib/python3.10/dist-packages/transformers/utils/import_utils.py](https://localhost:8080/#) in _get_module(self, module_name)
   1536         except Exception as e:
-> 1537             raise RuntimeError(
   1538                 f"Failed to import {self.__name__}.{module_name} because of the following error (look up to see its"

RuntimeError: Failed to import transformers.generation.utils because of the following error (look up to see its traceback):
No module named 'torch.distributed.checkpoint.format_utils'

The above exception was the direct cause of the following exception:

RuntimeError                              Traceback (most recent call last)
[<ipython-input-4-0fb942283b15>](https://localhost:8080/#) in <cell line: 4>()
      2 device = "cuda" # the device to load the model onto
      3 
----> 4 model = AutoModelForCausalLM.from_pretrained(
      5     "/content/drive/MyDrive/LLM_MODELS/CodeQwen1.5-7B-Chat-model",
      6     torch_dtype="auto",

[/usr/local/lib/python3.10/dist-packages/transformers/models/auto/auto_factory.py](https://localhost:8080/#) in from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs)
    560             )
    561         elif type(config) in cls._model_mapping.keys():
--> 562             model_class = _get_model_class(config, cls._model_mapping)
    563             return model_class.from_pretrained(
    564                 pretrained_model_name_or_path, *model_args, config=config, **hub_kwargs, **kwargs

[/usr/local/lib/python3.10/dist-packages/transformers/models/auto/auto_factory.py](https://localhost:8080/#) in _get_model_class(config, model_mapping)
    381 
    382 def _get_model_class(config, model_mapping):
--> 383     supported_models = model_mapping[type(config)]
    384     if not isinstance(supported_models, (list, tuple)):
    385         return supported_models

[/usr/local/lib/python3.10/dist-packages/transformers/models/auto/auto_factory.py](https://localhost:8080/#) in __getitem__(self, key)
    732         if model_type in self._model_mapping:
    733             model_name = self._model_mapping[model_type]
--> 734             return self._load_attr_from_module(model_type, model_name)
    735 
    736         # Maybe there was several model types associated with this config.

[/usr/local/lib/python3.10/dist-packages/transformers/models/auto/auto_factory.py](https://localhost:8080/#) in _load_attr_from_module(self, model_type, attr)
    746         if module_name not in self._modules:
    747             self._modules[module_name] = importlib.import_module(f".{module_name}", "transformers.models")
--> 748         return getattribute_from_module(self._modules[module_name], attr)
    749 
    750     def keys(self):

[/usr/local/lib/python3.10/dist-packages/transformers/models/auto/auto_factory.py](https://localhost:8080/#) in getattribute_from_module(module, attr)
    690     if isinstance(attr, tuple):
    691         return tuple(getattribute_from_module(module, a) for a in attr)
--> 692     if hasattr(module, attr):
    693         return getattr(module, attr)
    694     # Some of the mappings have entries model_type -> object of another model type. In that case we try to grab the

[/usr/local/lib/python3.10/dist-packages/transformers/utils/import_utils.py](https://localhost:8080/#) in __getattr__(self, name)
   1523             value = self._get_module(name)
   1524         elif name in self._class_to_module.keys():
-> 1525             module = self._get_module(self._class_to_module[name])
   1526             value = getattr(module, name)
   1527         else:

[/usr/local/lib/python3.10/dist-packages/transformers/utils/import_utils.py](https://localhost:8080/#) in _get_module(self, module_name)
   1535             return importlib.import_module("." + module_name, self.__name__)
   1536         except Exception as e:
-> 1537             raise RuntimeError(
   1538                 f"Failed to import {self.__name__}.{module_name} because of the following error (look up to see its"
   1539                 f" traceback):\n{e}"

RuntimeError: Failed to import transformers.models.qwen2.modeling_qwen2 because of the following error (look up to see its traceback):
Failed to import transformers.generation.utils because of the following error (look up to see its traceback):
No module named 'torch.distributed.checkpoint.format_utils'

Comments:
cc @muellerzr @pacman100 As this looks like it's coming from accelerate

Hi @johnmendez2 , can you please try upgrading your pytorch version by `pip install --upgrade torch torchvision` in your environment and see if the error still persists

> Hi @johnmendez2 , can you please try upgrading your pytorch version by `pip install --upgrade torch torchvision` in your environment and see if the error still persists

This worked thank you :)

Cheers!