From 00065320d86e4720773be9e277ee157c357ae2e8 Mon Sep 17 00:00:00 2001
From: 51616 <rujikorn.ch@gmail.com>
Date: Sat, 21 Jun 2025 01:24:05 +0900
Subject: [PATCH] compile max autotune

---
 intx_sft.py | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/intx_sft.py b/intx_sft.py
index 6c5fafe..feb526f 100755
--- a/intx_sft.py
+++ b/intx_sft.py
@@ -240,7 +240,8 @@ def main():
             raise ValueError("base model contains trainable parameters")
         # we can't compile the base model bc we will be
         # adding/removing forward hooks during training
-        model.hypernet = torch.compile(model.hypernet)
+        model.hypernet.compile(fullgraph=True, mode="max-autotune")
+
     else:
         # activate LoRA
         base_model_config = AutoConfig.from_pretrained(
@@ -424,9 +425,9 @@ def main():
 
         _apply_liger_kernel_to_instance(model=base_model.model)
         # compile_linear(base_model)
-        base_model.compile(fullgraph=True)
+        base_model.compile(fullgraph=True, mode="max-autotune")
         base_model.loss_function = torch.compile(
-            base_model.loss_function, fullgraph=True
+            base_model.loss_function, fullgraph=True, mode="max-autotune"
         )
 
     if LOCAL_RANK == 0: