rixty/KoboldAI-Client
1
1
Fork 0
mirror of https://github.com/KoboldAI/KoboldAI-Client.git synced 2025-06-05 21:59:24 +02:00
Code Issues Projects Releases Wiki Activity

Seems to work

Browse Source
This commit is contained in:
somebody
2023-05-31 14:42:05 -05:00
parent 24b0b32829
commit f326fc07e8
2 changed files with 4 additions and 121 deletions
Download Patch File Download Diff File Expand all files Collapse all files

87
modeling/patches.py
View File

@@ -148,29 +148,6 @@ def patch_transformers_for_lazyload() -> None:
limitations under the License.
"""
import torch
import accelerate
# _old_set_module_tensor_to_device = (
# accelerate.utils.modeling.set_module_tensor_to_device
# )
# def _set_module_tensor_to_device(
# module: torch.nn.Module,
# tensor_name: str,
# device: Union[int, str, torch.device],
# value: Optional[torch.Tensor] = None,
# dtype: Optional[Union[str, torch.dtype]] = None,
# ):
# if isinstance(value, LazyTensor):
# value = value.materialize()
# print("HEY!", dtype)
# return _old_set_module_tensor_to_device(
# module, tensor_name, device, value, dtype
# )
# accelerate.utils.modeling.set_module_tensor_to_device = _set_module_tensor_to_device
from accelerate.utils.modeling import named_module_tensors
from accelerate.utils import set_module_tensor_to_device, offload_weight
def _load_state_dict_into_meta_model(
@@ -225,10 +202,7 @@ def patch_transformers_for_lazyload() -> None:
for old_key, new_key in zip(old_keys, new_keys):
state_dict[new_key] = state_dict.pop(old_key)
# BEGIN PATCH
# TODO: Based on config
# dtype = torch.float16
# BEGIN PATCH
for param_name, param in sorted(
state_dict.items(),
# State dict must be ordered in this manner to make the caching in
@@ -243,7 +217,7 @@ def patch_transformers_for_lazyload() -> None:
if isinstance(param, LazyTensor):
# Should always be true
param = param.materialize()
# END PATCH
# END PATCH
# First part of the test is always true as load_state_dict_keys always contains state_dict keys.
if (
@@ -338,63 +312,6 @@ def patch_transformers_for_lazyload() -> None:
_load_state_dict_into_meta_model
)
# # Patch AlignDevicesHook to hack around OPT lm_head
# HACK_ZERO_ON_FAIL_TENSORS = ["lm_head.weight"]
# def _init_hook(self, module):
# if not self.offload and self.execution_device is not None:
# # BEGIN PATCH
# for name, tensor in named_module_tensors(
# module, recurse=self.place_submodules
# ):
# try:
# set_module_tensor_to_device(module, name, self.execution_device)
# except ValueError:
# # ValueError: weight is on the meta device, we need a `value` to put in on 0.
# # bleuuuuuuuuuuuuuuuhhh
# if name in HACK_ZERO_ON_FAIL_TENSORS:
# logger.warning(
# f"Couldn't find value for weight {name}, zeroing."
# )
# set_module_tensor_to_device(
# module,
# name,
# self.execution_device,
# value=torch.zeros(tensor.shape),
# )
# # END PATCH
# elif self.offload:
# self.original_devices = {
# name: param.device
# for name, param in named_module_tensors(
# module, recurse=self.place_submodules
# )
# }
# if self.weights_map is None:
# self.weights_map = {
# name: param.to("cpu")
# for name, param in named_module_tensors(
# module,
# include_buffers=self.offload_buffers,
# recurse=self.place_submodules,
# )
# }
# for name, _ in named_module_tensors(
# module,
# include_buffers=self.offload_buffers,
# recurse=self.place_submodules,
# ):
# set_module_tensor_to_device(module, name, "meta")
# if not self.offload_buffers and self.execution_device is not None:
# for name, _ in module.named_buffers(recurse=self.place_submodules):
# set_module_tensor_to_device(module, name, self.execution_device)
# return module
# accelerate.hooks.AlignDevicesHook.init_hook = _init_hook
def patch_transformers() -> None:
patch_transformers_download()