make sure torch_compile is enabled with SAC

This commit is contained in:
Wing Lian
2025-05-23 11:15:44 -04:00
parent 5930c91a12
commit 3ae0f7c08e

View File

@@ -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):