We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
[tensor] fix some unittests

[tensor] fix some unittests

[tensor] fix some unittests


--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `colossalai/lazy/pretrained.py`
Content:
```
1 import os
2 from typing import Callable, Optional, Union
3 
4 import torch
5 from torch.nn import Module
6 
7 from colossalai.interface import pretrained as pretrained_interface
8 
9 
10 class PretrainedManager:
11     old_from_pretrained: Optional[Callable] = None
12 
13     @staticmethod
14     def inject() -> None:
15         try:
16             from transformers.modeling_utils import PreTrainedModel
17         except ImportError:
18             return
19         # recover bound method to plain function
20         PretrainedManager.old_from_pretrained = PreTrainedModel.from_pretrained.__func__
21         PreTrainedModel.from_pretrained = new_from_pretrained
22 
23     @staticmethod
24     def recover() -> None:
25         try:
26             from transformers.modeling_utils import PreTrainedModel
27         except ImportError:
28             return
29         # convert plain function to class method
30         PreTrainedModel.from_pretrained = classmethod(PretrainedManager.old_from_pretrained)
31         PretrainedManager.old_from_pretrained = None
32 
33 
34 @classmethod
35 def new_from_pretrained(
36     cls, pretrained_model_name_or_path: Optional[Union[str, os.PathLike]], *model_args, **kwargs
37 ) -> Module:
38     from transformers import GenerationConfig
39     from transformers.configuration_utils import PretrainedConfig
40     from transformers.modeling_utils import (
41         ContextManagers,
42         _add_variant,
43         cached_file,
44         download_url,
45         has_file,
46         is_offline_mode,
47         is_remote_url,
48         no_init_weights,
49     )
50     from transformers.utils import (
51         SAFE_WEIGHTS_INDEX_NAME,
52         SAFE_WEIGHTS_NAME,
53         WEIGHTS_INDEX_NAME,
54         WEIGHTS_NAME,
55         is_safetensors_available,
56         logging,
57     )
58 
59     logger = logging.get_logger(__name__)
60 
61     config = kwargs.pop("config", None)
62     cache_dir = kwargs.pop("cache_dir", None)
63     force_download = kwargs.pop("force_download", False)
64     resume_download = kwargs.pop("resume_download", False)
65     proxies = kwargs.pop("proxies", None)
66     local_files_only = kwargs.pop("local_files_only", False)
67     use_auth_token = kwargs.pop("use_auth_token", None)
68     revision = kwargs.pop("revision", None)
69     _ = kwargs.pop("mirror", None)
70     from_pipeline = kwargs.pop("_from_pipeline", None)
71     from_auto_class = kwargs.pop("_from_auto", False)
72     _fast_init = kwargs.pop("_fast_init", True)
73     torch_dtype = kwargs.pop("torch_dtype", None)
74     subfolder = kwargs.pop("subfolder", "")
75     commit_hash = kwargs.pop("_commit_hash", None)
76     variant = kwargs.pop("variant", None)
77     use_safetensors = kwargs.pop("use_safetensors", None if is_safetensors_available() else False)
78 
79     if len(kwargs) > 0:
80         logger.warning(f"Below kwargs may be ignored: {list(kwargs.keys())}")
81 
82     from_pt = True
83 
84     user_agent = {"file_type": "model", "framework": "pytorch", "from_auto_class": from_auto_class}
85     if from_pipeline is not None:
86         user_agent["using_pipeline"] = from_pipeline
87 
88     if is_offline_mode() and not local_files_only:
89         logger.info("Offline mode: forcing local_files_only=True")
90         local_files_only = True
91 
92     # Load config if we don't provide a configuration
93     if not isinstance(config, PretrainedConfig):
94         config_path = config if config is not None else pretrained_model_name_or_path
95         config, model_kwargs = cls.config_class.from_pretrained(
96             config_path,
97             cache_dir=cache_dir,
98             return_unused_kwargs=True,
99             force_download=force_download,
100             resume_download=resume_download,
101             proxies=proxies,
102             local_files_only=local_files_only,
103             use_auth_token=use_auth_token,
104             revision=revision,
105             subfolder=subfolder,
106             _from_auto=from_auto_class,
107             _from_pipeline=from_pipeline,
108             **kwargs,
109         )
110     else:
111         model_kwargs = kwargs
112 
113     if commit_hash is None:
114         commit_hash = getattr(config, "_commit_hash", None)
115 
116     # This variable will flag if we're loading a sharded checkpoint. In this case the archive file is just the
117     # index of the files.
118 
119     if pretrained_model_name_or_path is not None:
120         pretrained_model_name_or_path = str(pretrained_model_name_or_path)
121         is_local = os.path.isdir(pretrained_model_name_or_path)
122         if is_local:
123             if use_safetensors is not False and os.path.isfile(
124                 os.path.join(pretrained_model_name_or_path, subfolder, _add_variant(SAFE_WEIGHTS_NAME, variant))
125             ):
126                 # Load from a safetensors checkpoint
127                 archive_file = os.path.join(
128                     pretrained_model_name_or_path, subfolder, _add_variant(SAFE_WEIGHTS_NAME, variant)
129                 )
130             elif use_safetensors is not False and os.path.isfile(
131                 os.path.join(pretrained_model_name_or_path, subfolder, _add_variant(SAFE_WEIGHTS_INDEX_NAME, variant))
132             ):
133                 # Load from a sharded safetensors checkpoint
134                 archive_file = os.path.join(
135                     pretrained_model_name_or_path, subfolder, _add_variant(SAFE_WEIGHTS_INDEX_NAME, variant)
136                 )
137             elif os.path.isfile(
138                 os.path.join(pretrained_model_name_or_path, subfolder, _add_variant(WEIGHTS_NAME, variant))
139             ):
140                 # Load from a PyTorch checkpoint
141                 archive_file = os.path.join(
142                     pretrained_model_name_or_path, subfolder, _add_variant(WEIGHTS_NAME, variant)
143                 )
144             elif os.path.isfile(
145                 os.path.join(pretrained_model_name_or_path, subfolder, _add_variant(WEIGHTS_INDEX_NAME, variant))
146             ):
147                 # Load from a sharded PyTorch checkpoint
148                 archive_file = os.path.join(
149                     pretrained_model_name_or_path, subfolder, _add_variant(WEIGHTS_INDEX_NAME, variant)
150                 )
151             else:
152                 raise EnvironmentError(
153                     f"Error no file named {_add_variant(WEIGHTS_NAME, variant)} found in directory"
154                     f" {pretrained_model_name_or_path}."
155                 )
156         elif os.path.isfile(os.path.join(subfolder, pretrained_model_name_or_path)):
157             archive_file = pretrained_model_name_or_path
158             is_local = True
159         elif is_remote_url(pretrained_model_name_or_path):
160             filename = pretrained_model_name_or_path
161             resolved_archive_file = download_url(pretrained_model_name_or_path)
162         else:
163             # set correct filename
164             if use_safetensors is not False:
165                 filename = _add_variant(SAFE_WEIGHTS_NAME, variant)
166             else:
167                 filename = _add_variant(WEIGHTS_NAME, variant)
168 
169             try:
170                 # Load from URL or cache if already cached
171                 cached_file_kwargs = {
172                     "cache_dir": cache_dir,
173                     "force_download": force_download,
174                     "proxies": proxies,
175                     "resume_download": resume_download,
176                     "local_files_only": local_files_only,
177                     "use_auth_token": use_auth_token,
178                     "user_agent": user_agent,
179                     "revision": revision,
180                     "subfolder": subfolder,
181                     "_raise_exceptions_for_missing_entries": False,
182                     "_commit_hash": commit_hash,
183                 }
184                 resolved_archive_file = cached_file(pretrained_model_name_or_path, filename, **cached_file_kwargs)
185 
186                 # Since we set _raise_exceptions_for_missing_entries=False, we don't get an exception but a None
187                 # result when internet is up, the repo and revision exist, but the file does not.
188                 if resolved_archive_file is None and filename == _add_variant(SAFE_WEIGHTS_NAME, variant):
189                     # Maybe the checkpoint is sharded, we try to grab the index name in this case.
190                     resolved_archive_file = cached_file(
191                         pretrained_model_name_or_path,
192                         _add_variant(SAFE_WEIGHTS_INDEX_NAME, variant),
193                         **cached_file_kwargs,
194                     )
195                     if resolved_archive_file is not None:
196                         pass
197                     elif use_safetensors:
198                         raise EnvironmentError(
199                             f" {_add_variant(SAFE_WEIGHTS_NAME, variant)} or {_add_variant(SAFE_WEIGHTS_INDEX_NAME, variant)} and thus cannot be loaded with `safetensors`. Please make sure that the model has been saved with `safe_serialization=True` or do not set `use_safetensors=True`."
200                         )
201                     else:
202                         # This repo has no safetensors file of any kind, we switch to PyTorch.
203                         filename = _add_variant(WEIGHTS_NAME, variant)
204                         resolved_archive_file = cached_file(
205                             pretrained_model_name_or_path, filename, **cached_file_kwargs
206                         )
207                 if resolved_archive_file is None and filename == _add_variant(WEIGHTS_NAME, variant):
208                     # Maybe the checkpoint is sharded, we try to grab the index name in this case.
209                     resolved_archive_file = cached_file(
210                         pretrained_model_name_or_path,
211                         _add_variant(WEIGHTS_INDEX_NAME, variant),
212                         **cached_file_kwargs,
213                     )
214                     if resolved_archive_file is not None:
215                         pass
216                 if resolved_archive_file is None:
217                     # Otherwise, maybe there is a TF or Flax model file.  We try those to give a helpful error
218                     # message.
219                     has_file_kwargs = {
220                         "revision": revision,
221                         "proxies": proxies,
222                         "use_auth_token": use_auth_token,
223                     }
224                     if variant is not None and has_file(pretrained_model_name_or_path, WEIGHTS_NAME, **has_file_kwargs):
225                         raise EnvironmentError(
226                             f"{pretrained_model_name_or_path} does not appear to have a file named"
227                             f" {_add_variant(WEIGHTS_NAME, variant)} but there is a file without the variant"
228                             f" {variant}. Use `variant=None` to load this model from those weights."
229                         )
230                     else:
231                         raise EnvironmentError(
232                             f"{pretrained_model_name_or_path} does not appear to have a file named"
233                             f" {_add_variant(WEIGHTS_NAME, variant)}"
234                         )
235             except EnvironmentError:
236                 # Raise any environment error raise by `cached_file`. It will have a helpful error message adapted
237                 # to the original exception.
238                 raise
239             except Exception:
240                 # For any other exception, we throw a generic error.
241                 raise EnvironmentError(
242                     f"Can't load the model for '{pretrained_model_name_or_path}'. If you were trying to load it"
243                     " from 'https://huggingface.co/models', make sure you don't have a local directory with the"
244                     f" same name. Otherwise, make sure '{pretrained_model_name_or_path}' is the correct path to a"
245                     f" directory containing a file named {_add_variant(WEIGHTS_NAME, variant)}."
246                 )
247 
248         if is_local:
249             logger.info(f"loading weights file {archive_file}")
250             resolved_archive_file = archive_file
251         else:
252             logger.info(f"loading weights file {filename} from cache at {resolved_archive_file}")
253     else:
254         resolved_archive_file = None
255 
256     if from_pt:
257         # set dtype to instantiate the model under:
258         # 1. If torch_dtype is not None, we use that dtype
259         dtype_orig = None
260 
261         if torch_dtype is not None:
262             if not isinstance(torch_dtype, torch.dtype):
263                 raise ValueError(f"`torch_dtype` can be either `torch.dtype` or `None`, but received {torch_dtype}")
264             dtype_orig = cls._set_default_torch_dtype(torch_dtype)
265 
266     config.name_or_path = pretrained_model_name_or_path
267 
268     # Instantiate model.
269     init_contexts = [no_init_weights(_enable=_fast_init)]
270 
271     with ContextManagers(init_contexts):
272         model = cls(config, *model_args, **model_kwargs)
273 
274     if from_pt:
275         # restore default dtype
276         if dtype_orig is not None:
277             torch.set_default_dtype(dtype_orig)
278 
279     # make sure token embedding weights are still tied if needed
280     model.tie_weights()
281 
282     # Set model in evaluation mode to deactivate DropOut modules by default
283     model.eval()
284 
285     # If it is a model with generation capabilities, attempt to load the generation config
286     if model.can_generate():
287         try:
288             model.generation_config = GenerationConfig.from_pretrained(
289                 pretrained_model_name_or_path,
290                 cache_dir=cache_dir,
291                 force_download=force_download,
292                 resume_download=resume_download,
293                 proxies=proxies,
294                 local_files_only=local_files_only,
295                 use_auth_token=use_auth_token,
296                 revision=revision,
297                 subfolder=subfolder,
298                 _from_auto=from_auto_class,
299                 _from_pipeline=from_pipeline,
300                 **kwargs,
301             )
302         except (OSError, TypeError):
303             logger.info("Generation config file not found, using a generation config created from the model config.")
304 
305     # set pretrained path
306     if resolved_archive_file:
307         pretrained_interface.set_pretrained_path(model, resolved_archive_file)
308 
309     return model
310
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/colossalai/lazy/pretrained.py b/colossalai/lazy/pretrained.py
--- a/colossalai/lazy/pretrained.py
+++ b/colossalai/lazy/pretrained.py
@@ -1,3 +1,4 @@
+import copy
 import os
 from typing import Callable, Optional, Union
 
@@ -74,6 +75,24 @@
     subfolder = kwargs.pop("subfolder", "")
     commit_hash = kwargs.pop("_commit_hash", None)
     variant = kwargs.pop("variant", None)
+
+    kwargs.pop("state_dict", None)
+    kwargs.pop("from_tf", False)
+    kwargs.pop("from_flax", False)
+    kwargs.pop("output_loading_info", False)
+    kwargs.pop("trust_remote_code", None)
+    kwargs.pop("low_cpu_mem_usage", None)
+    kwargs.pop("device_map", None)
+    kwargs.pop("max_memory", None)
+    kwargs.pop("offload_folder", None)
+    kwargs.pop("offload_state_dict", False)
+    kwargs.pop("load_in_8bit", False)
+    kwargs.pop("load_in_4bit", False)
+    kwargs.pop("quantization_config", None)
+    kwargs.pop("adapter_kwargs", {})
+    kwargs.pop("adapter_name", "default")
+    kwargs.pop("use_flash_attention_2", False)
+
     use_safetensors = kwargs.pop("use_safetensors", None if is_safetensors_available() else False)
 
     if len(kwargs) > 0:
@@ -108,6 +127,10 @@
             **kwargs,
         )
     else:
+        config = copy.deepcopy(config)
+        kwarg_attn_imp = kwargs.pop("attn_implementation", None)
+        if kwarg_attn_imp is not None and config._attn_implementation != kwarg_attn_imp:
+            config._attn_implementation = kwarg_attn_imp
         model_kwargs = kwargs
 
     if commit_hash is None:

{"golden_diff": "diff --git a/colossalai/lazy/pretrained.py b/colossalai/lazy/pretrained.py\n--- a/colossalai/lazy/pretrained.py\n+++ b/colossalai/lazy/pretrained.py\n@@ -1,3 +1,4 @@\n+import copy\n import os\n from typing import Callable, Optional, Union\n \n@@ -74,6 +75,24 @@\n     subfolder = kwargs.pop(\"subfolder\", \"\")\n     commit_hash = kwargs.pop(\"_commit_hash\", None)\n     variant = kwargs.pop(\"variant\", None)\n+\n+    kwargs.pop(\"state_dict\", None)\n+    kwargs.pop(\"from_tf\", False)\n+    kwargs.pop(\"from_flax\", False)\n+    kwargs.pop(\"output_loading_info\", False)\n+    kwargs.pop(\"trust_remote_code\", None)\n+    kwargs.pop(\"low_cpu_mem_usage\", None)\n+    kwargs.pop(\"device_map\", None)\n+    kwargs.pop(\"max_memory\", None)\n+    kwargs.pop(\"offload_folder\", None)\n+    kwargs.pop(\"offload_state_dict\", False)\n+    kwargs.pop(\"load_in_8bit\", False)\n+    kwargs.pop(\"load_in_4bit\", False)\n+    kwargs.pop(\"quantization_config\", None)\n+    kwargs.pop(\"adapter_kwargs\", {})\n+    kwargs.pop(\"adapter_name\", \"default\")\n+    kwargs.pop(\"use_flash_attention_2\", False)\n+\n     use_safetensors = kwargs.pop(\"use_safetensors\", None if is_safetensors_available() else False)\n \n     if len(kwargs) > 0:\n@@ -108,6 +127,10 @@\n             **kwargs,\n         )\n     else:\n+        config = copy.deepcopy(config)\n+        kwarg_attn_imp = kwargs.pop(\"attn_implementation\", None)\n+        if kwarg_attn_imp is not None and config._attn_implementation != kwarg_attn_imp:\n+            config._attn_implementation = kwarg_attn_imp\n         model_kwargs = kwargs\n \n     if commit_hash is None:\n", "issue": "[tensor] fix some unittests\n\n[tensor] fix some unittests\n\n[tensor] fix some unittests\n\n", "before_files": [{"content": "import os\nfrom typing import Callable, Optional, Union\n\nimport torch\nfrom torch.nn import Module\n\nfrom colossalai.interface import pretrained as pretrained_interface\n\n\nclass PretrainedManager:\n    old_from_pretrained: Optional[Callable] = None\n\n    @staticmethod\n    def inject() -> None:\n        try:\n            from transformers.modeling_utils import PreTrainedModel\n        except ImportError:\n            return\n        # recover bound method to plain function\n        PretrainedManager.old_from_pretrained = PreTrainedModel.from_pretrained.__func__\n        PreTrainedModel.from_pretrained = new_from_pretrained\n\n    @staticmethod\n    def recover() -> None:\n        try:\n            from transformers.modeling_utils import PreTrainedModel\n        except ImportError:\n            return\n        # convert plain function to class method\n        PreTrainedModel.from_pretrained = classmethod(PretrainedManager.old_from_pretrained)\n        PretrainedManager.old_from_pretrained = None\n\n\n@classmethod\ndef new_from_pretrained(\n    cls, pretrained_model_name_or_path: Optional[Union[str, os.PathLike]], *model_args, **kwargs\n) -> Module:\n    from transformers import GenerationConfig\n    from transformers.configuration_utils import PretrainedConfig\n    from transformers.modeling_utils import (\n        ContextManagers,\n        _add_variant,\n        cached_file,\n        download_url,\n        has_file,\n        is_offline_mode,\n        is_remote_url,\n        no_init_weights,\n    )\n    from transformers.utils import (\n        SAFE_WEIGHTS_INDEX_NAME,\n        SAFE_WEIGHTS_NAME,\n        WEIGHTS_INDEX_NAME,\n        WEIGHTS_NAME,\n        is_safetensors_available,\n        logging,\n    )\n\n    logger = logging.get_logger(__name__)\n\n    config = kwargs.pop(\"config\", None)\n    cache_dir = kwargs.pop(\"cache_dir\", None)\n    force_download = kwargs.pop(\"force_download\", False)\n    resume_download = kwargs.pop(\"resume_download\", False)\n    proxies = kwargs.pop(\"proxies\", None)\n    local_files_only = kwargs.pop(\"local_files_only\", False)\n    use_auth_token = kwargs.pop(\"use_auth_token\", None)\n    revision = kwargs.pop(\"revision\", None)\n    _ = kwargs.pop(\"mirror\", None)\n    from_pipeline = kwargs.pop(\"_from_pipeline\", None)\n    from_auto_class = kwargs.pop(\"_from_auto\", False)\n    _fast_init = kwargs.pop(\"_fast_init\", True)\n    torch_dtype = kwargs.pop(\"torch_dtype\", None)\n    subfolder = kwargs.pop(\"subfolder\", \"\")\n    commit_hash = kwargs.pop(\"_commit_hash\", None)\n    variant = kwargs.pop(\"variant\", None)\n    use_safetensors = kwargs.pop(\"use_safetensors\", None if is_safetensors_available() else False)\n\n    if len(kwargs) > 0:\n        logger.warning(f\"Below kwargs may be ignored: {list(kwargs.keys())}\")\n\n    from_pt = True\n\n    user_agent = {\"file_type\": \"model\", \"framework\": \"pytorch\", \"from_auto_class\": from_auto_class}\n    if from_pipeline is not None:\n        user_agent[\"using_pipeline\"] = from_pipeline\n\n    if is_offline_mode() and not local_files_only:\n        logger.info(\"Offline mode: forcing local_files_only=True\")\n        local_files_only = True\n\n    # Load config if we don't provide a configuration\n    if not isinstance(config, PretrainedConfig):\n        config_path = config if config is not None else pretrained_model_name_or_path\n        config, model_kwargs = cls.config_class.from_pretrained(\n            config_path,\n            cache_dir=cache_dir,\n            return_unused_kwargs=True,\n            force_download=force_download,\n            resume_download=resume_download,\n            proxies=proxies,\n            local_files_only=local_files_only,\n            use_auth_token=use_auth_token,\n            revision=revision,\n            subfolder=subfolder,\n            _from_auto=from_auto_class,\n            _from_pipeline=from_pipeline,\n            **kwargs,\n        )\n    else:\n        model_kwargs = kwargs\n\n    if commit_hash is None:\n        commit_hash = getattr(config, \"_commit_hash\", None)\n\n    # This variable will flag if we're loading a sharded checkpoint. In this case the archive file is just the\n    # index of the files.\n\n    if pretrained_model_name_or_path is not None:\n        pretrained_model_name_or_path = str(pretrained_model_name_or_path)\n        is_local = os.path.isdir(pretrained_model_name_or_path)\n        if is_local:\n            if use_safetensors is not False and os.path.isfile(\n                os.path.join(pretrained_model_name_or_path, subfolder, _add_variant(SAFE_WEIGHTS_NAME, variant))\n            ):\n                # Load from a safetensors checkpoint\n                archive_file = os.path.join(\n                    pretrained_model_name_or_path, subfolder, _add_variant(SAFE_WEIGHTS_NAME, variant)\n                )\n            elif use_safetensors is not False and os.path.isfile(\n                os.path.join(pretrained_model_name_or_path, subfolder, _add_variant(SAFE_WEIGHTS_INDEX_NAME, variant))\n            ):\n                # Load from a sharded safetensors checkpoint\n                archive_file = os.path.join(\n                    pretrained_model_name_or_path, subfolder, _add_variant(SAFE_WEIGHTS_INDEX_NAME, variant)\n                )\n            elif os.path.isfile(\n                os.path.join(pretrained_model_name_or_path, subfolder, _add_variant(WEIGHTS_NAME, variant))\n            ):\n                # Load from a PyTorch checkpoint\n                archive_file = os.path.join(\n                    pretrained_model_name_or_path, subfolder, _add_variant(WEIGHTS_NAME, variant)\n                )\n            elif os.path.isfile(\n                os.path.join(pretrained_model_name_or_path, subfolder, _add_variant(WEIGHTS_INDEX_NAME, variant))\n            ):\n                # Load from a sharded PyTorch checkpoint\n                archive_file = os.path.join(\n                    pretrained_model_name_or_path, subfolder, _add_variant(WEIGHTS_INDEX_NAME, variant)\n                )\n            else:\n                raise EnvironmentError(\n                    f\"Error no file named {_add_variant(WEIGHTS_NAME, variant)} found in directory\"\n                    f\" {pretrained_model_name_or_path}.\"\n                )\n        elif os.path.isfile(os.path.join(subfolder, pretrained_model_name_or_path)):\n            archive_file = pretrained_model_name_or_path\n            is_local = True\n        elif is_remote_url(pretrained_model_name_or_path):\n            filename = pretrained_model_name_or_path\n            resolved_archive_file = download_url(pretrained_model_name_or_path)\n        else:\n            # set correct filename\n            if use_safetensors is not False:\n                filename = _add_variant(SAFE_WEIGHTS_NAME, variant)\n            else:\n                filename = _add_variant(WEIGHTS_NAME, variant)\n\n            try:\n                # Load from URL or cache if already cached\n                cached_file_kwargs = {\n                    \"cache_dir\": cache_dir,\n                    \"force_download\": force_download,\n                    \"proxies\": proxies,\n                    \"resume_download\": resume_download,\n                    \"local_files_only\": local_files_only,\n                    \"use_auth_token\": use_auth_token,\n                    \"user_agent\": user_agent,\n                    \"revision\": revision,\n                    \"subfolder\": subfolder,\n                    \"_raise_exceptions_for_missing_entries\": False,\n                    \"_commit_hash\": commit_hash,\n                }\n                resolved_archive_file = cached_file(pretrained_model_name_or_path, filename, **cached_file_kwargs)\n\n                # Since we set _raise_exceptions_for_missing_entries=False, we don't get an exception but a None\n                # result when internet is up, the repo and revision exist, but the file does not.\n                if resolved_archive_file is None and filename == _add_variant(SAFE_WEIGHTS_NAME, variant):\n                    # Maybe the checkpoint is sharded, we try to grab the index name in this case.\n                    resolved_archive_file = cached_file(\n                        pretrained_model_name_or_path,\n                        _add_variant(SAFE_WEIGHTS_INDEX_NAME, variant),\n                        **cached_file_kwargs,\n                    )\n                    if resolved_archive_file is not None:\n                        pass\n                    elif use_safetensors:\n                        raise EnvironmentError(\n                            f\" {_add_variant(SAFE_WEIGHTS_NAME, variant)} or {_add_variant(SAFE_WEIGHTS_INDEX_NAME, variant)} and thus cannot be loaded with `safetensors`. Please make sure that the model has been saved with `safe_serialization=True` or do not set `use_safetensors=True`.\"\n                        )\n                    else:\n                        # This repo has no safetensors file of any kind, we switch to PyTorch.\n                        filename = _add_variant(WEIGHTS_NAME, variant)\n                        resolved_archive_file = cached_file(\n                            pretrained_model_name_or_path, filename, **cached_file_kwargs\n                        )\n                if resolved_archive_file is None and filename == _add_variant(WEIGHTS_NAME, variant):\n                    # Maybe the checkpoint is sharded, we try to grab the index name in this case.\n                    resolved_archive_file = cached_file(\n                        pretrained_model_name_or_path,\n                        _add_variant(WEIGHTS_INDEX_NAME, variant),\n                        **cached_file_kwargs,\n                    )\n                    if resolved_archive_file is not None:\n                        pass\n                if resolved_archive_file is None:\n                    # Otherwise, maybe there is a TF or Flax model file.  We try those to give a helpful error\n                    # message.\n                    has_file_kwargs = {\n                        \"revision\": revision,\n                        \"proxies\": proxies,\n                        \"use_auth_token\": use_auth_token,\n                    }\n                    if variant is not None and has_file(pretrained_model_name_or_path, WEIGHTS_NAME, **has_file_kwargs):\n                        raise EnvironmentError(\n                            f\"{pretrained_model_name_or_path} does not appear to have a file named\"\n                            f\" {_add_variant(WEIGHTS_NAME, variant)} but there is a file without the variant\"\n                            f\" {variant}. Use `variant=None` to load this model from those weights.\"\n                        )\n                    else:\n                        raise EnvironmentError(\n                            f\"{pretrained_model_name_or_path} does not appear to have a file named\"\n                            f\" {_add_variant(WEIGHTS_NAME, variant)}\"\n                        )\n            except EnvironmentError:\n                # Raise any environment error raise by `cached_file`. It will have a helpful error message adapted\n                # to the original exception.\n                raise\n            except Exception:\n                # For any other exception, we throw a generic error.\n                raise EnvironmentError(\n                    f\"Can't load the model for '{pretrained_model_name_or_path}'. If you were trying to load it\"\n                    \" from 'https://huggingface.co/models', make sure you don't have a local directory with the\"\n                    f\" same name. Otherwise, make sure '{pretrained_model_name_or_path}' is the correct path to a\"\n                    f\" directory containing a file named {_add_variant(WEIGHTS_NAME, variant)}.\"\n                )\n\n        if is_local:\n            logger.info(f\"loading weights file {archive_file}\")\n            resolved_archive_file = archive_file\n        else:\n            logger.info(f\"loading weights file {filename} from cache at {resolved_archive_file}\")\n    else:\n        resolved_archive_file = None\n\n    if from_pt:\n        # set dtype to instantiate the model under:\n        # 1. If torch_dtype is not None, we use that dtype\n        dtype_orig = None\n\n        if torch_dtype is not None:\n            if not isinstance(torch_dtype, torch.dtype):\n                raise ValueError(f\"`torch_dtype` can be either `torch.dtype` or `None`, but received {torch_dtype}\")\n            dtype_orig = cls._set_default_torch_dtype(torch_dtype)\n\n    config.name_or_path = pretrained_model_name_or_path\n\n    # Instantiate model.\n    init_contexts = [no_init_weights(_enable=_fast_init)]\n\n    with ContextManagers(init_contexts):\n        model = cls(config, *model_args, **model_kwargs)\n\n    if from_pt:\n        # restore default dtype\n        if dtype_orig is not None:\n            torch.set_default_dtype(dtype_orig)\n\n    # make sure token embedding weights are still tied if needed\n    model.tie_weights()\n\n    # Set model in evaluation mode to deactivate DropOut modules by default\n    model.eval()\n\n    # If it is a model with generation capabilities, attempt to load the generation config\n    if model.can_generate():\n        try:\n            model.generation_config = GenerationConfig.from_pretrained(\n                pretrained_model_name_or_path,\n                cache_dir=cache_dir,\n                force_download=force_download,\n                resume_download=resume_download,\n                proxies=proxies,\n                local_files_only=local_files_only,\n                use_auth_token=use_auth_token,\n                revision=revision,\n                subfolder=subfolder,\n                _from_auto=from_auto_class,\n                _from_pipeline=from_pipeline,\n                **kwargs,\n            )\n        except (OSError, TypeError):\n            logger.info(\"Generation config file not found, using a generation config created from the model config.\")\n\n    # set pretrained path\n    if resolved_archive_file:\n        pretrained_interface.set_pretrained_path(model, resolved_archive_file)\n\n    return model\n", "path": "colossalai/lazy/pretrained.py"}], "after_files": [{"content": "import copy\nimport os\nfrom typing import Callable, Optional, Union\n\nimport torch\nfrom torch.nn import Module\n\nfrom colossalai.interface import pretrained as pretrained_interface\n\n\nclass PretrainedManager:\n    old_from_pretrained: Optional[Callable] = None\n\n    @staticmethod\n    def inject() -> None:\n        try:\n            from transformers.modeling_utils import PreTrainedModel\n        except ImportError:\n            return\n        # recover bound method to plain function\n        PretrainedManager.old_from_pretrained = PreTrainedModel.from_pretrained.__func__\n        PreTrainedModel.from_pretrained = new_from_pretrained\n\n    @staticmethod\n    def recover() -> None:\n        try:\n            from transformers.modeling_utils import PreTrainedModel\n        except ImportError:\n            return\n        # convert plain function to class method\n        PreTrainedModel.from_pretrained = classmethod(PretrainedManager.old_from_pretrained)\n        PretrainedManager.old_from_pretrained = None\n\n\n@classmethod\ndef new_from_pretrained(\n    cls, pretrained_model_name_or_path: Optional[Union[str, os.PathLike]], *model_args, **kwargs\n) -> Module:\n    from transformers import GenerationConfig\n    from transformers.configuration_utils import PretrainedConfig\n    from transformers.modeling_utils import (\n        ContextManagers,\n        _add_variant,\n        cached_file,\n        download_url,\n        has_file,\n        is_offline_mode,\n        is_remote_url,\n        no_init_weights,\n    )\n    from transformers.utils import (\n        SAFE_WEIGHTS_INDEX_NAME,\n        SAFE_WEIGHTS_NAME,\n        WEIGHTS_INDEX_NAME,\n        WEIGHTS_NAME,\n        is_safetensors_available,\n        logging,\n    )\n\n    logger = logging.get_logger(__name__)\n\n    config = kwargs.pop(\"config\", None)\n    cache_dir = kwargs.pop(\"cache_dir\", None)\n    force_download = kwargs.pop(\"force_download\", False)\n    resume_download = kwargs.pop(\"resume_download\", False)\n    proxies = kwargs.pop(\"proxies\", None)\n    local_files_only = kwargs.pop(\"local_files_only\", False)\n    use_auth_token = kwargs.pop(\"use_auth_token\", None)\n    revision = kwargs.pop(\"revision\", None)\n    _ = kwargs.pop(\"mirror\", None)\n    from_pipeline = kwargs.pop(\"_from_pipeline\", None)\n    from_auto_class = kwargs.pop(\"_from_auto\", False)\n    _fast_init = kwargs.pop(\"_fast_init\", True)\n    torch_dtype = kwargs.pop(\"torch_dtype\", None)\n    subfolder = kwargs.pop(\"subfolder\", \"\")\n    commit_hash = kwargs.pop(\"_commit_hash\", None)\n    variant = kwargs.pop(\"variant\", None)\n\n    kwargs.pop(\"state_dict\", None)\n    kwargs.pop(\"from_tf\", False)\n    kwargs.pop(\"from_flax\", False)\n    kwargs.pop(\"output_loading_info\", False)\n    kwargs.pop(\"trust_remote_code\", None)\n    kwargs.pop(\"low_cpu_mem_usage\", None)\n    kwargs.pop(\"device_map\", None)\n    kwargs.pop(\"max_memory\", None)\n    kwargs.pop(\"offload_folder\", None)\n    kwargs.pop(\"offload_state_dict\", False)\n    kwargs.pop(\"load_in_8bit\", False)\n    kwargs.pop(\"load_in_4bit\", False)\n    kwargs.pop(\"quantization_config\", None)\n    kwargs.pop(\"adapter_kwargs\", {})\n    kwargs.pop(\"adapter_name\", \"default\")\n    kwargs.pop(\"use_flash_attention_2\", False)\n\n    use_safetensors = kwargs.pop(\"use_safetensors\", None if is_safetensors_available() else False)\n\n    if len(kwargs) > 0:\n        logger.warning(f\"Below kwargs may be ignored: {list(kwargs.keys())}\")\n\n    from_pt = True\n\n    user_agent = {\"file_type\": \"model\", \"framework\": \"pytorch\", \"from_auto_class\": from_auto_class}\n    if from_pipeline is not None:\n        user_agent[\"using_pipeline\"] = from_pipeline\n\n    if is_offline_mode() and not local_files_only:\n        logger.info(\"Offline mode: forcing local_files_only=True\")\n        local_files_only = True\n\n    # Load config if we don't provide a configuration\n    if not isinstance(config, PretrainedConfig):\n        config_path = config if config is not None else pretrained_model_name_or_path\n        config, model_kwargs = cls.config_class.from_pretrained(\n            config_path,\n            cache_dir=cache_dir,\n            return_unused_kwargs=True,\n            force_download=force_download,\n            resume_download=resume_download,\n            proxies=proxies,\n            local_files_only=local_files_only,\n            use_auth_token=use_auth_token,\n            revision=revision,\n            subfolder=subfolder,\n            _from_auto=from_auto_class,\n            _from_pipeline=from_pipeline,\n            **kwargs,\n        )\n    else:\n        config = copy.deepcopy(config)\n        kwarg_attn_imp = kwargs.pop(\"attn_implementation\", None)\n        if kwarg_attn_imp is not None and config._attn_implementation != kwarg_attn_imp:\n            config._attn_implementation = kwarg_attn_imp\n        model_kwargs = kwargs\n\n    if commit_hash is None:\n        commit_hash = getattr(config, \"_commit_hash\", None)\n\n    # This variable will flag if we're loading a sharded checkpoint. In this case the archive file is just the\n    # index of the files.\n\n    if pretrained_model_name_or_path is not None:\n        pretrained_model_name_or_path = str(pretrained_model_name_or_path)\n        is_local = os.path.isdir(pretrained_model_name_or_path)\n        if is_local:\n            if use_safetensors is not False and os.path.isfile(\n                os.path.join(pretrained_model_name_or_path, subfolder, _add_variant(SAFE_WEIGHTS_NAME, variant))\n            ):\n                # Load from a safetensors checkpoint\n                archive_file = os.path.join(\n                    pretrained_model_name_or_path, subfolder, _add_variant(SAFE_WEIGHTS_NAME, variant)\n                )\n            elif use_safetensors is not False and os.path.isfile(\n                os.path.join(pretrained_model_name_or_path, subfolder, _add_variant(SAFE_WEIGHTS_INDEX_NAME, variant))\n            ):\n                # Load from a sharded safetensors checkpoint\n                archive_file = os.path.join(\n                    pretrained_model_name_or_path, subfolder, _add_variant(SAFE_WEIGHTS_INDEX_NAME, variant)\n                )\n            elif os.path.isfile(\n                os.path.join(pretrained_model_name_or_path, subfolder, _add_variant(WEIGHTS_NAME, variant))\n            ):\n                # Load from a PyTorch checkpoint\n                archive_file = os.path.join(\n                    pretrained_model_name_or_path, subfolder, _add_variant(WEIGHTS_NAME, variant)\n                )\n            elif os.path.isfile(\n                os.path.join(pretrained_model_name_or_path, subfolder, _add_variant(WEIGHTS_INDEX_NAME, variant))\n            ):\n                # Load from a sharded PyTorch checkpoint\n                archive_file = os.path.join(\n                    pretrained_model_name_or_path, subfolder, _add_variant(WEIGHTS_INDEX_NAME, variant)\n                )\n            else:\n                raise EnvironmentError(\n                    f\"Error no file named {_add_variant(WEIGHTS_NAME, variant)} found in directory\"\n                    f\" {pretrained_model_name_or_path}.\"\n                )\n        elif os.path.isfile(os.path.join(subfolder, pretrained_model_name_or_path)):\n            archive_file = pretrained_model_name_or_path\n            is_local = True\n        elif is_remote_url(pretrained_model_name_or_path):\n            filename = pretrained_model_name_or_path\n            resolved_archive_file = download_url(pretrained_model_name_or_path)\n        else:\n            # set correct filename\n            if use_safetensors is not False:\n                filename = _add_variant(SAFE_WEIGHTS_NAME, variant)\n            else:\n                filename = _add_variant(WEIGHTS_NAME, variant)\n\n            try:\n                # Load from URL or cache if already cached\n                cached_file_kwargs = {\n                    \"cache_dir\": cache_dir,\n                    \"force_download\": force_download,\n                    \"proxies\": proxies,\n                    \"resume_download\": resume_download,\n                    \"local_files_only\": local_files_only,\n                    \"use_auth_token\": use_auth_token,\n                    \"user_agent\": user_agent,\n                    \"revision\": revision,\n                    \"subfolder\": subfolder,\n                    \"_raise_exceptions_for_missing_entries\": False,\n                    \"_commit_hash\": commit_hash,\n                }\n                resolved_archive_file = cached_file(pretrained_model_name_or_path, filename, **cached_file_kwargs)\n\n                # Since we set _raise_exceptions_for_missing_entries=False, we don't get an exception but a None\n                # result when internet is up, the repo and revision exist, but the file does not.\n                if resolved_archive_file is None and filename == _add_variant(SAFE_WEIGHTS_NAME, variant):\n                    # Maybe the checkpoint is sharded, we try to grab the index name in this case.\n                    resolved_archive_file = cached_file(\n                        pretrained_model_name_or_path,\n                        _add_variant(SAFE_WEIGHTS_INDEX_NAME, variant),\n                        **cached_file_kwargs,\n                    )\n                    if resolved_archive_file is not None:\n                        pass\n                    elif use_safetensors:\n                        raise EnvironmentError(\n                            f\" {_add_variant(SAFE_WEIGHTS_NAME, variant)} or {_add_variant(SAFE_WEIGHTS_INDEX_NAME, variant)} and thus cannot be loaded with `safetensors`. Please make sure that the model has been saved with `safe_serialization=True` or do not set `use_safetensors=True`.\"\n                        )\n                    else:\n                        # This repo has no safetensors file of any kind, we switch to PyTorch.\n                        filename = _add_variant(WEIGHTS_NAME, variant)\n                        resolved_archive_file = cached_file(\n                            pretrained_model_name_or_path, filename, **cached_file_kwargs\n                        )\n                if resolved_archive_file is None and filename == _add_variant(WEIGHTS_NAME, variant):\n                    # Maybe the checkpoint is sharded, we try to grab the index name in this case.\n                    resolved_archive_file = cached_file(\n                        pretrained_model_name_or_path,\n                        _add_variant(WEIGHTS_INDEX_NAME, variant),\n                        **cached_file_kwargs,\n                    )\n                    if resolved_archive_file is not None:\n                        pass\n                if resolved_archive_file is None:\n                    # Otherwise, maybe there is a TF or Flax model file.  We try those to give a helpful error\n                    # message.\n                    has_file_kwargs = {\n                        \"revision\": revision,\n                        \"proxies\": proxies,\n                        \"use_auth_token\": use_auth_token,\n                    }\n                    if variant is not None and has_file(pretrained_model_name_or_path, WEIGHTS_NAME, **has_file_kwargs):\n                        raise EnvironmentError(\n                            f\"{pretrained_model_name_or_path} does not appear to have a file named\"\n                            f\" {_add_variant(WEIGHTS_NAME, variant)} but there is a file without the variant\"\n                            f\" {variant}. Use `variant=None` to load this model from those weights.\"\n                        )\n                    else:\n                        raise EnvironmentError(\n                            f\"{pretrained_model_name_or_path} does not appear to have a file named\"\n                            f\" {_add_variant(WEIGHTS_NAME, variant)}\"\n                        )\n            except EnvironmentError:\n                # Raise any environment error raise by `cached_file`. It will have a helpful error message adapted\n                # to the original exception.\n                raise\n            except Exception:\n                # For any other exception, we throw a generic error.\n                raise EnvironmentError(\n                    f\"Can't load the model for '{pretrained_model_name_or_path}'. If you were trying to load it\"\n                    \" from 'https://huggingface.co/models', make sure you don't have a local directory with the\"\n                    f\" same name. Otherwise, make sure '{pretrained_model_name_or_path}' is the correct path to a\"\n                    f\" directory containing a file named {_add_variant(WEIGHTS_NAME, variant)}.\"\n                )\n\n        if is_local:\n            logger.info(f\"loading weights file {archive_file}\")\n            resolved_archive_file = archive_file\n        else:\n            logger.info(f\"loading weights file {filename} from cache at {resolved_archive_file}\")\n    else:\n        resolved_archive_file = None\n\n    if from_pt:\n        # set dtype to instantiate the model under:\n        # 1. If torch_dtype is not None, we use that dtype\n        dtype_orig = None\n\n        if torch_dtype is not None:\n            if not isinstance(torch_dtype, torch.dtype):\n                raise ValueError(f\"`torch_dtype` can be either `torch.dtype` or `None`, but received {torch_dtype}\")\n            dtype_orig = cls._set_default_torch_dtype(torch_dtype)\n\n    config.name_or_path = pretrained_model_name_or_path\n\n    # Instantiate model.\n    init_contexts = [no_init_weights(_enable=_fast_init)]\n\n    with ContextManagers(init_contexts):\n        model = cls(config, *model_args, **model_kwargs)\n\n    if from_pt:\n        # restore default dtype\n        if dtype_orig is not None:\n            torch.set_default_dtype(dtype_orig)\n\n    # make sure token embedding weights are still tied if needed\n    model.tie_weights()\n\n    # Set model in evaluation mode to deactivate DropOut modules by default\n    model.eval()\n\n    # If it is a model with generation capabilities, attempt to load the generation config\n    if model.can_generate():\n        try:\n            model.generation_config = GenerationConfig.from_pretrained(\n                pretrained_model_name_or_path,\n                cache_dir=cache_dir,\n                force_download=force_download,\n                resume_download=resume_download,\n                proxies=proxies,\n                local_files_only=local_files_only,\n                use_auth_token=use_auth_token,\n                revision=revision,\n                subfolder=subfolder,\n                _from_auto=from_auto_class,\n                _from_pipeline=from_pipeline,\n                **kwargs,\n            )\n        except (OSError, TypeError):\n            logger.info(\"Generation config file not found, using a generation config created from the model config.\")\n\n    # set pretrained path\n    if resolved_archive_file:\n        pretrained_interface.set_pretrained_path(model, resolved_archive_file)\n\n    return model\n", "path": "colossalai/lazy/pretrained.py"}]}