File size: 2,921 Bytes
668789a |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 |
import torch
from torch import nn
from .lora import FrozenBNBLinear, FrozenBNBEmbedding
from .config import GPTJLoraConfig
import transformers
def add_adapters(model, adapter_dim=16):
assert adapter_dim > 0
for module in model.modules():
if isinstance(module, FrozenBNBLinear):
module.adapter = nn.Sequential(
nn.Linear(module.in_features, adapter_dim, bias=False),
nn.Linear(adapter_dim, module.out_features, bias=False),
)
nn.init.zeros_(module.adapter[1].weight)
elif isinstance(module, FrozenBNBEmbedding):
module.adapter = nn.Sequential(
nn.Embedding(module.num_embeddings, adapter_dim),
nn.Linear(adapter_dim, module.embedding_dim, bias=False),
)
nn.init.zeros_(module.adapter[1].weight)
def convert_to_int8(model):
"""Convert linear and embedding modules to 8-bit with optional adapters"""
for module in list(model.modules()):
for name, child in module.named_children():
if isinstance(child, nn.Linear):
setattr(
module,
name,
FrozenBNBLinear(
weight=torch.zeros(child.out_features, child.in_features, dtype=torch.uint8),
absmax=torch.zeros((child.weight.numel() - 1) // 4096 + 1),
code=torch.zeros(256),
bias=child.bias,
),
)
elif isinstance(child, nn.Embedding):
setattr(
module,
name,
FrozenBNBEmbedding(
weight=torch.zeros(child.num_embeddings, child.embedding_dim, dtype=torch.uint8),
absmax=torch.zeros((child.weight.numel() - 1) // 4096 + 1),
code=torch.zeros(256),
)
)
class GPTJLoraBlock(transformers.models.gptj.modeling_gptj.GPTJBlock):
config_class = GPTJLoraConfig
def __init__(self, config):
super().__init__(config)
self.config_class = GPTJLoraConfig
convert_to_int8(self.attn)
convert_to_int8(self.mlp)
class GPTJModel(transformers.models.gptj.modeling_gptj.GPTJModel):
config_class = GPTJLoraConfig
def __init__(self, config):
super().__init__(config)
self.config_class = GPTJLoraConfig
convert_to_int8(self)
class GPTJLoraForCausalLM(transformers.models.gptj.modeling_gptj.GPTJForCausalLM):
config_class = GPTJLoraConfig
def __init__(self, config):
super().__init__(config)
self.config_class = GPTJLoraConfig
convert_to_int8(self)
if config.add_apapters:
add_adapters(self)
transformers.models.gptj.modeling_gptj.GPTJBlock = GPTJLoraBlock # monkey-patch GPT-J
|