From 3ae0f7c08e17d609ffb074e19d5e8d34351a73d4 Mon Sep 17 00:00:00 2001
From: Wing Lian <wing@axolotl.ai>
Date: Fri, 23 May 2025 11:15:44 -0400
Subject: [PATCH] make sure torch_compile is enabled with SAC

---
 src/axolotl/utils/schemas/config.py | 13 +++++++++++++
 1 file changed, 13 insertions(+)

diff --git a/src/axolotl/utils/schemas/config.py b/src/axolotl/utils/schemas/config.py
index 6422bde82..283bc1727 100644
--- a/src/axolotl/utils/schemas/config.py
+++ b/src/axolotl/utils/schemas/config.py
@@ -1080,6 +1080,19 @@ class AxolotlInputConfig(
             )
         return data
 
+    @model_validator(mode="before")
+    @classmethod
+    def check_activation_memory_budget_w_compile(cls, data):
+        if data.get("activation_memory_budget") is not None and not data.get(
+            "torch_compile"
+        ):
+            LOG.warning(
+                "activation_memory_budget is enabled, but torch_compile is not set. "
+                "Automatically setting torch_compile to true."
+            )
+            data["torch_compile"] = True
+        return data
+
     @model_validator(mode="before")
     @classmethod
     def check_npu_config(cls, data):