Merge pull request #119 from drunkenbot-ai/develop

update dimentions related issue

Author: drunkenbot-ai

engine

@@ -109,6 +109,13 @@ def _update_architecture_advisor(
             label = "Advisor: balanced"
             notes.append(f"Token budget is about {tokens_per_param:.1f} tokens per parameter.")
         if model_config is not None:
+            head_dim = model_config.embedding_size // max(1, model_config.head_count)
+            if head_dim % 8 != 0:
+                notes.append(
+                    f"Head dimension d_k = {head_dim} (n_embd={model_config.embedding_size} / n_head={model_config.head_count}) "
+                    "is not divisible by 8. FlashAttention will fall back to slow Math attention. Recommended: 64 or 128 (e.g. n_embd: 512, n_head: 8)."
+                )
+                label = f"Advisor: d_k={head_dim} (slow)"
             if model_config.context_length >= 2048 and model_config.embedding_size <= 256:
                 notes.append("Long context with a small embedding can be memory-heavy without adding much capacity.")
             if model_config.attention_type in {"grouped_query", "multi_query"}:

interface/screens/training_estimation.py

@@ -193,6 +193,19 @@ def _run_training_preflight(self, model_config: ModelConfig, training_config: Tr
             warnings.append("Free disk space is close to the estimated training storage need.")
         if checkpoint_count > 50:
             warnings.append("Save interval may create many checkpoints. Increase Save every or clean old checkpoints.")
+        head_dim = model_config.embedding_size // max(1, model_config.head_count)
+        if head_dim % 8 != 0:
+            warnings.append(
+                f"Head dimension d_k = {head_dim} (n_embd={model_config.embedding_size} / n_head={model_config.head_count}) "
+                "is not divisible by 8. PyTorch FlashAttention will fall back to slow Math attention. "
+                "For 10x-20x faster training, adjust settings so d_k is 64 or 128 (e.g. n_embd: 512, n_head: 8)."
+            )
+        if training_config.compile_model and training_config.activation_checkpointing:
+            warnings.append(
+                "Both Torch compile and Activation checkpointing are enabled. "
+                "Activation checkpointing forces extra forward passes and can slow down compilation. "
+                "Disable Activation checkpointing unless GPU memory is nearly exhausted."
+            )
 
         log.clear()
         log.append("Training checklist")

interface/screens/training_run.py

@@ -87,13 +87,23 @@ def build_training_tab(window) -> QWidget:
         "Leave checked unless you are deliberately experimenting with this.",
     )
     window.n_embd = window._spin(32, 4096, 128)
-    window._tip(window.n_embd, "Embedding size, also called n_embd. Larger values increase model capacity and memory usage.")
+    window._tip(
+        window.n_embd,
+        "Embedding size (n_embd). Larger values increase capacity and memory usage. "
+        "For maximum GPU FlashAttention speed, ensure (n_embd / n_head) is a multiple of 8, "
+        "ideally 64 or 128 (e.g. 512, 768).",
+    )
     window.architecture_style.currentTextChanged.connect(
         lambda text: window.rope_theta.setEnabled(text == "Llama-like")
     )
     window.rope_theta.setEnabled(window.architecture_style.currentText() == "Llama-like")
     window.n_head = window._spin(1, 64, 4)
-    window._tip(window.n_head, "Attention head count. More heads can model varied relationships, but n_embd must divide evenly by n_head.")
+    window._tip(
+        window.n_head,
+        "Attention head count. n_embd must divide evenly by n_head. "
+        "For fast GPU FlashAttention, head dimension (n_embd / n_head) must be a multiple of 8, "
+        "ideally 64 or 128 (e.g. 512 / 8 = 64, 768 / 12 = 64).",
+    )
     window.attention_type = QComboBox()
     window.attention_type.addItems(["Multi-head", "Grouped-query", "Multi-query"])
     window.attention_type.setMaximumWidth(260)

interface/tabs/training_tab.py

@@ -10,7 +10,7 @@
 from interface.core.project_state import ProjectStateMixin
 from interface.core.project_state_apply import ProjectStateApplyMixin
 
-from engine.config import DatasetConfig
+from engine.config import DatasetConfig, ModelConfig, TrainingConfig
 from engine.data import Document
 from engine.dataset_corpus import _StreamingCorpusBuilder
 from engine.dataset_mixture import (
@@ -25,7 +25,7 @@
     token_dtype_for_vocab,
     train_tokenizer,
 )
-from engine.training import TokenDataset
+from engine.training import TokenDataset, train_model
 
 
 class DiversityFilterTests(unittest.TestCase):
@@ -181,5 +181,60 @@ def test_validation_stride_equals_context_length(self) -> None:
         self.assertEqual(y1.tolist(), list(range(65, 129)))
 
 
+class TrainingDiagnosticsAndTelemetryTests(unittest.TestCase):
+    """Tests for preflight warnings and telemetry emission enhancements."""
+
+    def test_head_dim_not_divisible_by_eight_detected(self) -> None:
+        model_config = ModelConfig(
+            vocab_size=32,
+            context_length=16,
+            embedding_size=560,
+            head_count=8,
+            layer_count=2,
+        )
+        head_dim = model_config.embedding_size // model_config.head_count
+        self.assertEqual(head_dim, 70)
+        self.assertNotEqual(head_dim % 8, 0)
+
+    def test_milestone_step_emits_event(self) -> None:
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            model_config = ModelConfig(
+                vocab_size=16,
+                context_length=8,
+                embedding_size=16,
+                head_count=2,
+                layer_count=1,
+                dropout=0.0,
+            )
+            training_config = TrainingConfig(
+                output_dir=Path(tmp_dir),
+                epochs=1,
+                batch_size=1,
+                learning_rate=1e-3,
+                sample_stride=8,
+                warmup_steps=0,
+                eval_interval=0,
+                save_interval=0,
+                use_amp=False,
+                precision="fp32",
+                device="cpu",
+                resume=False,
+                early_stopping=False,
+            )
+            events: list[dict] = []
+            train_model(
+                model_config,
+                training_config,
+                [index % 16 for index in range(32)],
+                [],
+                pad_token_id=-1,
+                progress=events.append,
+            )
+            step_events = [e for e in events if e.get("event_type") == "step"]
+            self.assertGreaterEqual(len(step_events), 1)
+            self.assertIn("Step 1", step_events[0]["message"])
+
+
 if __name__ == "__main__":
     unittest.main()
+