From 1546b9efaa95d7388b853d25ff3e8b756526f44e Mon Sep 17 00:00:00 2001 From: somebody Date: Sat, 27 May 2023 16:31:53 -0500 Subject: [PATCH] Hello its breaking breakmodel time --- aiserver.py | 3 + breakmodel.py | 955 ------------------ .../generic_hf_torch/class.py | 37 +- modeling/inference_models/hf.py | 6 +- modeling/inference_models/hf_torch.py | 148 +-- modeling/lazy_loader.py | 10 +- modeling/patches.py | 168 +++ utils.py | 6 +- 8 files changed, 236 insertions(+), 1097 deletions(-) delete mode 100644 breakmodel.py diff --git a/aiserver.py b/aiserver.py index faffe0e2..ad1efdab 100644 --- a/aiserver.py +++ b/aiserver.py @@ -30,6 +30,9 @@ logging.getLogger("urllib3").setLevel(logging.ERROR) import attention_bias attention_bias.do_patches() +from modeling import patches +patches.patch_transformers_for_lazyload() + from os import path, getcwd import time import re diff --git a/breakmodel.py b/breakmodel.py deleted file mode 100644 index 75bc03cc..00000000 --- a/breakmodel.py +++ /dev/null @@ -1,955 +0,0 @@ -''' -This is a MODIFIED version of arrmansa's low VRAM patch. -https://github.com/arrmansa/Basic-UI-for-GPT-J-6B-with-low-vram/blob/main/GPT-J-6B-Low-Vram-UI.ipynb -The ORIGINAL version of the patch is released under the Apache License 2.0 -Copyright 2021 arrmansa -Copyright 2021 finetuneanon -Copyright 2018, 2022 The Hugging Face team - - - Apache License - Version 2.0, January 2004 - http://www.apache.org/licenses/ - - TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION - - 1. Definitions. - - "License" shall mean the terms and conditions for use, reproduction, - and distribution as defined by Sections 1 through 9 of this document. - - "Licensor" shall mean the copyright owner or entity authorized by - the copyright owner that is granting the License. - - "Legal Entity" shall mean the union of the acting entity and all - other entities that control, are controlled by, or are under common - control with that entity. For the purposes of this definition, - "control" means (i) the power, direct or indirect, to cause the - direction or management of such entity, whether by contract or - otherwise, or (ii) ownership of fifty percent (50%) or more of the - outstanding shares, or (iii) beneficial ownership of such entity. - - "You" (or "Your") shall mean an individual or Legal Entity - exercising permissions granted by this License. - - "Source" form shall mean the preferred form for making modifications, - including but not limited to software source code, documentation - source, and configuration files. - - "Object" form shall mean any form resulting from mechanical - transformation or translation of a Source form, including but - not limited to compiled object code, generated documentation, - and conversions to other media types. - - "Work" shall mean the work of authorship, whether in Source or - Object form, made available under the License, as indicated by a - copyright notice that is included in or attached to the work - (an example is provided in the Appendix below). - - "Derivative Works" shall mean any work, whether in Source or Object - form, that is based on (or derived from) the Work and for which the - editorial revisions, annotations, elaborations, or other modifications - represent, as a whole, an original work of authorship. For the purposes - of this License, Derivative Works shall not include works that remain - separable from, or merely link (or bind by name) to the interfaces of, - the Work and Derivative Works thereof. - - "Contribution" shall mean any work of authorship, including - the original version of the Work and any modifications or additions - to that Work or Derivative Works thereof, that is intentionally - submitted to Licensor for inclusion in the Work by the copyright owner - or by an individual or Legal Entity authorized to submit on behalf of - the copyright owner. For the purposes of this definition, "submitted" - means any form of electronic, verbal, or written communication sent - to the Licensor or its representatives, including but not limited to - communication on electronic mailing lists, source code control systems, - and issue tracking systems that are managed by, or on behalf of, the - Licensor for the purpose of discussing and improving the Work, but - excluding communication that is conspicuously marked or otherwise - designated in writing by the copyright owner as "Not a Contribution." - - "Contributor" shall mean Licensor and any individual or Legal Entity - on behalf of whom a Contribution has been received by Licensor and - subsequently incorporated within the Work. - - 2. Grant of Copyright License. Subject to the terms and conditions of - this License, each Contributor hereby grants to You a perpetual, - worldwide, non-exclusive, no-charge, royalty-free, irrevocable - copyright license to reproduce, prepare Derivative Works of, - publicly display, publicly perform, sublicense, and distribute the - Work and such Derivative Works in Source or Object form. - - 3. Grant of Patent License. Subject to the terms and conditions of - this License, each Contributor hereby grants to You a perpetual, - worldwide, non-exclusive, no-charge, royalty-free, irrevocable - (except as stated in this section) patent license to make, have made, - use, offer to sell, sell, import, and otherwise transfer the Work, - where such license applies only to those patent claims licensable - by such Contributor that are necessarily infringed by their - Contribution(s) alone or by combination of their Contribution(s) - with the Work to which such Contribution(s) was submitted. If You - institute patent litigation against any entity (including a - cross-claim or counterclaim in a lawsuit) alleging that the Work - or a Contribution incorporated within the Work constitutes direct - or contributory patent infringement, then any patent licenses - granted to You under this License for that Work shall terminate - as of the date such litigation is filed. - - 4. Redistribution. You may reproduce and distribute copies of the - Work or Derivative Works thereof in any medium, with or without - modifications, and in Source or Object form, provided that You - meet the following conditions: - - (a) You must give any other recipients of the Work or - Derivative Works a copy of this License; and - - (b) You must cause any modified files to carry prominent notices - stating that You changed the files; and - - (c) You must retain, in the Source form of any Derivative Works - that You distribute, all copyright, patent, trademark, and - attribution notices from the Source form of the Work, - excluding those notices that do not pertain to any part of - the Derivative Works; and - - (d) If the Work includes a "NOTICE" text file as part of its - distribution, then any Derivative Works that You distribute must - include a readable copy of the attribution notices contained - within such NOTICE file, excluding those notices that do not - pertain to any part of the Derivative Works, in at least one - of the following places: within a NOTICE text file distributed - as part of the Derivative Works; within the Source form or - documentation, if provided along with the Derivative Works; or, - within a display generated by the Derivative Works, if and - wherever such third-party notices normally appear. The contents - of the NOTICE file are for informational purposes only and - do not modify the License. You may add Your own attribution - notices within Derivative Works that You distribute, alongside - or as an addendum to the NOTICE text from the Work, provided - that such additional attribution notices cannot be construed - as modifying the License. - - You may add Your own copyright statement to Your modifications and - may provide additional or different license terms and conditions - for use, reproduction, or distribution of Your modifications, or - for any such Derivative Works as a whole, provided Your use, - reproduction, and distribution of the Work otherwise complies with - the conditions stated in this License. - - 5. Submission of Contributions. Unless You explicitly state otherwise, - any Contribution intentionally submitted for inclusion in the Work - by You to the Licensor shall be under the terms and conditions of - this License, without any additional terms or conditions. - Notwithstanding the above, nothing herein shall supersede or modify - the terms of any separate license agreement you may have executed - with Licensor regarding such Contributions. - - 6. Trademarks. This License does not grant permission to use the trade - names, trademarks, service marks, or product names of the Licensor, - except as required for reasonable and customary use in describing the - origin of the Work and reproducing the content of the NOTICE file. - - 7. Disclaimer of Warranty. Unless required by applicable law or - agreed to in writing, Licensor provides the Work (and each - Contributor provides its Contributions) on an "AS IS" BASIS, - WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or - implied, including, without limitation, any warranties or conditions - of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A - PARTICULAR PURPOSE. You are solely responsible for determining the - appropriateness of using or redistributing the Work and assume any - risks associated with Your exercise of permissions under this License. - - 8. Limitation of Liability. In no event and under no legal theory, - whether in tort (including negligence), contract, or otherwise, - unless required by applicable law (such as deliberate and grossly - negligent acts) or agreed to in writing, shall any Contributor be - liable to You for damages, including any direct, indirect, special, - incidental, or consequential damages of any character arising as a - result of this License or out of the use or inability to use the - Work (including but not limited to damages for loss of goodwill, - work stoppage, computer failure or malfunction, or any and all - other commercial damages or losses), even if such Contributor - has been advised of the possibility of such damages. - - 9. Accepting Warranty or Additional Liability. While redistributing - the Work or Derivative Works thereof, You may choose to offer, - and charge a fee for, acceptance of support, warranty, indemnity, - or other liability obligations and/or rights consistent with this - License. However, in accepting such obligations, You may act only - on Your own behalf and on Your sole responsibility, not on behalf - of any other Contributor, and only if You agree to indemnify, - defend, and hold each Contributor harmless for any liability - incurred by, or claims asserted against, such Contributor by reason - of your accepting any such warranty or additional liability. - - END OF TERMS AND CONDITIONS - - APPENDIX: How to apply the Apache License to your work. - - To apply the Apache License to your work, attach the following - boilerplate notice, with the fields enclosed by brackets "[]" - replaced with your own identifying information. (Don't include - the brackets!) The text should be enclosed in the appropriate - comment syntax for the file format. We also recommend that a - file or class name and description of purpose be included on the - same "printed page" as the copyright notice for easier - identification within third-party archives. - - Copyright [yyyy] [name of copyright owner] - - Licensed under the Apache License, Version 2.0 (the "License"); - you may not use this file except in compliance with the License. - You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - - Unless required by applicable law or agreed to in writing, software - distributed under the License is distributed on an "AS IS" BASIS, - WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - See the License for the specific language governing permissions and - limitations under the License. -''' - - -import torch -from torch import nn -import torch.cuda.comm -import copy -import gc -import os -import sys -import itertools -import bisect -import random -import utils -from typing import Dict, List, Optional, Union - -from transformers.modeling_outputs import BaseModelOutputWithPast, BaseModelOutputWithPastAndCrossAttentions - -from transformers.utils import logging -logger = logging.get_logger(__name__) - - -breakmodel = True -gpu_blocks = [] -disk_blocks = 0 -primary_device = 0 if torch.cuda.device_count() > 0 else "cpu" - -from accelerate.hooks import attach_align_device_hook_on_blocks -from accelerate.utils import OffloadedWeightsLoader, check_device_map, extract_submodules_state_dict, offload_state_dict -from accelerate import dispatch_model - -def dispatch_model_ex( - model: nn.Module, - device_map: Dict[str, Union[str, int, torch.device]], - main_device: Optional[torch.device] = None, - state_dict: Optional[Dict[str, torch.Tensor]] = None, - offload_dir: Union[str, os.PathLike] = None, - offload_buffers: bool = False, - **kwargs, -): - """ - This is a modified version of - https://github.com/huggingface/accelerate/blob/eeaba598f455fbd2c48661d7e816d3ff25ab050b/src/accelerate/big_modeling.py#L130 - that still works when the main device is the CPU. - - Dispatches a model according to a given device map. Layers of the model might be spread across GPUs, offloaded on - the CPU or even the disk. - - Args: - model (`torch.nn.Module`): - The model to dispatch. - device_map (`Dict[str, Union[str, int, torch.device]]`): - A dictionary mapping module names in the models `state_dict` to the device they should go to. Note that - `"disk"` is accepted even if it's not a proper value for `torch.device`. - main_device (`str`, `int` or `torch.device`, *optional*): - The main execution device. Will default to the first device in the `device_map` different from `"cpu"` or - `"disk"`. - state_dict (`Dict[str, torch.Tensor]`, *optional*): - The state dict of the part of the model that will be kept on CPU. - offload_dir (`str` or `os.PathLike`): - The folder in which to offload the model weights (or where the model weights are already offloaded). - offload_buffers (`bool`, *optional*, defaults to `False`): - Whether or not to offload the buffers with the model parameters. - preload_module_classes (`List[str]`, *optional*): - A list of classes whose instances should load all their weights (even in the submodules) at the beginning - of the forward. This should only be used for classes that have submodules which are registered but not - called directly during the forward, for instance if a `dense` linear layer is registered, but at forward, - `dense.weight` and `dense.bias` are used in some operations instead of calling `dense` directly. - """ - if main_device != "cpu": - return dispatch_model(model, device_map, main_device, state_dict, offload_dir=offload_dir, offload_buffers=offload_buffers, **kwargs) - - # Error early if the device map is incomplete. - check_device_map(model, device_map) - - offload_devices = ["cpu", "disk"] if main_device != "cpu" else ["disk"] - - if main_device is None: - main_device = [d for d in device_map.values() if d not in offload_devices][0] - - cpu_modules = [name for name, device in device_map.items() if device == "cpu"] if main_device != "cpu" else [] - if state_dict is None and len(cpu_modules) > 0: - state_dict = extract_submodules_state_dict(model.state_dict(), cpu_modules) - - disk_modules = [name for name, device in device_map.items() if device == "disk"] - if offload_dir is None and len(disk_modules) > 0: - raise ValueError( - "We need an `offload_dir` to dispatch this model according to this `device_map`, the following submodules " - f"need to be offloaded: {', '.join(disk_modules)}." - ) - if len(disk_modules) > 0 and ( - not os.path.isdir(offload_dir) or not os.path.isfile(os.path.join(offload_dir, "index.json")) - ): - disk_state_dict = extract_submodules_state_dict(model.state_dict(), disk_modules) - offload_state_dict(offload_dir, disk_state_dict) - - execution_device = { - name: main_device if device in offload_devices else device for name, device in device_map.items() - } - offload = {name: device in offload_devices for name, device in device_map.items()} - save_folder = offload_dir if len(disk_modules) > 0 else None - if state_dict is not None or save_folder is not None: - weights_map = OffloadedWeightsLoader(state_dict=state_dict, save_folder=save_folder) - else: - weights_map = None - - attach_align_device_hook_on_blocks( - model, - execution_device=execution_device, - offload=offload, - offload_buffers=offload_buffers, - weights_map=weights_map, - **kwargs, - ) - model.hf_device_map = device_map - return model - - -# Copied from transformers.models.bart.modeling_bart._expand_mask -def _expand_mask(mask: torch.Tensor, dtype: torch.dtype, tgt_len: Optional[int] = None): - """ - Expands attention_mask from `[bsz, seq_len]` to `[bsz, 1, tgt_seq_len, src_seq_len]`. - """ - bsz, src_len = mask.size() - tgt_len = tgt_len if tgt_len is not None else src_len - - expanded_mask = mask[:, None, None, :].expand(bsz, 1, tgt_len, src_len).to(dtype) - - inverted_mask = 1.0 - expanded_mask - - return inverted_mask.masked_fill(inverted_mask.bool(), torch.finfo(dtype).min) - - -def move_hidden_layers(transformer, h=None): - if h is None: - h = transformer.h - - assert len(gpu_blocks) <= torch.cuda.device_count() - assert sum(gpu_blocks) <= len(h) - ram_blocks = len(h) - sum(gpu_blocks) - - transformer.extrastorage = {} - torch.cuda.empty_cache() - - able_to_pin_layers = True - for i in range(ram_blocks): - h[i].to("cpu") - transformer.extrastorage[i] = copy.deepcopy(h[i]) - smalltensor = torch.tensor(0).to(primary_device) - for param1 in h[i].parameters(): - param1.data = smalltensor - h[i].to(primary_device) - for param in transformer.extrastorage[i].parameters(): - param.requires_grad = False - param.data = param.data.detach() - if able_to_pin_layers: - try: - param.data = param.data.pin_memory() - except: - able_to_pin_layers = False - print(f"WARNING: You only have enough shared GPU memory for {i} out of {ram_blocks} CPU layers. Expect suboptimal speed.", file=sys.stderr) - gc.collect() - torch.cuda.empty_cache() - - if ram_blocks: - for param1,param2 in zip(h[0].parameters(),transformer.extrastorage[0].parameters()): - param1.data = param2.data.to(primary_device, non_blocking=False).detach() - - for param1,param2 in zip(h[ram_blocks-1].parameters(),transformer.extrastorage[ram_blocks-1].parameters()): - param1.data = param2.data.to(primary_device, non_blocking=False).detach() - - i = ram_blocks - for j in range(len(gpu_blocks)): - for _ in range(gpu_blocks[j]): - h[i].to(j) - i += 1 - - -def new_forward_neo( - self, - input_ids=None, - past_key_values=None, - attention_mask=None, - token_type_ids=None, - position_ids=None, - head_mask=None, - inputs_embeds=None, - use_cache=None, - output_attentions=None, - output_hidden_states=None, - return_dict=None, - embs=None, -): - assert len(gpu_blocks) <= torch.cuda.device_count() - assert sum(gpu_blocks) <= len(self.h) - ram_blocks = len(self.h) - sum(gpu_blocks) - cumulative_gpu_blocks = tuple(itertools.accumulate(gpu_blocks)) - - - output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions - output_hidden_states = ( - output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states - ) - use_cache = use_cache if use_cache is not None else self.config.use_cache - return_dict = return_dict if return_dict is not None else self.config.use_return_dict - - if input_ids is not None and inputs_embeds is not None: - raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time") - elif input_ids is not None: - input_shape = input_ids.size() - input_ids = input_ids.view(-1, input_shape[-1]) - batch_size = input_ids.shape[0] - elif inputs_embeds is not None: - input_shape = inputs_embeds.size()[:-1] - batch_size = inputs_embeds.shape[0] - else: - raise ValueError("You have to specify either input_ids or inputs_embeds") - - device = input_ids.device if input_ids is not None else inputs_embeds.device - - if token_type_ids is not None: - token_type_ids = token_type_ids.view(-1, input_shape[-1]) - if position_ids is not None: - position_ids = position_ids.view(-1, input_shape[-1]) - - if past_key_values is None: - past_length = 0 - past_key_values = tuple([None] * len(self.h)) - else: - past_length = past_key_values[0][0].size(-2) - - device = primary_device if breakmodel else input_ids.device if input_ids is not None else inputs_embeds.device - if position_ids is None: - position_ids = torch.arange(past_length, input_shape[-1] + past_length, dtype=torch.long, device=device) - position_ids = position_ids.unsqueeze(0).view(-1, input_shape[-1]) - - # Attention mask. - if attention_mask is not None: - assert batch_size > 0, "batch_size has to be defined and > 0" - attention_mask = attention_mask.view(batch_size, -1) - # We create a 3D attention mask from a 2D tensor mask. - # Sizes are [batch_size, 1, 1, to_seq_length] - # So we can broadcast to [batch_size, num_heads, from_seq_length, to_seq_length] - # this attention mask is more simple than the triangular masking of causal attention - # used in OpenAI GPT, we just need to prepare the broadcast dimension here. - attention_mask = attention_mask[:, None, None, :] - - # Since attention_mask is 1.0 for positions we want to attend and 0.0 for - # masked positions, this operation will create a tensor which is 0.0 for - # positions we want to attend and -10000.0 for masked positions. - # Since we are adding it to the raw scores before the softmax, this is - # effectively the same as removing these entirely. - attention_mask = attention_mask.to(dtype=self.dtype) # fp16 compatibility - attention_mask = (1.0 - attention_mask) * -10000.0 - - # Prepare head mask if needed - # 1.0 in head_mask indicate we keep the head - # attention_probs has shape bsz x num_heads x N x N - # head_mask has shape n_layer x batch x num_heads x N x N - head_mask = self.get_head_mask(head_mask, getattr(self.config, "num_layers", None) or self.config.n_layer) - - if inputs_embeds is None: - if breakmodel: - input_ids = input_ids.to(primary_device) - inputs_embeds = self.wte(input_ids) - - if embs is not None and not (use_cache is not None and use_cache and past_key_values is not None and len(past_key_values) > 0 and past_key_values[0] is not None): - offset = 0 - for pos, emb in embs: - pos += offset - if len(emb.shape) == 2: - emb = emb.repeat(input_shape[0], 1, 1) - inputs_embeds[:, pos:pos+emb.shape[1]] = emb - offset += emb.shape[1] - - if getattr(self, "wpe", None) is None: - hidden_states = inputs_embeds - else: - if breakmodel: - position_ids = position_ids.to(primary_device) - position_embeds = self.wpe(position_ids) - if breakmodel: - position_embeds = position_embeds.to(primary_device) - hidden_states = inputs_embeds + position_embeds - - if token_type_ids is not None: - token_type_embeds = self.wte(token_type_ids) - hidden_states = hidden_states + token_type_embeds - - hidden_states = self.drop(hidden_states) - - output_shape = input_shape + (hidden_states.size(-1),) - - presents = () if use_cache else None - all_self_attentions = () if output_attentions else None - all_hidden_states = () if output_hidden_states else None - - if breakmodel and ram_blocks: - copystream = torch.cuda.Stream(device=primary_device, priority=-1) - - for i, (block, layer_past) in enumerate(zip(self.h, past_key_values)): - - if breakmodel: - if i in range(ram_blocks): - index1 = (i+1)%ram_blocks - for param1,param2 in zip(self.h[index1].parameters(),self.h[(i-1)%ram_blocks].parameters()): - param1.data = param2.data - for param1,param2 in zip(self.h[index1].parameters(),self.extrastorage[index1].parameters()): - with torch.cuda.stream(copystream): - torch.cuda.comm.broadcast(param2.data,out = [param1.data]) - - if output_hidden_states: - all_hidden_states = all_hidden_states + (hidden_states.cpu(),) - - if getattr(self.config, "gradient_checkpointing", False) and self.training: - - if use_cache: - logger.warning( - "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..." - ) - use_cache = False - - def create_custom_forward(module): - def custom_forward(*inputs): - # None for past_key_value - return module(*inputs, use_cache, output_attentions) - - return custom_forward - - outputs = torch.utils.checkpoint.checkpoint( - create_custom_forward(block), - hidden_states, - None, - attention_mask, - head_mask[i], - ) - else: - if breakmodel: - device = primary_device if i < ram_blocks else bisect.bisect_right(cumulative_gpu_blocks, i - ram_blocks) - outputs = block( - hidden_states.to(device) if breakmodel and hidden_states is not None else hidden_states, - layer_past=tuple(v.to(device) for v in layer_past if v is not None) if breakmodel and layer_past is not None and i >= ram_blocks and len(layer_past) and layer_past[0].device.index != device else layer_past, - attention_mask=attention_mask.to(device) if breakmodel and attention_mask is not None else attention_mask, - head_mask=head_mask[i].to(device) if breakmodel and head_mask[i] is not None else head_mask[i], - use_cache=use_cache, - output_attentions=output_attentions, - ) - - hidden_states = outputs[0] - if use_cache is True: - presents = presents + (outputs[1],) - - if output_attentions: - all_self_attentions = all_self_attentions + (outputs[2 if use_cache else 1],) - - - if breakmodel: - if i in range(ram_blocks): - torch.cuda.synchronize() - torch.cuda.empty_cache() - - if breakmodel: - if ram_blocks: - del copystream - torch.cuda.empty_cache() - hidden_states = hidden_states.to(primary_device) - hidden_states = self.ln_f(hidden_states) - if breakmodel: - hidden_states = hidden_states.to(primary_device) - - hidden_states = hidden_states.view(*output_shape) - # Add last hidden state - if output_hidden_states: - all_hidden_states = all_hidden_states + (hidden_states,) - - if not return_dict: - return tuple(v for v in [hidden_states, presents, all_hidden_states, all_self_attentions] if v is not None) - return BaseModelOutputWithPast( - last_hidden_state=hidden_states, - past_key_values=presents, - hidden_states=all_hidden_states, - attentions=all_self_attentions, - ) - - -def new_forward_xglm( - self, - input_ids=None, - attention_mask=None, - encoder_hidden_states=None, - encoder_attention_mask=None, - head_mask=None, - cross_attn_head_mask=None, - past_key_values=None, - inputs_embeds=None, - use_cache=None, - output_attentions=None, - output_hidden_states=None, - return_dict=None, -): - assert len(gpu_blocks) <= torch.cuda.device_count() - assert sum(gpu_blocks) <= len(self.layers) - ram_blocks = len(self.layers) - sum(gpu_blocks) - cumulative_gpu_blocks = tuple(itertools.accumulate(gpu_blocks)) - - - output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions - output_hidden_states = ( - output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states - ) - use_cache = use_cache if use_cache is not None else self.config.use_cache - return_dict = return_dict if return_dict is not None else self.config.use_return_dict - - # retrieve input_ids and inputs_embeds - if input_ids is not None and inputs_embeds is not None: - raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time") - elif input_ids is not None: - input_shape = input_ids.size() - input_ids = input_ids.view(-1, input_shape[-1]) - elif inputs_embeds is not None: - input_shape = inputs_embeds.size()[:-1] - else: - raise ValueError("You have to specify either input_ids or inputs_embeds") - - # past_key_values_length - past_key_values_length = past_key_values[0][0].shape[2] if past_key_values is not None else 0 - - if inputs_embeds is None: - if breakmodel: - input_ids = input_ids.to(primary_device) - inputs_embeds = self.embed_tokens(input_ids) * self.embed_scale - - attention_mask = self._prepare_decoder_attention_mask( - attention_mask, input_shape, inputs_embeds, past_key_values_length - ) - - # expand encoder attention mask - if encoder_hidden_states is not None and encoder_attention_mask is not None: - # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len] - encoder_attention_mask = _expand_mask(encoder_attention_mask, inputs_embeds.dtype, tgt_len=input_shape[-1]) - - # embed positions - if breakmodel: - inputs_embeds = inputs_embeds.to(primary_device) - positions = self.embed_positions(input_ids, inputs_embeds, past_key_values_length) - if breakmodel: - positions = positions.to(primary_device) - - hidden_states = inputs_embeds + positions - - hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training) - - # decoder layers - all_hidden_states = () if output_hidden_states else None - all_self_attns = () if output_attentions else None - all_cross_attentions = () if (output_attentions and encoder_hidden_states is not None) else None - next_decoder_cache = () if use_cache else None - - if breakmodel and ram_blocks: - copystream = torch.cuda.Stream(device=primary_device, priority=-1) - - # check if head_mask/cross_attn_head_mask has a correct number of layers specified if desired - for attn_mask, mask_name in zip([head_mask, cross_attn_head_mask], ["head_mask", "cross_attn_head_mask"]): - if attn_mask is not None: - assert attn_mask.size()[0] == ( - len(self.layers) - ), f"The `{mask_name}` should be specified for {len(self.layers)} layers, but it is for {head_mask.size()[0]}." - for idx, decoder_layer in enumerate(self.layers): - i = idx - if breakmodel: - if i in range(ram_blocks): - index1 = (i+1)%ram_blocks - for param1,param2 in zip(self.layers[index1].parameters(),self.layers[(i-1)%ram_blocks].parameters()): - param1.data = param2.data - for param1,param2 in zip(self.layers[index1].parameters(),self.extrastorage[index1].parameters()): - with torch.cuda.stream(copystream): - torch.cuda.comm.broadcast(param2.data,out = [param1.data]) - - # add LayerDrop (see https://arxiv.org/abs/1909.11556 for description) - if output_hidden_states: - all_hidden_states += (hidden_states,) - dropout_probability = random.uniform(0, 1) - if self.training and (dropout_probability < self.layerdrop): - continue - - past_key_value = past_key_values[idx] if past_key_values is not None else None - - if self.gradient_checkpointing and self.training: - - if use_cache: - logger.warning( - "`use_cache = True` is incompatible with gradient checkpointing`. Setting `use_cache = False`..." - ) - use_cache = False - - def create_custom_forward(module): - def custom_forward(*inputs): - # None for past_key_value - return module(*inputs, output_attentions, use_cache) - - return custom_forward - - layer_outputs = torch.utils.checkpoint.checkpoint( - create_custom_forward(decoder_layer), - hidden_states, - attention_mask, - encoder_hidden_states, - encoder_attention_mask, - head_mask[idx] if head_mask is not None else None, - cross_attn_head_mask[idx] if cross_attn_head_mask is not None else None, - None, - ) - else: - if breakmodel: - device = primary_device if i < ram_blocks else bisect.bisect_right(cumulative_gpu_blocks, i - ram_blocks) - layer_outputs = decoder_layer( - hidden_states.to(device) if breakmodel and hidden_states is not None else hidden_states, - attention_mask=attention_mask.to(device) if breakmodel and attention_mask is not None else attention_mask, - encoder_hidden_states=encoder_hidden_states.to(device) if breakmodel and encoder_hidden_states is not None else encoder_hidden_states, - encoder_attention_mask=encoder_attention_mask.to(device) if breakmodel and encoder_attention_mask is not None else encoder_attention_mask, - layer_head_mask=((head_mask[idx].to(device) if breakmodel and head_mask[idx] is not None else head_mask[idx]) if head_mask is not None else None), - cross_attn_layer_head_mask=( - (cross_attn_head_mask[idx].to(device) if breakmodel and cross_attn_head_mask[idx] is not None else cross_attn_head_mask[idx]) if cross_attn_head_mask is not None else None - ), - past_key_value=tuple(v.to(device) for v in past_key_value if v is not None) if breakmodel and past_key_value is not None and i >= ram_blocks and len(past_key_value) and past_key_value[0].device.index != device else past_key_value, - output_attentions=output_attentions, - use_cache=use_cache, - ) - hidden_states = layer_outputs[0] - - if use_cache: - next_decoder_cache += (layer_outputs[3 if output_attentions else 1],) - - if output_attentions: - all_self_attns += (layer_outputs[1],) - - if encoder_hidden_states is not None: - all_cross_attentions += (layer_outputs[2],) - - if breakmodel: - if i in range(ram_blocks): - torch.cuda.synchronize() - torch.cuda.empty_cache() - - if breakmodel: - if ram_blocks: - del copystream - torch.cuda.empty_cache() - hidden_states = hidden_states.to(primary_device) - hidden_states = self.layer_norm(hidden_states) - if breakmodel: - hidden_states = hidden_states.to(primary_device) - - # add hidden states from the last decoder layer - if output_hidden_states: - all_hidden_states += (hidden_states,) - - next_cache = next_decoder_cache if use_cache else None - if not return_dict: - return tuple( - v - for v in [hidden_states, next_cache, all_hidden_states, all_self_attns, all_cross_attentions] - if v is not None - ) - return BaseModelOutputWithPastAndCrossAttentions( - last_hidden_state=hidden_states, - past_key_values=next_cache, - hidden_states=all_hidden_states, - attentions=all_self_attns, - cross_attentions=all_cross_attentions, - ) - - -def new_forward_opt( - self, - input_ids=None, - attention_mask=None, - head_mask=None, - past_key_values=None, - inputs_embeds=None, - use_cache=None, - output_attentions=None, - output_hidden_states=None, - return_dict=None, -): - assert len(gpu_blocks) <= torch.cuda.device_count() - assert sum(gpu_blocks) <= len(self.layers) - ram_blocks = len(self.layers) - sum(gpu_blocks) - cumulative_gpu_blocks = tuple(itertools.accumulate(gpu_blocks)) - - - output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions - output_hidden_states = ( - output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states - ) - use_cache = use_cache if use_cache is not None else self.config.use_cache - - return_dict = return_dict if return_dict is not None else self.config.use_return_dict - - # retrieve input_ids and inputs_embeds - if input_ids is not None and inputs_embeds is not None: - raise ValueError("You cannot specify both decoder_input_ids and decoder_inputs_embeds at the same time") - elif input_ids is not None: - input_shape = input_ids.size() - input_ids = input_ids.view(-1, input_shape[-1]) - elif inputs_embeds is not None: - input_shape = inputs_embeds.size()[:-1] - else: - raise ValueError("You have to specify either decoder_input_ids or decoder_inputs_embeds") - - past_key_values_length = past_key_values[0][0].shape[2] if past_key_values is not None else 0 - - if inputs_embeds is None: - if breakmodel: - input_ids = input_ids.to(primary_device) - inputs_embeds = self.embed_tokens(input_ids) - - # embed positions - if breakmodel: - inputs_embeds = inputs_embeds.to(primary_device) - if attention_mask is None: - attention_mask = torch.ones(inputs_embeds.shape[:2], dtype=torch.bool, device=inputs_embeds.device) - - positions = self.embed_positions(attention_mask)[:, past_key_values_length:, :] - if breakmodel: - positions = positions.to(primary_device) - - attention_mask = self._prepare_decoder_attention_mask( - attention_mask, input_shape, inputs_embeds, past_key_values_length - ) - - if self.project_in is not None: - inputs_embeds = self.project_in(inputs_embeds) - - hidden_states = inputs_embeds + positions - - hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training) - - # decoder layers - all_hidden_states = () if output_hidden_states else None - all_self_attns = () if output_attentions else None - next_decoder_cache = () if use_cache else None - - if breakmodel and ram_blocks: - copystream = torch.cuda.Stream(device=primary_device, priority=-1) - - # check if head_mask has a correct number of layers specified if desired - for attn_mask, mask_name in zip([head_mask], ["head_mask"]): - if attn_mask is not None: - if attn_mask.size()[0] != (len(self.layers)): - raise ValueError( - f"The `{mask_name}` should be specified for {len(self.layers)} layers, but it is for" - f" {head_mask.size()[0]}." - ) - - for idx, decoder_layer in enumerate(self.layers): - i = idx - if breakmodel: - if i in range(ram_blocks): - index1 = (i+1)%ram_blocks - for param1,param2 in zip(self.layers[index1].parameters(),self.layers[(i-1)%ram_blocks].parameters()): - param1.data = param2.data - for param1,param2 in zip(self.layers[index1].parameters(),self.extrastorage[index1].parameters()): - with torch.cuda.stream(copystream): - torch.cuda.comm.broadcast(param2.data,out = [param1.data]) - - # add LayerDrop (see https://arxiv.org/abs/1909.11556 for description) - if output_hidden_states: - all_hidden_states += (hidden_states,) - dropout_probability = random.uniform(0, 1) - if self.training and (dropout_probability < self.layerdrop): - continue - - past_key_value = past_key_values[idx] if past_key_values is not None else None - - if self.gradient_checkpointing and self.training: - - if use_cache: - logger.warning( - "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..." - ) - use_cache = False - - def create_custom_forward(module): - def custom_forward(*inputs): - # None for past_key_value - return module(*inputs, output_attentions, None) - - return custom_forward - - layer_outputs = torch.utils.checkpoint.checkpoint( - create_custom_forward(decoder_layer), - hidden_states, - attention_mask, - head_mask[idx] if head_mask is not None else None, - None, - ) - else: - if breakmodel: - device = primary_device if i < ram_blocks else bisect.bisect_right(cumulative_gpu_blocks, i - ram_blocks) - layer_outputs = decoder_layer( - hidden_states.to(device) if breakmodel and hidden_states is not None else hidden_states, - attention_mask=attention_mask.to(device) if breakmodel and attention_mask is not None else attention_mask, - layer_head_mask=((head_mask[idx].to(device) if breakmodel and head_mask[idx] is not None else head_mask[idx]) if head_mask is not None else None), - past_key_value=tuple(v.to(device) for v in past_key_value if v is not None) if breakmodel and past_key_value is not None and i >= ram_blocks and len(past_key_value) and past_key_value[0].device.index != device else past_key_value, - output_attentions=output_attentions, - use_cache=use_cache, - ) - - hidden_states = layer_outputs[0] - - if use_cache: - next_decoder_cache += (layer_outputs[2 if output_attentions else 1],) - - if output_attentions: - all_self_attns += (layer_outputs[1],) - - if breakmodel: - if i in range(ram_blocks): - torch.cuda.synchronize() - torch.cuda.empty_cache() - - if breakmodel: - if ram_blocks: - del copystream - torch.cuda.empty_cache() - hidden_states = hidden_states.to(primary_device) - if self.project_out is not None: - hidden_states = self.project_out(hidden_states) - if breakmodel: - hidden_states = hidden_states.to(primary_device) - - # add hidden states from the last decoder layer - if output_hidden_states: - all_hidden_states += (hidden_states,) - - next_cache = next_decoder_cache if use_cache else None - if not return_dict: - return tuple(v for v in [hidden_states, next_cache, all_hidden_states, all_self_attns] if v is not None) - return BaseModelOutputWithPast( - last_hidden_state=hidden_states, - past_key_values=next_cache, - hidden_states=all_hidden_states, - attentions=all_self_attns, - ) diff --git a/modeling/inference_models/generic_hf_torch/class.py b/modeling/inference_models/generic_hf_torch/class.py index fd4c2a1a..96d5f0c4 100644 --- a/modeling/inference_models/generic_hf_torch/class.py +++ b/modeling/inference_models/generic_hf_torch/class.py @@ -13,12 +13,6 @@ import modeling.lazy_loader as lazy_loader import koboldai_settings from logger import logger -try: - import breakmodel -except ModuleNotFoundError as e: - # Breakmodel is only expected to work on GPU - if not utils.koboldai_vars.use_colab_tpu: - raise e from modeling.inference_models.hf_torch import HFTorchInferenceModel @@ -70,14 +64,6 @@ class model_backend(HFTorchInferenceModel): # If we're using torch_lazy_loader, we need to get breakmodel config # early so that it knows where to load the individual model tensors logger.debug("lazy_load: {} hascuda: {} breakmodel: {} nobreakmode: {}".format(self.lazy_load, utils.koboldai_vars.hascuda, self.breakmodel, self.nobreakmodel)) - if ( - self.lazy_load - and utils.koboldai_vars.hascuda - and self.breakmodel - and not self.nobreakmodel - ): - logger.debug("loading breakmodel") - self.breakmodel_device_config(self.model_config) if self.lazy_load: # If we're using lazy loader, we need to figure out what the model's hidden layers are called @@ -141,7 +127,7 @@ class model_backend(HFTorchInferenceModel): self.get_local_model_path(ignore_existance=True) ) - if utils.koboldai_vars.fp32_model and not breakmodel.disk_blocks: + if utils.koboldai_vars.fp32_model: # Use save_pretrained to convert fp32 models to fp16, # unless we are using disk cache because save_pretrained # is not supported in that case @@ -247,27 +233,6 @@ class model_backend(HFTorchInferenceModel): shutil.rmtree("cache/") self.patch_embedding() - - - if utils.koboldai_vars.hascuda: - if self.usegpu: - # Use just VRAM - self.model = self.model.half().to(utils.koboldai_vars.gpu_device) - elif self.breakmodel: - # Use both RAM and VRAM (breakmodel) - if not self.lazy_load: - self.breakmodel_device_config(self.model.config) - self._move_to_devices() - elif breakmodel.disk_blocks > 0: - # Use disk - self._move_to_devices() - else: - # Use CPU - self.model = self.model.to("cpu").float() - elif breakmodel.disk_blocks > 0: - self._move_to_devices() - else: - self.model = self.model.to("cpu").float() self.model.kai_model = self diff --git a/modeling/inference_models/hf.py b/modeling/inference_models/hf.py index 4226d1b1..881ca604 100644 --- a/modeling/inference_models/hf.py +++ b/modeling/inference_models/hf.py @@ -157,7 +157,6 @@ class HFInferenceModel(InferenceModel): def set_input_parameters(self, parameters): if self.hf_torch and hasattr(self, "get_model_type") and self.get_model_type() != "gpt2": - import breakmodel layer_count = self.model_config["n_layer"] if isinstance(self.model_config, dict) else self.model_config.num_layers if hasattr(self.model_config, "num_layers") else self.model_config.n_layer if hasattr(self.model_config, "n_layer") else self.model_config.num_hidden_layers if hasattr(self.model_config, 'num_hidden_layers') else None if layer_count is not None and layer_count >= 0 and not self.nobreakmodel: gpu_count = torch.cuda.device_count() @@ -176,9 +175,8 @@ class HFInferenceModel(InferenceModel): self.disk_layers = parameters['Disk_Layers'] if 'Disk_Layers' in parameters else 0 if isinstance(self.disk_layers, str): self.disk_layers = int(self.disk_layers) if self.disk_layers.isnumeric() else 0 - breakmodel.gpu_blocks = layers - breakmodel.disk_blocks = self.disk_layers - self.usegpu = self.cpu_layers == 0 and breakmodel.disk_blocks == 0 and sum(self.layers)-self.layers[0] == 0 + print("TODO: Allow config") + # self.usegpu = self.cpu_layers == 0 and breakmodel.disk_blocks == 0 and sum(self.layers)-self.layers[0] == 0 self.model_type = self.get_model_type() self.breakmodel = ((self.model_type != 'gpt2') or self.model_type in ("gpt_neo", "gptj", "xglm", "opt")) and not self.nobreakmodel self.lazy_load = True diff --git a/modeling/inference_models/hf_torch.py b/modeling/inference_models/hf_torch.py index 2f575e73..abc3c54c 100644 --- a/modeling/inference_models/hf_torch.py +++ b/modeling/inference_models/hf_torch.py @@ -9,6 +9,7 @@ import functools import itertools import traceback import contextlib +from accelerate.utils.modeling import infer_auto_device_map, load_checkpoint_in_model from tqdm.auto import tqdm from typing import Dict, List, Optional, Union @@ -40,7 +41,6 @@ from modeling.inference_model import ( ) try: - import breakmodel import accelerate.utils except ModuleNotFoundError as e: if not utils.koboldai_vars.use_colab_tpu: @@ -125,17 +125,6 @@ class HFTorchInferenceModel(HFInferenceModel): else: return "Unknown" - def get_auxilary_device(self): - """Get device auxilary tensors like inputs should be stored on.""" - - # NOTE: TPU isn't a torch device, so TPU stuff gets sent to CPU. - if utils.koboldai_vars.hascuda and self.usegpu: - return utils.koboldai_vars.gpu_device - elif utils.koboldai_vars.hascuda and self.breakmodel: - import breakmodel - return breakmodel.primary_device - return "cpu" - def _post_load(m_self) -> None: if not utils.koboldai_vars.model_type: @@ -237,7 +226,7 @@ class HFTorchInferenceModel(HFInferenceModel): else: gen_in = prompt_tokens - device = self.get_auxilary_device() + device = utils.get_auxilary_device() gen_in = gen_in.to(device) additional_bad_words_ids = [self.tokenizer.encode("\n")] if single_line else [] @@ -254,8 +243,7 @@ class HFTorchInferenceModel(HFInferenceModel): len(prompt_tokens) + max_new, utils.koboldai_vars.max_length ), repetition_penalty=1.0, - bad_words_ids=self.badwordsids - + additional_bad_words_ids, + bad_words_ids=self.badwordsids + additional_bad_words_ids, use_cache=True, num_return_sequences=batch_count, ) @@ -286,7 +274,27 @@ class HFTorchInferenceModel(HFInferenceModel): # Try to determine model type from either AutoModel or falling back to legacy try: - return AutoModelForCausalLM.from_pretrained(location, **tf_kwargs) + model = AutoModelForCausalLM.from_config(self.model_config) + + # load_checkpoint_in_model( + # model.model, + # location, + # device_map=device_map + # offload_folder="accelerate-disk-cache", + # dtype="float16", + # offload_state_dict=True + # ) + # model.tie_weights() + + device_map = infer_auto_device_map( + model, + max_memory={0: "10GiB", 1: "7GiB", "cpu": "15GiB"}, + no_split_module_classes=["GPTJBlock"], + ) + + return AutoModelForCausalLM.from_pretrained( + location, device_map=device_map + ) # , **tf_kwargs) except Exception as e: traceback_string = traceback.format_exc().lower() @@ -325,49 +333,6 @@ class HFTorchInferenceModel(HFInferenceModel): return True - def _move_to_devices(self) -> None: - for key, value in self.model.state_dict().items(): - target_dtype = ( - torch.float32 if breakmodel.primary_device == "cpu" else torch.float16 - ) - if value.dtype is not target_dtype: - accelerate.utils.set_module_tensor_to_device( - self.model, - tensor_name=key, - device=torch.device(value.device), - value=value, - dtype=target_dtype, - ) - - disk_blocks = breakmodel.disk_blocks - gpu_blocks = breakmodel.gpu_blocks - ram_blocks = len(utils.layers_module_names) - sum(gpu_blocks) - cumulative_gpu_blocks = tuple(itertools.accumulate(gpu_blocks)) - device_map = {} - - for name in utils.layers_module_names: - layer = int(name.rsplit(".", 1)[1]) - device = ( - ("disk" if layer < disk_blocks else "cpu") - if layer < ram_blocks - else bisect.bisect_right(cumulative_gpu_blocks, layer - ram_blocks) - ) - device_map[name] = device - - for name in utils.get_missing_module_names(self.model, list(device_map.keys())): - device_map[name] = breakmodel.primary_device - - breakmodel.dispatch_model_ex( - self.model, - device_map, - main_device=breakmodel.primary_device, - offload_buffers=True, - offload_dir="accelerate-disk-cache", - ) - - gc.collect() - return - # Function to patch transformers to use our soft prompt def patch_embedding(self) -> None: if getattr(Embedding, "_koboldai_patch_causallm_model", None): @@ -413,11 +378,10 @@ class HFTorchInferenceModel(HFInferenceModel): if not self.lazy_load: return - - disk_blocks = breakmodel.disk_blocks - gpu_blocks = breakmodel.gpu_blocks - ram_blocks = ram_blocks = n_layers - sum(gpu_blocks) - cumulative_gpu_blocks = tuple(itertools.accumulate(gpu_blocks)) + # disk_blocks = breakmodel.disk_blocks + # gpu_blocks = breakmodel.gpu_blocks + # ram_blocks = ram_blocks = n_layers - sum(gpu_blocks) + # cumulative_gpu_blocks = tuple(itertools.accumulate(gpu_blocks)) def lazy_load_callback( model_dict: Dict[str, Union[lazy_loader.LazyTensor, torch.Tensor]], @@ -428,6 +392,7 @@ class HFTorchInferenceModel(HFInferenceModel): if lazy_load_callback.nested: return lazy_load_callback.nested = True + return device_map: Dict[str, Union[str, int]] = {} @@ -458,8 +423,7 @@ class HFTorchInferenceModel(HFInferenceModel): utils.koboldai_vars.gpu_device if utils.koboldai_vars.hascuda and self.usegpu else "cpu" - if not utils.koboldai_vars.hascuda - or not self.breakmodel + if not utils.koboldai_vars.hascuda or not self.breakmodel else breakmodel.primary_device ) else: @@ -479,8 +443,7 @@ class HFTorchInferenceModel(HFInferenceModel): else "disk" if layer < disk_blocks and layer < ram_blocks else "cpu" - if not utils.koboldai_vars.hascuda - or not self.breakmodel + if not utils.koboldai_vars.hascuda or not self.breakmodel else "shared" if layer < ram_blocks else bisect.bisect_right( @@ -519,7 +482,7 @@ class HFTorchInferenceModel(HFInferenceModel): total=num_tensors, desc="Loading model tensors", file=utils.UIProgressBarFile(), - position=1 + position=1, ) if not is_safetensors: @@ -550,7 +513,7 @@ class HFTorchInferenceModel(HFInferenceModel): f.close() ziproot = z.namelist()[0].split("/")[0] f = z.open(f"{ziproot}/data/{storage_key}") - + current_offset = 0 if current_offset != model_dict[key].seek_offset: f.read(model_dict[key].seek_offset - current_offset) @@ -574,7 +537,7 @@ class HFTorchInferenceModel(HFInferenceModel): ) ) # print(f"Transferring <{key}> to {f'({device.upper()})' if isinstance(device, str) else '[device ' + str(device) + ']'} ... ", end="", flush=True) - #logger.debug(f"Transferring <{key}> to {f'({device.upper()})' if isinstance(device, str) else '[device ' + str(device) + ']'} ... ") + # logger.debug(f"Transferring <{key}> to {f'({device.upper()})' if isinstance(device, str) else '[device ' + str(device) + ']'} ... ") model_dict[key] = model_dict[key].materialize( f, map_location="cpu" ) @@ -584,10 +547,7 @@ class HFTorchInferenceModel(HFInferenceModel): convert_to_float16 and breakmodel.primary_device != "cpu" and utils.koboldai_vars.hascuda - and ( - self.breakmodel - or self.usegpu - ) + and (self.breakmodel or self.usegpu) and model_dict[key].dtype is torch.float32 ): model_dict[key] = model_dict[key].to(torch.float16) @@ -630,15 +590,11 @@ class HFTorchInferenceModel(HFInferenceModel): convert_to_float16 and breakmodel.primary_device != "cpu" and utils.koboldai_vars.hascuda - and ( - self.breakmodel - or self.usegpu - ) + and (self.breakmodel or self.usegpu) ): dtype = torch.float16 if breakmodel.primary_device == "cpu" or ( - not self.usegpu - and not self.breakmodel + not self.usegpu and not self.breakmodel ): dtype = torch.float32 if ( @@ -693,10 +649,7 @@ class HFTorchInferenceModel(HFInferenceModel): convert_to_float16 and breakmodel.primary_device != "cpu" and utils.koboldai_vars.hascuda - and ( - self.breakmodel - or self.usegpu - ) + and (self.breakmodel or self.usegpu) and model_dict[key].dtype is torch.float32 ): model_dict[key] = model_dict[key].to(torch.float16) @@ -741,15 +694,11 @@ class HFTorchInferenceModel(HFInferenceModel): convert_to_float16 and breakmodel.primary_device != "cpu" and utils.koboldai_vars.hascuda - and ( - self.breakmodel - or self.usegpu - ) + and (self.breakmodel or self.usegpu) ): dtype = torch.float16 if breakmodel.primary_device == "cpu" or ( - not self.usegpu - and not self.breakmodel + not self.usegpu and not self.breakmodel ): dtype = torch.float32 if ( @@ -793,6 +742,7 @@ class HFTorchInferenceModel(HFInferenceModel): yield False def breakmodel_device_list(self, n_layers, primary=None, selected=None): + return # TODO: Find a better place for this or rework this device_count = torch.cuda.device_count() @@ -824,6 +774,7 @@ class HFTorchInferenceModel(HFInferenceModel): def breakmodel_device_config(self, config): # TODO: Find a better place for this or rework this + return global breakmodel, generator import breakmodel @@ -840,7 +791,7 @@ class HFTorchInferenceModel(HFInferenceModel): logger.info("Breakmodel not specified, assuming GPU 0") breakmodel.gpu_blocks = [n_layers] n_layers = 0 - + else: s = n_layers for i in range(len(breakmodel.gpu_blocks)): @@ -857,8 +808,14 @@ class HFTorchInferenceModel(HFInferenceModel): logger.init_ok("Final device configuration:", status="Info") self.breakmodel_device_list(n_layers, primary=breakmodel.primary_device) - with open("settings/{}.breakmodel".format(self.model_name.replace("/", "_")), "w") as file: - file.write("{}\n{}".format(",".join(map(str, breakmodel.gpu_blocks)), breakmodel.disk_blocks)) + with open( + "settings/{}.breakmodel".format(self.model_name.replace("/", "_")), "w" + ) as file: + file.write( + "{}\n{}".format( + ",".join(map(str, breakmodel.gpu_blocks)), breakmodel.disk_blocks + ) + ) # If all layers are on the same device, use the old GPU generation mode while len(breakmodel.gpu_blocks) and breakmodel.gpu_blocks[-1] == 0: @@ -876,9 +833,6 @@ class HFTorchInferenceModel(HFInferenceModel): if not breakmodel.gpu_blocks: logger.warning("Nothing assigned to a GPU, reverting to CPU only mode") - import breakmodel - - breakmodel.primary_device = "cpu" self.breakmodel = False self.usegpu = False return diff --git a/modeling/lazy_loader.py b/modeling/lazy_loader.py index 5a27d549..a0f67d4a 100644 --- a/modeling/lazy_loader.py +++ b/modeling/lazy_loader.py @@ -101,6 +101,7 @@ class TorchLazyTensor(LazyTensor): stride: Optional[Tuple[int, ...]] = None, requires_grad=False, backward_hooks: Any = None, + file_handle: Any = None ): self.storage_type = storage_type self.key = key @@ -111,6 +112,7 @@ class TorchLazyTensor(LazyTensor): self.stride = stride self.requires_grad = requires_grad self.backward_hooks = backward_hooks + self.file_handle = file_handle def __view(self, f: Callable): return f"{type(self).__name__}(storage_type={f(self.storage_type)}, key={f(self.key)}, location={f(self.location)}, dtype={f(self.dtype)}, seek_offset={f(self.seek_offset)}, shape={f(self.shape)}, stride={f(self.stride)}, requires_grad={f(self.requires_grad)}, backward_hooks={f(self.backward_hooks)})" @@ -120,11 +122,13 @@ class TorchLazyTensor(LazyTensor): def materialize( self, - checkpoint: Union[zipfile.ZipFile, zipfile.ZipExtFile], + checkpoint: Union[zipfile.ZipFile, zipfile.ZipExtFile] = None, map_location=None, no_grad=True, filename="pytorch_model.bin", ) -> torch.Tensor: + checkpoint = checkpoint or self.file_handle + filename = os.path.basename(os.path.normpath(filename)).split(".")[0] size = reduce(lambda x, y: x * y, self.shape, 1) dtype = self.dtype @@ -237,6 +241,8 @@ class _LazyUnpickler(RestrictedUnpickler): lazy_loaded_storages: Dict[str, LazyTensor] def __init__(self, *args, **kwargs): + # print(args, kwargs) + self.file_handle = args[0] self.lazy_loaded_storages = {} return super().__init__(*args, **kwargs) @@ -247,7 +253,7 @@ class _LazyUnpickler(RestrictedUnpickler): typename == "storage" ), f"Unknown typename for persistent_load, expected 'storage' but got '{typename}'" storage_type, key, location, _ = saved_id[1:] - return TorchLazyTensor(storage_type, key, location) + return TorchLazyTensor(storage_type, key, location, file_handle=self.file_handle) def load(self, *args, **kwargs): retval = super().load(*args, **kwargs) diff --git a/modeling/patches.py b/modeling/patches.py index f6b1ff0a..c8c070f8 100644 --- a/modeling/patches.py +++ b/modeling/patches.py @@ -10,6 +10,7 @@ from transformers import ( PreTrainedModel, modeling_utils, ) +from modeling.lazy_loader import LazyTensor import utils @@ -125,6 +126,173 @@ def patch_transformers_generation() -> None: transformers.generation.logits_process.NoBadWordsLogitsProcessor.__init__ = new_init +CURRENT_CHECKPOINT = None +def patch_transformers_for_lazyload() -> None: + import torch + import inspect + from accelerate.utils import set_module_tensor_to_device, offload_weight + + def _load_state_dict_into_meta_model( + model, + state_dict, + loaded_state_dict_keys, # left for now but could be removed, see below + start_prefix, + expected_keys, + device_map=None, + offload_folder=None, + offload_index=None, + state_dict_folder=None, + state_dict_index=None, + dtype=None, + load_in_8bit=False, + is_safetensors=False, + keep_in_fp32_modules=None, + ): + """ + This is somewhat similar to `_load_state_dict_into_model`, but deals with a model that has some or all of its + params on a `meta` device. It replaces the model params with the data from the `state_dict`, while moving the + params back to the normal device, but only for `loaded_state_dict_keys`. + + `start_prefix` is used for models which insert their name into model keys, e.g. `bert` in + `bert.pooler.dense.weight` + + """ + + print("DEVMAP", device_map) + + # XXX: remaining features to implement to be fully compatible with _load_state_dict_into_model + # - deepspeed zero 3 support + # - need to copy metadata if any - see _load_state_dict_into_model + # - handling error_msgs - mimicking the error handling in module._load_from_state_dict() + # - Is there a situation where some keys aren't in `loaded_state_dict_keys` and in which case + # they won't get loaded. + + if load_in_8bit: + from .utils.bitsandbytes import set_module_8bit_tensor_to_device + + error_msgs = [] + + old_keys = [] + new_keys = [] + for key in state_dict.keys(): + new_key = None + if "gamma" in key: + new_key = key.replace("gamma", "weight") + if "beta" in key: + new_key = key.replace("beta", "bias") + if new_key: + old_keys.append(key) + new_keys.append(new_key) + for old_key, new_key in zip(old_keys, new_keys): + state_dict[new_key] = state_dict.pop(old_key) + + for param_name, param in state_dict.items(): + + # BEGIN PATCH + if isinstance(param, LazyTensor): + print("Materializing", param_name) + param = param.materialize() + # END PATCH + + # First part of the test is always true as load_state_dict_keys always contains state_dict keys. + if ( + param_name not in loaded_state_dict_keys + or param_name not in expected_keys + ): + continue + + if param_name.startswith(start_prefix): + param_name = param_name[len(start_prefix) :] + + module_name = param_name + set_module_kwargs = {} + + # We convert floating dtypes to the `dtype` passed. We want to keep the buffers/params + # in int/uint/bool and not cast them. + if dtype is not None and torch.is_floating_point(param): + if ( + keep_in_fp32_modules is not None + and any( + module_to_keep_in_fp32 in param_name + for module_to_keep_in_fp32 in keep_in_fp32_modules + ) + and dtype == torch.float16 + ): + param = param.to(torch.float32) + + # For backward compatibility with older versions of `accelerate` + # TODO: @sgugger replace this check with version check at the next `accelerate` release + if "dtype" in list( + inspect.signature(set_module_tensor_to_device).parameters + ): + set_module_kwargs["dtype"] = torch.float32 + else: + param = param.to(dtype) + + # For compatibility with PyTorch load_state_dict which converts state dict dtype to existing dtype in model + if dtype is None: + old_param = model + splits = param_name.split(".") + for split in splits: + old_param = getattr(old_param, split) + if old_param is None: + break + + if old_param is not None: + param = param.to(old_param.dtype) + + set_module_kwargs["value"] = param + + if device_map is None: + param_device = "cpu" + else: + # find next higher level module that is defined in device_map: + # bert.lm_head.weight -> bert.lm_head -> bert -> '' + while len(module_name) > 0 and module_name not in device_map: + module_name = ".".join(module_name.split(".")[:-1]) + if module_name == "" and "" not in device_map: + # TODO: group all errors and raise at the end. + raise ValueError(f"{param_name} doesn't have any device set.") + param_device = device_map[module_name] + if param_device == "disk": + if not is_safetensors: + offload_index = offload_weight( + param, param_name, offload_folder, offload_index + ) + elif param_device == "cpu" and state_dict_index is not None: + state_dict_index = offload_weight( + param, param_name, state_dict_folder, state_dict_index + ) + elif not load_in_8bit: + # For backward compatibility with older versions of `accelerate` + set_module_tensor_to_device( + model, param_name, param_device, **set_module_kwargs + ) + else: + if ( + param.dtype == torch.int8 + and param_name.replace("weight", "SCB") in state_dict.keys() + ): + fp16_statistics = state_dict[param_name.replace("weight", "SCB")] + else: + fp16_statistics = None + + if "SCB" not in param_name: + set_module_8bit_tensor_to_device( + model, + param_name, + param_device, + value=param, + fp16_statistics=fp16_statistics, + ) + + return error_msgs, offload_index, state_dict_index + + transformers.modeling_utils._load_state_dict_into_meta_model = ( + _load_state_dict_into_meta_model + ) + + def patch_transformers() -> None: patch_transformers_download() patch_transformers_loader() diff --git a/utils.py b/utils.py index 13ebb6a3..863bda2f 100644 --- a/utils.py +++ b/utils.py @@ -656,9 +656,9 @@ def get_auxilary_device(): # NOTE: TPU isn't a torch device, so TPU stuff gets sent to CPU. if koboldai_vars.hascuda and koboldai_vars.usegpu: return koboldai_vars.gpu_device - elif koboldai_vars.hascuda and koboldai_vars.breakmodel: - import breakmodel - return breakmodel.primary_device + elif koboldai_vars.hascuda: + # TODO: Primary device + return "cuda" return "cpu" #==================================================================#