Compile FFT FSDP (#1573)

Co-authored-by: Felipe Mello <felipemello@fb.com>
pytorch · Sep 13, 2024 · ee343e6 · ee343e6
1 parent 4d3e48a
commit ee343e6
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/recipes/full_finetune_distributed.py b/recipes/full_finetune_distributed.py
@@ -362,6 +362,9 @@ def _setup_model(
         with training.set_default_dtype(self._dtype), torch.device("meta"):
             model = config.instantiate(cfg_model)
 
+        if self._compile:
+            training.compile_model(model, verbose=self._is_rank_zero)
+
         # We currently have two versions of activation checkpointing in this recipe
         # for testing and BC purposes. ``enable_activation_checkpointing`` controls
         # the older version of AC and this behavior is unchanged