Update HROM_Trainer.py

HROM_Trainer.py

@@ -28,30 +28,33 @@ logging.basicConfig(
 
 # Configuration
 CONFIG = {
+    # --- Scaled Parameters ---
     "dim": 768,
-    "n_layers": 8,
-    "n_heads": 8,
-    "ff_dim": 2048,
+    "n_layers": 16,
+    "n_heads": 16,
+    "ff_dim": 3072, # Explicitly set to 4 * dim
+
+    # --- Kept Parameters ---
     "dropout": 0.1,
     "max_seq_len": 512,
-    "batch_size": 16, # Keep batch size reasonable
+    "vocab_size": 32000, # Fixed by tokenizer
+
+    # --- Training/Dataset Parameters ---
+    "batch_size": 12,
     "checkpoint_interval": 2000,
     "debug_interval": 400,
-    # Reverted to training on all four datasets, using correct persona_chat identifier
+    # --- ADDED CoQA and QuAC ---
     "datasets": ["daily_dialog", "empathetic_dialogues", "blended_skill_talk", "AlekseyKorshuk/persona-chat"],
-    # Reverted to combined tokenizer name
-    "tokenizer_name": "hrom_tokenizer.json",
-    # Reverted to combined checkpoint dir
-    "checkpoint_dir": "checkpoints",
-    "vocab_size": 32000,
-    # Adjusted samples per dataset: with 4 datasets, 50k each gives 200k total samples
-    "tokenizer_train_samples_per_dataset": 50000,
-    "learning_rate": 2e-5,
+    "tokenizer_name": "hrom_tokenizer.json", # New name for expanded tokenizer
+    "checkpoint_dir": "checkpoints", # Separate directory for expanded data model
+    # --- Increased samples per dataset slightly for tokenizer ---
+    "tokenizer_train_samples_per_dataset": 100000, # Use same limit for all, incl. new ones
+    "learning_rate": 1e-5,
     "warmup_steps": 1000,
-    "max_turns": 8, # Max turns applied per dialogue
+    "max_turns": 8, # Keep max_turns limit for Q&A datasets too
     "max_checkpoints": 5,
     "num_epochs": 30,
-    "grad_accum_steps": 8 # Keep grad accum reasonable
+    "grad_accum_steps": 16
 }
 
 # --- Model Definition (HROM, HROMBlock, HROMAttention, SwiGLU, RoPE) ---