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

Merge pull request #127 from huggingface/smangrul/make-activation-checkpointing-optional

Browse Source 
fixes `prepare_for_int8_training`
This commit is contained in:
Sourab Mangrulkar
2023-02-25 02:13:04 +05:30
committed by GitHub
parent a3537160dc 83d6d55d4b
commit e19ee681ac
1 changed files with 12 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/peft/utils/other.py
+12 -10
View File
@@ -30,7 +30,9 @@ 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, layer_norm_names=["layer_norm"]
):
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
@@ -48,20 +50,20 @@ def prepare_model_for_int8_training(model, output_embedding_layer_name="lm_head"
if loaded_in_8bit:
# cast layer norm in fp32 for stability for 8bit models
if param.ndim == 1 and "layer_norm" in name:
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)
# 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()