gradient_checkpointing_enable()

wukaixingxp · wukaixingxp · commit 50dff0b78e89 · 2024-09-24T12:22:47.000-07:00
diff --git a/src/llama_recipes/finetuning.py b/src/llama_recipes/finetuning.py
@@ -208,7 +208,7 @@ def main(**kwargs):
         )
         if fsdp_config.fsdp_activation_checkpointing:            
             model.enable_input_require_grads()
-            #model.gradient_checkpointing_enable()
+            model.gradient_checkpointing_enable()
             apply_fsdp_checkpointing(model)                      
     elif not train_config.quantization and not train_config.enable_fsdp:
         if is_xpu_available():

Original file line number	Diff line number	Diff line change
`@@ -208,7 +208,7 @@ def main(**kwargs):`
`208`	`208`	`)`
`209`	`209`	`if fsdp_config.fsdp_activation_checkpointing:`
`210`	`210`	`model.enable_input_require_grads()`
`211`		`- #model.gradient_checkpointing_enable()`
	`211`	`+ model.gradient_checkpointing_enable()`
`212`	`212`	`apply_fsdp_checkpointing(model)`
`213`	`213`	`elif not train_config.quantization and not train_config.enable_fsdp:`
`214`	`214`	`if is_xpu_available():`