Spaces:
Runtime error
Runtime error
| import torch | |
| from typing import TYPE_CHECKING, List, Optional | |
| from llmtuner.extras.constants import LAYERNORM_NAMES | |
| if TYPE_CHECKING: | |
| from transformers.modeling_utils import PreTrainedModel | |
| def find_all_linear_modules( | |
| model: "PreTrainedModel", | |
| quantization_bit: Optional[int] = None, | |
| output_layer_name: Optional[str] = "lm_head" | |
| ) -> List[str]: | |
| if quantization_bit is not None: | |
| import bitsandbytes as bnb | |
| linear_cls = bnb.nn.Linear4bit if quantization_bit == 4 else bnb.nn.Linear8bitLt | |
| else: | |
| linear_cls = torch.nn.Linear | |
| module_names = set() | |
| for name, module in model.named_modules(): | |
| if output_layer_name not in name and isinstance(module, linear_cls): | |
| module_names.add(name.split(".")[-1]) | |
| if output_layer_name in module_names: | |
| module_names.pop(output_layer_name) | |
| return list(module_names) | |
| def prepare_model_for_training( | |
| model: "PreTrainedModel", | |
| finetuning_type: str, | |
| output_layer_name: Optional[str] = "lm_head", | |
| use_gradient_checkpointing: Optional[bool] = True, | |
| layer_norm_names: Optional[List[str]] = LAYERNORM_NAMES | |
| ) -> "PreTrainedModel": | |
| r""" | |
| Includes: | |
| (1) cast the layernorm in fp32 | |
| (2) make output embedding layer require grads | |
| (3) upcast the lm_head to fp32 | |
| Inspired by: https://github.com/huggingface/peft/blob/v0.2.0/src/peft/utils/other.py#L33 | |
| """ | |
| for name, param in model.named_parameters(): | |
| if param.ndim == 1 and any(layer_norm_name in name for layer_norm_name in layer_norm_names): | |
| param.data = param.data.to(torch.float32) | |
| if use_gradient_checkpointing: | |
| if hasattr(model, "enable_input_require_grads"): | |
| model.enable_input_require_grads() | |
| else: | |
| def make_inputs_require_grad(module, input, output): | |
| output.requires_grad_(True) | |
| model.get_input_embeddings().register_forward_hook(make_inputs_require_grad) | |
| model.gradient_checkpointing_enable() | |
| model.config.use_cache = False # turn off when gradient checkpointing is enabled | |
| if finetuning_type != "full" and hasattr(model, output_layer_name): | |
| output_layer: torch.nn.Linear = getattr(model, output_layer_name) | |
| input_dtype = output_layer.weight.dtype | |
| class CastOutputToFloat(torch.nn.Sequential): | |
| def forward(self, x: torch.Tensor) -> torch.Tensor: | |
| return super().forward(x.to(input_dtype)).to(torch.float32) | |
| setattr(model, output_layer_name, CastOutputToFloat(output_layer)) | |
| return model | |