Merge pull request #117 from drunkenbot-ai/develop

fix triton related issue

Author: drunkenbot-ai

engine

@@ -101,6 +101,7 @@ def _default_project_state(self) -> dict[str, Any]:
                 "data_loader_workers": 0,
                 "max_grad_norm": 1.0,
                 "activation_checkpointing": False,
+                "compile_model": False,
                 "seed": 1337,
                 "device": self.device.currentText(),
                 "use_amp": self.use_amp_default,
@@ -350,6 +351,7 @@ def _project_state_dict(self, project_name: str, project_dir: Path) -> dict[str,
                 "data_loader_workers": self.data_loader_workers.value(),
                 "max_grad_norm": self.max_grad_norm.value(),
                 "activation_checkpointing": self.activation_checkpointing.isChecked(),
+                "compile_model": self.compile_model.isChecked() if hasattr(self, "compile_model") else False,
                 "seed": self.seed.value(),
                 "device": self.device.currentText(),
                 "use_amp": self.use_amp.isChecked(),

interface/core/project_state.py

@@ -171,6 +171,8 @@ def _apply_project_state(self, data: dict[str, Any]) -> None:
         self.data_loader_workers.setValue(int(training.get("data_loader_workers", self.data_loader_workers.value())))
         self.max_grad_norm.setValue(float(training.get("max_grad_norm", self.max_grad_norm.value())))
         self.activation_checkpointing.setChecked(bool(training.get("activation_checkpointing", False)))
+        if hasattr(self, "compile_model"):
+            self.compile_model.setChecked(bool(training.get("compile_model", False)))
         self.seed.setValue(int(training.get("seed", self.seed.value())))
         self._set_combo_text(self.device, str(training.get("device", self.device.currentText())))
         self.use_amp.setChecked(bool(training.get("use_amp", self.use_amp.isChecked())))

interface/core/project_state_apply.py

@@ -67,6 +67,7 @@ def _current_training_config(
             data_loader_workers=self.data_loader_workers.value(),
             max_grad_norm=self.max_grad_norm.value(),
             activation_checkpointing=self.activation_checkpointing.isChecked(),
+            compile_model=self.compile_model.isChecked() if hasattr(self, "compile_model") else False,
             device=self.device.currentText(),
             use_amp=self.use_amp.isChecked(),
             precision=self._precision_value(),

interface/screens/training_config.py

@@ -195,6 +195,12 @@ def build_training_tab(window) -> QWidget:
         window.activation_checkpointing,
         "Recompute transformer activations during backpropagation to lower VRAM use. Training becomes slower, but this is useful when memory is the constraint.",
     )
+    window.compile_model = QCheckBox("Torch compile")
+    window.compile_model.setChecked(False)
+    window._tip(
+        window.compile_model,
+        "Compile model using PyTorch Inductor for kernel fusion. Requires supported CUDA environment and Triton.",
+    )
     window.seed = window._spin(1, 2_147_483_647, 1337)
     window._tip(window.seed, "Random seed for reproducible initialization and sampling order.")
     window.device = QComboBox()
@@ -263,6 +269,7 @@ def build_training_tab(window) -> QWidget:
     right.addRow("CPU workers", window.data_loader_workers)
     right.addRow("Max grad", window.max_grad_norm)
     right.addRow("VRAM saver", window.activation_checkpointing)
+    right.addRow("Kernel fusion", window.compile_model)
     right.addRow("Seed", window.seed)
     runtime = QFormLayout()
     window._configure_form(runtime)

interface/tabs/training_tab.py

@@ -0,0 +1,37 @@
+from __future__ import annotations
+
+import unittest
+from pathlib import Path
+from unittest.mock import MagicMock, patch
+
+import torch
+
+from engine.config import ModelConfig, TrainingConfig
+from engine.training_impl import _try_compile_model
+
+
+class CompileModelTests(unittest.TestCase):
+    def test_compile_model_defaults_to_false(self) -> None:
+        config = TrainingConfig(output_dir=Path("tmp"))
+        self.assertFalse(config.compile_model)
+
+    def test_try_compile_model_returns_eager_when_disabled(self) -> None:
+        model = torch.nn.Linear(10, 10)
+        compiled = _try_compile_model(model, "cuda", enabled=False)
+        self.assertIs(compiled, model)
+
+    def test_try_compile_model_skips_on_cpu(self) -> None:
+        model = torch.nn.Linear(10, 10)
+        compiled = _try_compile_model(model, "cpu", enabled=True)
+        self.assertIs(compiled, model)
+
+    @patch("torch.compile", side_effect=RuntimeError("Inductor compiler error"))
+    def test_try_compile_model_catches_compiler_error_and_falls_back(self, _mock_compile) -> None:
+        model = torch.nn.Linear(10, 10)
+        with patch("torch.cuda.is_available", return_value=True):
+            compiled = _try_compile_model(model, "cuda", enabled=True)
+        self.assertIs(compiled, model)
+
+
+if __name__ == "__main__":
+    unittest.main()