Small changes to reduce peak memory. (#389)

Co-authored-by: Taylor Robie <taylor.robie@lightning.ai>
Co-authored-by: Luca Antiga <luca@lightning.ai>

Author: 3 people

finetune/full.py

@@ -55,7 +55,7 @@ def main(
 ):
 
     auto_wrap_policy = partial(transformer_auto_wrap_policy, transformer_layer_cls={Block})
-    strategy = FSDPStrategy(auto_wrap_policy=auto_wrap_policy, activation_checkpointing=Block)
+    strategy = FSDPStrategy(auto_wrap_policy=auto_wrap_policy, activation_checkpointing=Block, limit_all_gathers=True)
 
     fabric = L.Fabric(accelerator="cuda", devices=devices, precision="bf16-mixed", strategy=strategy)
     fabric.launch()
@@ -79,7 +79,7 @@ def main(
 
     model = fabric.setup_module(model)
 
-    optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate)
+    optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate, foreach=False)
     optimizer = fabric.setup_optimizers(optimizer)
 
     train(fabric, model, optimizer, train_data, val_data, out_dir)

pretrain/redpajama.py

@@ -69,7 +69,7 @@ def main(
         transformer_auto_wrap_policy, transformer_layer_cls={Block}
     )
     strategy = FSDPStrategy(
-        auto_wrap_policy=auto_wrap_policy, activation_checkpointing=Block
+        auto_wrap_policy=auto_wrap_policy, activation_checkpointing=Block, limit_all_gathers=True
     )
 
     fabric = L.Fabric(
@@ -110,6 +110,7 @@ def main(
         lr=learning_rate,
         weight_decay=weight_decay,
         betas=(beta1, beta2),
+        foreach=False,
     )
 
     model, optimizer = fabric.setup(model, optimizer)

pretrain/shakespeare.py

@@ -47,7 +47,7 @@
 
 def main() -> None:
     auto_wrap_policy = partial(transformer_auto_wrap_policy, transformer_layer_cls={Block})
-    strategy = FSDPStrategy(auto_wrap_policy=auto_wrap_policy, activation_checkpointing=Block)
+    strategy = FSDPStrategy(auto_wrap_policy=auto_wrap_policy, activation_checkpointing=Block, limit_all_gathers=True)
 
     fabric = L.Fabric(accelerator="cuda", devices=4, precision="bf16-mixed", strategy=strategy)
     fabric.launch()
@@ -70,7 +70,7 @@ def main() -> None:
 
     model = fabric.setup_module(model)
 
-    optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate, weight_decay=weight_decay, betas=(beta1, beta2))
+    optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate, weight_decay=weight_decay, betas=(beta1, beta2), foreach=False)
     optimizer = fabric.setup_optimizers(optimizer)
 
     train(fabric, model, optimizer, train_data, val_data)