wassname/peft
1
0
Fork 0
mirror of https://github.com/wassname/peft.git synced 2026-09-09 11:28:32 +08:00
Code Issues Packages Projects Releases Wiki Activity

make gradient checkpointing optional when using PEFT+INT8

Browse Source
This commit is contained in:
Sourab Mangrulkar
2023-02-24 13:02:40 +05:30
parent a3537160dc
commit 7dfb472424
1 changed files with 9 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/peft/utils/other.py
+9 -9
View File
@@ -30,7 +30,7 @@ def bloom_model_postprocess_past_key_value(past_key_values):
return tuple(zip(keys, values))
def prepare_model_for_int8_training(model, output_embedding_layer_name="lm_head"):
def prepare_model_for_int8_training(model, output_embedding_layer_name="lm_head", use_gradient_checkpointing=True):
r"""
This method wrapps the entire protocol for preparing a model before running a training. This includes:
1- Cast the layernorm in fp32 2- making output embedding layer require grads 3- Add the upcasting of the lm
@@ -51,17 +51,17 @@ def prepare_model_for_int8_training(model, output_embedding_layer_name="lm_head"
if param.ndim == 1 and "layer_norm" in name:
param.data = param.data.to(torch.float32)
# For backward compatibility
if hasattr(model, "enable_input_require_grads"):
model.enable_input_require_grads()
else:
if loaded_in_8bit and use_gradient_checkpointing:
# For backward compatibility
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)
def make_inputs_require_grad(module, input, output):
output.requires_grad_(True)
model.get_input_embeddings().register_forward_hook(make_inputs_require_grad)
model.get_input_embeddings().register_forward_hook(make_inputs_require_grad)
if loaded_in_8bit:
# enable gradient checkpointing for memory efficiency
model.gradient_checkpointing_enable()