Update on "compile optimizer"

Compiling optimizer helps perf of Llama4 Scout Model
3.8 tokens_per_second -> 9 tokens_per_second (max value of tokens per second in the first ~10 iterations)
peak memory is the same

```
tune run --nproc_per_node 8 \
  full_finetune_distributed \
  --config recipes/configs/llama4/scout_17B_16E_full.yaml
```


PS:
Current repo compilation fails if to set `skip_rope_interval=4,`, have to test with `skip_rope_interval=None,`

[ghstack-poisoned]

Author: IvanKobzarev

recipes/configs/llama4/scout_17B_16E_full.yaml

@@ -69,12 +69,14 @@ device: cuda
 enable_activation_checkpointing: True
 enable_activation_offloading: False
 fsdp_cpu_offload: True
-compile: False # torch.compile, set to true for perf/memory improvement
-
-compile_components:
-  model: True
-  loss: True
-  optimizer_step: False
+# compile True means use torch.compile for all components
+# compile False means no torch.compile
+# compile Dictionary with keys: "model", "loss", "optimizer_step"
+# enables torch.compile only for specified components.
+compile: False
+#    model: True
+#    loss: True
+#    optimizer_step: False
 
 # Reduced precision
 dtype: bf16

recipes/full_finetune_distributed.py

@@ -306,19 +306,17 @@ def setup(self, cfg: DictConfig) -> None:
         # Load the base model
         checkpoint_dict = self._checkpoint_client.load_base_checkpoint()
 
-        self._compile = cfg.get("compile", False)
+        compile = cfg.get("compile")
+        compile_bool = bool(compile)
         self._compile_backend = os.environ.get("TORCH_COMPILE_BACKEND", "inductor")
 
-        self._compile_model = False
-        self._compile_loss = False
-        self._compile_optimizer_step = False
-        compile_components = cfg.get("compile_components")
-        if self._compile and compile_components:
-            self._compile_model = compile_components.get("model", True)
-            self._compile_loss = compile_components.get("loss", True)
-            self._compile_optimizer_step = compile_components.get(
-                "optimizer_step", False
-            )
+        self._compile_model = compile_bool
+        self._compile_loss = compile_bool
+        self._compile_optimizer_step = compile_bool
+        if isinstance(compile, dict):
+            self._compile_model = compile.get("model", True)
+            self._compile_loss = compile.get("loss", True)
+            self._compile_optimizer_step = compile.get("optimizer_step", False)
 
         self._model = self._setup_model(
             cfg_model=cfg.model,