add fused = true to adam, except pagedAdam (#1575)

Co-authored-by: Felipe Mello <felipemello@fb.com>

recipes/configs/code_llama2/7B_lora_single_device.yaml

@@ -60,6 +60,7 @@ batch_size: 2
 gradient_accumulation_steps: 16
 optimizer:
   _component_: torch.optim.AdamW
+  fused: True
   weight_decay: 0.01
   lr: 3e-4
 lr_scheduler:

recipes/configs/code_llama2/7B_qlora_single_device.yaml

@@ -60,6 +60,7 @@ batch_size: 2
 gradient_accumulation_steps: 16
 optimizer:
   _component_: torch.optim.AdamW
+  fused: True
   weight_decay: 0.01
   lr: 3e-4
 lr_scheduler:

recipes/configs/dev/8B_full_experimental.yaml

@@ -52,8 +52,7 @@ epochs: 3
 optimizer:
   _component_: torch.optim.AdamW
   lr: 2e-5
-  foreach: False
-
+  fused: True
 loss:
   _component_: torchtune.modules.loss.CEWithChunkedOutputLoss
 max_steps_per_epoch: null

recipes/configs/gemma/2B_full.yaml

@@ -48,6 +48,7 @@ batch_size: 2
 epochs: 3
 optimizer:
   _component_: torch.optim.AdamW
+  fused: True
   lr: 2e-5
 loss:
   _component_: torchtune.modules.loss.CEWithChunkedOutputLoss

recipes/configs/gemma/2B_lora.yaml

@@ -51,6 +51,7 @@ save_adapter_weights_only: False
 
 optimizer:
   _component_: torch.optim.AdamW
+  fused: True
   lr: 2e-5
 
 lr_scheduler:

recipes/configs/gemma/2B_lora_single_device.yaml

@@ -50,6 +50,7 @@ save_adapter_weights_only: False
 
 optimizer:
   _component_: torch.optim.AdamW
+  fused: True
   lr: 2e-5
 
 lr_scheduler:

recipes/configs/gemma/2B_qlora_single_device.yaml

@@ -50,6 +50,7 @@ save_adapter_weights_only: False
 
 optimizer:
   _component_: torch.optim.AdamW
+  fused: True
   lr: 2e-5
 
 lr_scheduler:

recipes/configs/gemma/7B_full.yaml

@@ -50,6 +50,7 @@ batch_size: 1
 epochs: 1
 optimizer:
   _component_: torch.optim.AdamW
+  fused: True
   lr: 2e-5
 loss:
   _component_: torchtune.modules.loss.CEWithChunkedOutputLoss

recipes/configs/gemma/7B_lora.yaml

@@ -53,6 +53,7 @@ save_adapter_weights_only: False
 
 optimizer:
   _component_: torch.optim.AdamW
+  fused: True
   lr: 2e-5
 
 lr_scheduler:

recipes/configs/gemma/7B_lora_single_device.yaml

@@ -52,6 +52,7 @@ save_adapter_weights_only: False
 
 optimizer:
   _component_: torch.optim.AdamW
+  fused: True
   lr: 5e-5
 
 lr_scheduler:

recipes/configs/gemma/7B_qlora_single_device.yaml

@@ -52,6 +52,7 @@ save_adapter_weights_only: False
 
 optimizer:
   _component_: torch.optim.AdamW
+  fused: True
   lr: 2e-5
 
 lr_scheduler:

recipes/configs/llama2/13B_full.yaml

@@ -52,6 +52,7 @@ batch_size: 2
 epochs: 3
 optimizer:
   _component_: torch.optim.AdamW
+  fused: True
   lr: 2e-5
 loss:
   _component_: torchtune.modules.loss.CEWithChunkedOutputLoss

recipes/configs/llama2/13B_lora.yaml

@@ -60,6 +60,7 @@ batch_size: 2
 # Optimizer and Scheduler
 optimizer:
   _component_: torch.optim.AdamW
+  fused: True
   weight_decay: 0.01
   lr: 2e-4
 lr_scheduler:

recipes/configs/llama2/13B_qlora_single_device.yaml

@@ -55,6 +55,7 @@ batch_size: 2
 # Optimizer and Scheduler
 optimizer:
   _component_: torch.optim.AdamW
+  fused: True
   weight_decay: 0.01
   lr: 3e-4
 lr_scheduler:

recipes/configs/llama2/70B_lora.yaml

@@ -60,6 +60,7 @@ batch_size: 2
 # Optimizer and Scheduler
 optimizer:
   _component_: torch.optim.AdamW
+  fused: True
   weight_decay: 0.01
   lr: 3e-4
 lr_scheduler:

recipes/configs/llama2/70B_qlora.yaml

@@ -66,6 +66,7 @@ batch_size: 2
 # Optimizer and Scheduler
 optimizer:
   _component_: torch.optim.AdamW
+  fused: True
   weight_decay: 0.01
   lr: 3e-4
 lr_scheduler:

recipes/configs/llama2/7B_full.yaml

@@ -51,6 +51,7 @@ batch_size: 2
 epochs: 3
 optimizer:
   _component_: torch.optim.AdamW
+  fused: True
   lr: 2e-5
 loss:
   _component_: torchtune.modules.loss.CEWithChunkedOutputLoss

recipes/configs/llama2/7B_lora.yaml

@@ -57,6 +57,7 @@ batch_size: 2
 # Optimizer and Scheduler
 optimizer:
   _component_: torch.optim.AdamW
+  fused: True
   weight_decay: 0.01
   lr: 3e-4
 lr_scheduler:

recipes/configs/llama2/7B_lora_dpo.yaml

@@ -54,6 +54,7 @@ batch_size: 4
 # Optimizer and Scheduler
 optimizer:
   _component_: torch.optim.AdamW
+  fused: True
   weight_decay: 0.05
   lr: 5e-4
 lr_scheduler:

recipes/configs/llama2/7B_lora_dpo_single_device.yaml

@@ -53,6 +53,7 @@ batch_size: 4
 # Optimizer and Scheduler
 optimizer:
   _component_: torch.optim.AdamW
+  fused: True
   weight_decay: 0.05
   lr: 5e-4
 lr_scheduler:

recipes/configs/llama2/7B_lora_single_device.yaml

@@ -55,6 +55,7 @@ batch_size: 2
 # Optimizer and Scheduler
 optimizer:
   _component_: torch.optim.AdamW
+  fused: True
   weight_decay: 0.01
   lr: 3e-4
 lr_scheduler:

recipes/configs/llama2/7B_qat_full.yaml

@@ -47,6 +47,7 @@ batch_size: 2
 epochs: 3
 optimizer:
   _component_: torch.optim.AdamW
+  fused: True
   lr: 2e-5
 loss:
   _component_: torchtune.modules.loss.CEWithChunkedOutputLoss

recipes/configs/llama2/7B_qlora.yaml

@@ -57,6 +57,7 @@ batch_size: 2
 # Optimizer and Scheduler
 optimizer:
   _component_: torch.optim.AdamW
+  fused: True
   weight_decay: 0.01
   lr: 3e-4
 lr_scheduler:

recipes/configs/llama2/7B_qlora_single_device.yaml

@@ -54,6 +54,7 @@ batch_size: 2
 # Optimizer and Scheduler
 optimizer:
   _component_: torch.optim.AdamW
+  fused: True
   weight_decay: 0.01
   lr: 3e-4
 lr_scheduler:

recipes/configs/llama3/70B_full.yaml

@@ -79,8 +79,7 @@ epochs: 3
 optimizer:
   _component_: torch.optim.AdamW
   lr: 2e-5
-  foreach: False
-  # Note: highly recommended to use fused=True optimizer flag
+  fused: True  # Note: highly recommended to use fused=True optimizer flag
   # with CPU offload for faster optimizer step.
   fused: True
 

recipes/configs/llama3/70B_lora.yaml

@@ -75,6 +75,7 @@ batch_size: 2
 # Optimizer and Scheduler
 optimizer:
   _component_: torch.optim.AdamW
+  fused: True
   weight_decay: 0.01
   lr: 3e-4
 lr_scheduler:

recipes/configs/llama3/8B_dora.yaml

@@ -50,6 +50,7 @@ batch_size: 2
 # Optimizer and Scheduler
 optimizer:
   _component_: torch.optim.AdamW
+  fused: True
   weight_decay: 0.01
   lr: 3e-4
 lr_scheduler:

recipes/configs/llama3/8B_dora_single_device.yaml

@@ -52,6 +52,7 @@ batch_size: 2
 # Optimizer and Scheduler
 optimizer:
   _component_: torch.optim.AdamW
+  fused: True
   weight_decay: 0.01
   lr: 3e-4
 lr_scheduler:

recipes/configs/llama3/8B_full.yaml

@@ -52,8 +52,7 @@ epochs: 3
 optimizer:
   _component_: torch.optim.AdamW
   lr: 2e-5
-  foreach: False
-
+  fused: True
 loss:
   _component_: torchtune.modules.loss.CEWithChunkedOutputLoss
 max_steps_per_epoch: null

recipes/configs/llama3/8B_lora.yaml

@@ -55,6 +55,7 @@ batch_size: 2
 # Optimizer and Scheduler
 optimizer:
   _component_: torch.optim.AdamW
+  fused: True
   weight_decay: 0.01
   lr: 3e-4
 lr_scheduler:

recipes/configs/llama3/8B_lora_single_device.yaml

@@ -54,6 +54,7 @@ batch_size: 2
 # Optimizer and Scheduler
 optimizer:
   _component_: torch.optim.AdamW
+  fused: True
   weight_decay: 0.01
   lr: 3e-4
 lr_scheduler:

recipes/configs/llama3/8B_qat_full.yaml

@@ -52,8 +52,7 @@ quantizer:
 optimizer:
   _component_: torch.optim.AdamW
   lr: 2e-5
-  foreach: False
-
+  fused: True
 loss:
   _component_: torchtune.modules.loss.CEWithChunkedOutputLoss
 max_steps_per_epoch: null

recipes/configs/llama3/8B_qdora_single_device.yaml

@@ -53,6 +53,7 @@ batch_size: 2
 # Optimizer and Scheduler
 optimizer:
   _component_: torch.optim.AdamW
+  fused: True
   weight_decay: 0.01
   lr: 3e-4
 lr_scheduler:

recipes/configs/llama3/8B_qlora_single_device.yaml

@@ -53,6 +53,7 @@ batch_size: 2
 # Optimizer and Scheduler
 optimizer:
   _component_: torch.optim.AdamW
+  fused: True
   weight_decay: 0.01
   lr: 3e-4
 lr_scheduler:

recipes/configs/llama3_1/70B_full.yaml

@@ -79,7 +79,6 @@ epochs: 3
 optimizer:
   _component_: torch.optim.AdamW
   lr: 2e-5
-  foreach: False
   # Note: highly recommended to use fused=True optimizer flag
   # with CPU offload for faster optimizer step.
   fused: True

recipes/configs/llama3_1/70B_lora.yaml

@@ -74,6 +74,7 @@ batch_size: 2
 # Optimizer and Scheduler
 optimizer:
   _component_: torch.optim.AdamW
+  fused: True
   weight_decay: 0.01
   lr: 3e-4
 lr_scheduler:

recipes/configs/llama3_1/8B_full.yaml

@@ -55,8 +55,7 @@ epochs: 3
 optimizer:
   _component_: torch.optim.AdamW
   lr: 2e-5
-  foreach: False
-
+  fused: True
 loss:
   _component_: torchtune.modules.loss.CEWithChunkedOutputLoss
 max_steps_per_epoch: null

recipes/configs/llama3_1/8B_lora.yaml

@@ -58,6 +58,7 @@ batch_size: 2
 # Optimizer and Scheduler
 optimizer:
   _component_: torch.optim.AdamW
+  fused: True
   weight_decay: 0.01
   lr: 3e-4
 lr_scheduler:

recipes/configs/llama3_1/8B_lora_single_device.yaml

@@ -57,6 +57,7 @@ batch_size: 2
 # Optimizer and Scheduler
 optimizer:
   _component_: torch.optim.AdamW
+  fused: True
   weight_decay: 0.01
   lr: 3e-4
 lr_scheduler:

recipes/configs/llama3_1/8B_qlora_single_device.yaml

@@ -56,6 +56,7 @@ batch_size: 2
 # Optimizer and Scheduler
 optimizer:
   _component_: torch.optim.AdamW
+  fused: True
   weight_decay: 0.01
   lr: 3e-4
 lr_scheduler:

recipes/configs/mistral/7B_full.yaml

@@ -54,6 +54,7 @@ batch_size: 2
 epochs: 3
 optimizer:
   _component_: torch.optim.AdamW
+  fused: True
   lr: 5e-6
 loss:
   _component_: torchtune.modules.loss.CEWithChunkedOutputLoss

recipes/configs/mistral/7B_lora.yaml

@@ -59,6 +59,7 @@ save_adapter_weights_only: False
 
 optimizer:
   _component_: torch.optim.AdamW
+  fused: True
   lr: 2e-5
 
 lr_scheduler:

recipes/configs/mistral/7B_lora_single_device.yaml

@@ -56,6 +56,7 @@ save_adapter_weights_only: False
 
 optimizer:
   _component_: torch.optim.AdamW
+  fused: True
   lr: 2e-5
 
 lr_scheduler:

recipes/configs/mistral/7B_qlora_single_device.yaml

@@ -57,6 +57,7 @@ save_adapter_weights_only: False
 
 optimizer:
   _component_: torch.optim.AdamW
+  fused: True
   lr: 2e-5
 
 lr_scheduler:

recipes/configs/phi3/mini_full.yaml

@@ -53,6 +53,7 @@ batch_size: 2
 gradient_accumulation_steps: 16
 optimizer:
   _component_: torch.optim.AdamW
+  fused: True
   lr: 5e-6
 loss:
   _component_: torchtune.modules.loss.CEWithChunkedOutputLoss

recipes/configs/phi3/mini_lora.yaml

@@ -60,6 +60,7 @@ batch_size: 2
 gradient_accumulation_steps: 16
 optimizer:
   _component_: torch.optim.AdamW
+  fused: True
   weight_decay: 0.01
   lr: 3e-4
 lr_scheduler:

recipes/configs/phi3/mini_lora_single_device.yaml

@@ -58,6 +58,7 @@ batch_size: 2
 gradient_accumulation_steps: 16
 optimizer:
   _component_: torch.optim.AdamW
+  fused: True
   weight_decay: 0.01
   lr: 3e-4
 lr_scheduler:

recipes/configs/phi3/mini_qlora_single_device.yaml

@@ -58,6 +58,7 @@ batch_size: 2
 gradient_accumulation_steps: 16
 optimizer:
   _component_: torch.optim.AdamW
+  fused: True
   weight_decay: 0.01
   lr: 3e-4
 lr_scheduler: