Activation offloading for fullfinetuning + fix tied embedding (#1847)

Co-authored-by: Felipe Mello <[email protected]>

Author: felipemello1

recipes/configs/code_llama2/7B_full_low_memory.yaml

@@ -69,6 +69,7 @@ device: cuda
 
 # Memory management
 enable_activation_checkpointing: True
+enable_activation_offloading: True  # True reduces memory
 dtype: bf16
 
 # Logging

recipes/configs/code_llama2/7B_lora_single_device.yaml

@@ -77,7 +77,7 @@ device: cuda
 
 # Memory management
 enable_activation_checkpointing: True
-enable_activation_offloading: False
+enable_activation_offloading: False  # True reduces memory
 dtype: bf16
 
 # Logging

recipes/configs/code_llama2/7B_qlora_single_device.yaml

@@ -76,7 +76,7 @@ device: cuda
 
 # Memory management
 enable_activation_checkpointing: True
-enable_activation_offloading: False
+enable_activation_offloading: False  # True reduces memory
 dtype: bf16
 
 # Logging

recipes/configs/dev/8B_full_experimental.yaml

@@ -65,6 +65,7 @@ device: cuda
 
 # Memory management
 enable_activation_checkpointing: False
+enable_activation_offloading: False  # True reduces memory
 ac_mode: 'selective'  # ['selective', 'full']
 ac_option: 2 # [int] = ac every positive int layer
 memory_efficient_fsdp_wrap: False

recipes/configs/gemma/2B_full.yaml

@@ -62,6 +62,7 @@ device: cuda
 
 # Memory management
 enable_activation_checkpointing: True
+enable_activation_offloading: False  # True reduces memory
 
 # Reduced precision
 dtype: bf16

recipes/configs/gemma/2B_lora.yaml

@@ -74,6 +74,7 @@ device: cuda
 
 # Memory management
 enable_activation_checkpointing: True
+enable_activation_offloading: False  # True reduces memory
 
 # Reduced precision
 dtype: bf16

recipes/configs/gemma/2B_lora_single_device.yaml

@@ -73,7 +73,7 @@ device: cuda
 
 # Memory management
 enable_activation_checkpointing: True
-enable_activation_offloading: False
+enable_activation_offloading: False  # True reduces memory
 
 # Reduced precision
 dtype: bf16

recipes/configs/gemma/2B_qlora_single_device.yaml

@@ -73,7 +73,7 @@ device: cuda
 
 # Memory management
 enable_activation_checkpointing: True
-enable_activation_offloading: False
+enable_activation_offloading: False  # True reduces memory
 
 # Reduced precision
 dtype: bf16

recipes/configs/gemma/7B_full.yaml

@@ -64,6 +64,7 @@ device: cuda
 
 # Memory management
 enable_activation_checkpointing: True
+enable_activation_offloading: False  # True reduces memory
 
 # Reduced precision
 dtype: bf16

recipes/configs/gemma/7B_lora.yaml

@@ -76,6 +76,7 @@ device: cuda
 
 # Memory management
 enable_activation_checkpointing: True
+enable_activation_offloading: False  # True reduces memory
 
 # Reduced precision
 dtype: bf16