From 6d55527085e386b2e25d6213f95bbeacb6709de8 Mon Sep 17 00:00:00 2001
From: Hossein Kavianihamedani <hosseinkh@meta.com>
Date: Tue, 23 Dec 2025 11:18:59 -0800
Subject: [PATCH] Move compile config from training to top-level

- Move compile from training.compile to compile.enable in llama3_8b.yaml
- Move compile from training.compile to compile.enable in qwen3_8b.yaml
- Update main.py to use job_config.compile.enable instead of job_config.training.compile

This aligns the YAML config structure with ForgeJobConfig's expected schema,
where compile is a separate top-level dataclass, not nested under training.
---
 apps/sft/llama3_8b.yaml | 4 +++-
 apps/sft/main.py        | 6 +++---
 apps/sft/qwen3_8b.yaml  | 4 +++-
 3 files changed, 9 insertions(+), 5 deletions(-)

diff --git a/apps/sft/llama3_8b.yaml b/apps/sft/llama3_8b.yaml
index d9a5a9783..f82c26e1b 100644
--- a/apps/sft/llama3_8b.yaml
+++ b/apps/sft/llama3_8b.yaml
@@ -26,12 +26,14 @@ optimizer:
 lr_scheduler:
   warmup_steps: 200
 
+compile:
+  enable: false
+
 training:
   local_batch_size: 8
   seq_len: 2048
   max_norm: 1.0
   steps: 1000
-  compile: false
   datasets:
     - path: "yahma/alpaca-cleaned"
       split: "train[:95%]"
diff --git a/apps/sft/main.py b/apps/sft/main.py
index 4f2a7be74..c39145893 100644
--- a/apps/sft/main.py
+++ b/apps/sft/main.py
@@ -96,11 +96,11 @@ def record_batch_metrics(self, data_metrics: list):
     @endpoint
     async def setup(self):
         # Validate that compile is only used with flex attention
-        if self.job_config.training.compile:
+        if self.job_config.compile.enable:
             raise ValueError(
-                "training.compile=True is not currently supported. "
+                "compile.enable=True is not currently supported. "
                 "Compile is only supported with flex attention enabled, which requires PyTorch nightly. "
-                "Please set training.compile=false in your config."
+                "Please set compile.enable=false in your config."
             )
 
         # all ranks should record loss, except when PP=True. Then, only the last stage should record loss.
diff --git a/apps/sft/qwen3_8b.yaml b/apps/sft/qwen3_8b.yaml
index a8d2244a9..a90c4fd23 100644
--- a/apps/sft/qwen3_8b.yaml
+++ b/apps/sft/qwen3_8b.yaml
@@ -25,12 +25,14 @@ optimizer:
 lr_scheduler:
   warmup_steps: 200
 
+compile:
+  enable: false
+
 training:
   local_batch_size: 8
   seq_len: 2048
   max_norm: 1.0
   steps: 1000
-  compile: false
   datasets:
     - path: "yahma/alpaca-cleaned"
       split: "train[:95%]"