Account for DoRA
This commit is contained in:
Daniel Han-Chen
parent
2e87755a5c
commit
cb0d937469
2 changed files with 61 additions and 12 deletions
|
|
@ -95,7 +95,7 @@ def prepare_model_for_kbit_training(
|
|||
|
||||
# Freeze all parameters except LoRA
|
||||
for name, param in model.named_parameters():
|
||||
if ".lora_A." in name or ".lora_B." in name:
|
||||
if ".lora_A." in name or ".lora_B." in name or ".lora_magnitude_vector" in name:
|
||||
param.requires_grad_(True)
|
||||
else:
|
||||
param.requires_grad_(False)
|
||||
|
|
|
|||
|
|
@ -511,6 +511,7 @@ def LlamaModel_fast_forward(
|
|||
|
||||
# Mormalized from Gemma
|
||||
IS_GEMMA = self.config.model_type == "gemma"
|
||||
train_embed_tokens = self.embed_tokens.weight.requires_grad
|
||||
|
||||
if IS_GEMMA:
|
||||
# Match Gemma exactly by casting to bfloat16 / float16
|
||||
|
|
@ -519,7 +520,7 @@ def LlamaModel_fast_forward(
|
|||
# & 2048**0.5 = 45.2500 in bfloat16, whilst 45.2548 in float32
|
||||
normalizer = torch.tensor(math_sqrt(self.config.hidden_size), dtype = inputs_embeds.dtype)
|
||||
|
||||
if self.embed_tokens.weight.requires_grad:
|
||||
if train_embed_tokens:
|
||||
# Careful we must not do an inplace op!
|
||||
inputs_embeds = inputs_embeds * normalizer
|
||||
else:
|
||||
|
|
@ -538,7 +539,8 @@ def LlamaModel_fast_forward(
|
|||
|
||||
# Fix up attention mask by setting elements to 0
|
||||
# Specifically for DPO
|
||||
if self._has_no_labels and (attention_mask is not None) and (past_key_values is None):
|
||||
if self._has_no_labels and (attention_mask is not None) and (past_key_values is None) and \
|
||||
(not train_embed_tokens):
|
||||
# Careful for inference the attention_mask is size (1, kv_seq_len)
|
||||
# Whilst the input_embeds is size (1, 1, 4096)
|
||||
inputs_requires_grad = inputs_embeds.requires_grad
|
||||
|
|
@ -1234,6 +1236,7 @@ class FastLlamaModel:
|
|||
random_state = 3407,
|
||||
max_seq_length = 2048, # not used anymore
|
||||
use_rslora = False,
|
||||
modules_to_save = None,
|
||||
init_lora_weights = True,
|
||||
loftq_config = {},
|
||||
**kwargs,
|
||||
|
|
@ -1320,15 +1323,45 @@ class FastLlamaModel:
|
|||
accepted_modules = frozenset(("q_proj", "k_proj", "v_proj", "o_proj",
|
||||
"gate_proj", "up_proj", "down_proj",),)
|
||||
model.config.update({"unsloth_version" : __version__})
|
||||
|
||||
train_lm_head = False
|
||||
train_embed_tokens = False
|
||||
final_modules = []
|
||||
for module in target_modules:
|
||||
assert(module in accepted_modules)
|
||||
if module == "lm_head":
|
||||
logger.warning_once(
|
||||
"Unsloth: `lm_head` should be placed in `modules_to_save` and not `target_modules`."\
|
||||
"We shall do it for you!"
|
||||
)
|
||||
train_lm_head = True
|
||||
|
||||
elif module == "embed_tokens":
|
||||
logger.warning_once(
|
||||
"Unsloth: `embed_tokens` should be placed in `modules_to_save` and not `target_modules`."\
|
||||
"We shall do it for you!"
|
||||
)
|
||||
train_embed_tokens = True
|
||||
|
||||
else:
|
||||
assert(module in accepted_modules)
|
||||
final_modules.append(module)
|
||||
pass
|
||||
|
||||
# Check modules_to_save
|
||||
if modules_to_save is not None:
|
||||
for module in modules_to_save:
|
||||
if module == "lm_head":
|
||||
train_lm_head = True
|
||||
elif module == "embed_tokens":
|
||||
train_embed_tokens = True
|
||||
pass
|
||||
pass
|
||||
|
||||
# Get LoRA
|
||||
arguments = dict(
|
||||
r = r,
|
||||
lora_alpha = lora_alpha,
|
||||
target_modules = target_modules,
|
||||
target_modules = final_modules,
|
||||
lora_dropout = lora_dropout,
|
||||
bias = bias,
|
||||
task_type = TaskType.CAUSAL_LM,
|
||||
|
|
@ -1336,6 +1369,7 @@ class FastLlamaModel:
|
|||
init_lora_weights = init_lora_weights,
|
||||
loftq_config = loftq_config,
|
||||
use_rslora = use_rslora,
|
||||
modules_to_save = modules_to_save,
|
||||
**kwargs,
|
||||
)
|
||||
if not SUPPORTS_LOFTQ: del arguments["loftq_config"]
|
||||
|
|
@ -1345,6 +1379,14 @@ class FastLlamaModel:
|
|||
model = _get_peft_model(model, lora_config)
|
||||
|
||||
model = FastLlamaModel.patch_peft_model(model, use_gradient_checkpointing)
|
||||
|
||||
# Now patch lm_head and embed_tokens
|
||||
if train_embed_tokens:
|
||||
model.model.model.embed_tokens.requires_grad_(True)
|
||||
if train_lm_head:
|
||||
model.model.lm_head.requires_grad_(True)
|
||||
pass
|
||||
|
||||
return model
|
||||
pass
|
||||
|
||||
|
|
@ -1435,9 +1477,12 @@ class FastLlamaModel:
|
|||
if hasattr(gate_proj, "lora_A") and \
|
||||
hasattr( up_proj, "lora_A") and \
|
||||
hasattr(down_proj, "lora_A") and \
|
||||
(gate_proj.base_layer if hasattr(gate_proj, "base_layer") else gate_proj).bias is None and \
|
||||
( up_proj.base_layer if hasattr( up_proj, "base_layer") else up_proj).bias is None and \
|
||||
(down_proj.base_layer if hasattr(down_proj, "base_layer") else down_proj).bias is None:
|
||||
((gate_proj.base_layer if hasattr(gate_proj, "base_layer") else gate_proj).bias is None) and \
|
||||
(( up_proj.base_layer if hasattr( up_proj, "base_layer") else up_proj).bias is None) and \
|
||||
((down_proj.base_layer if hasattr(down_proj, "base_layer") else down_proj).bias is None) and \
|
||||
(not hasattr(gate_proj, "lora_magnitude_vector")) and \
|
||||
(not hasattr( up_proj, "lora_magnitude_vector")) and \
|
||||
(not hasattr(down_proj, "lora_magnitude_vector")):
|
||||
|
||||
# https://stackoverflow.com/questions/50599045/python-replacing-a-function-within-a-class-of-a-module
|
||||
layer.mlp.forward = types.MethodType(apply_lora_mlp, layer.mlp)
|
||||
|
|
@ -1456,9 +1501,12 @@ class FastLlamaModel:
|
|||
if hasattr(q_proj, "lora_A") and \
|
||||
hasattr(k_proj, "lora_A") and \
|
||||
hasattr(v_proj, "lora_A") and \
|
||||
(q_proj.base_layer if hasattr(q_proj, "base_layer") else q_proj).bias is None and \
|
||||
(k_proj.base_layer if hasattr(k_proj, "base_layer") else k_proj).bias is None and \
|
||||
(v_proj.base_layer if hasattr(v_proj, "base_layer") else v_proj).bias is None:
|
||||
((q_proj.base_layer if hasattr(q_proj, "base_layer") else q_proj).bias is None) and \
|
||||
((k_proj.base_layer if hasattr(k_proj, "base_layer") else k_proj).bias is None) and \
|
||||
((v_proj.base_layer if hasattr(v_proj, "base_layer") else v_proj).bias is None) and \
|
||||
(not hasattr(q_proj, "lora_magnitude_vector")) and \
|
||||
(not hasattr(k_proj, "lora_magnitude_vector")) and \
|
||||
(not hasattr(v_proj, "lora_magnitude_vector")):
|
||||
|
||||
layer.self_attn.apply_qkv = apply_lora_qkv
|
||||
n_qkv += 1
|
||||
|
|
@ -1472,7 +1520,8 @@ class FastLlamaModel:
|
|||
# O attention patching
|
||||
o_proj = layer.self_attn.o_proj
|
||||
if hasattr(o_proj, "lora_A") and \
|
||||
(o_proj.base_layer if hasattr(o_proj, "base_layer") else o_proj).bias is None:
|
||||
((o_proj.base_layer if hasattr(o_proj, "base_layer") else o_proj).bias is None) and \
|
||||
(not hasattr(o_proj, "lora_magnitude_vector")):
|
||||
|
||||
layer.self_attn.apply_o = apply_lora_o
|
||||
n_o += 1
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue