Fixes

Model_Architecture/config.json

@@ -30,7 +30,7 @@
         "tokenizer_name": "turkish"
     },
     "training": {
-        "learning_rate": 3e-4,
+        "learning_rate": 3e-5,
         "weight_decay": 0.1,
         "beta1": 0.9,
         "beta2": 0.95,

Model_Architecture/model.py

@@ -129,7 +129,7 @@ def linear(x: torch.Tensor, weight: torch.Tensor, bias: Optional[torch.Tensor] =
 
 
 class Linear(nn.Module):
-    dtype = torch.bfloat16
+    dtype = torch.float32
     scale_fmt: Optional[str] = None
 
     def __init__(self, in_features: int, out_features: int, bias: bool = False, dtype = None):
@@ -187,7 +187,7 @@ class RMSNorm(nn.Module):
         super().__init__()
         self.dim = dim
         self.eps = eps
-        self.weight = nn.Parameter(torch.ones(dim, dtype=torch.bfloat16)) 
+        self.weight = nn.Parameter(torch.ones(dim, dtype=torch.float32)) 
 
     def forward(self, x: torch.Tensor):
         output = F.rms_norm(x, (self.dim,), self.weight, self.eps)
@@ -500,12 +500,17 @@ class ismail(nn.Module):
         self.n_layers = args.n_layers
 
         self.tok_embeddings = nn.Embedding(args.vocab_size, args.dim)
+        nn.init.normal_(self.tok_embeddings.weight, mean=0.0, std=0.02)
+        
         self.layers = nn.ModuleList([Block(i, args) for i in range(args.n_layers)])
         self.norm = RMSNorm(args.dim)
         self.output = Linear(args.dim, args.vocab_size, bias=False)
         self.use_checkpointing = False
 
         self.register_buffer("freqs_cis", precompute_freqs_cis(args), persistent=False)
+        
+        if hasattr(self.output, 'weight'):
+            nn.init.normal_(self.output.weight, mean=0.0, std=0.02 / math.sqrt(args.n_layers))
 
     def set_active_expert(self, expert_idx: Optional[int]):
         """Set active expert for all MoE layers (for sequential training)"""

Model_Architecture/train.py

@@ -294,9 +294,21 @@ def save_checkpoint(model, optimizer, step, config, expert_idx=None):
 def train_step(model, input_mb, target_mb, device, config, scaler=None):
     """Process a SINGLE micro-batch (already sliced)"""
     
+    # 🚨 Validate data before processing
     if input_mb.size(0) == 0:
         return 0.0, 0.0
     
+    # Check for invalid token IDs (outside vocab range)
+    vocab_size = config["model"]["vocab_size"]
+    if input_mb.max() >= vocab_size or target_mb.max() >= vocab_size:
+        print(f"🚨 Invalid token detected! Max token: {input_mb.max().item()}, Vocab size: {vocab_size}")
+        return 0.0, 0.0
+    
+    # Check for NaN in data
+    if torch.isnan(input_mb).any() or torch.isnan(target_mb).any():
+        print("🚨 NaN detected in input data!")
+        return 0.0, 0.0
+    
     input_mb = input_mb.to(device, non_blocking=True)
     target_mb = target_mb.to(device, non_blocking=True)
 
@@ -309,9 +321,11 @@ def train_step(model, input_mb, target_mb, device, config, scaler=None):
             logits = output
             lb_loss = 0.0
         
-        # 🚨 DEBUG: Check for NaN in logits
+        # 🚨 Check for NaN in logits before computing loss
         if torch.isnan(logits).any():
-            print(f"🚨 NaN detected in logits! Scale: {logits.abs().max().item():.2f}")
+            print(f"🚨 NaN detected in logits! Scale: {logits.abs().max().item()}")
+            print(f"    Input range: [{input_mb.min().item()}, {input_mb.max().item()}]")
+            return 0.0, 0.0
         
         lm_loss = F.cross_entropy(
             logits.view(-1, logits.size(-1)),
@@ -319,19 +333,22 @@ def train_step(model, input_mb, target_mb, device, config, scaler=None):
             ignore_index=-1,
         )
         
+        # 🚨 Check for NaN in loss components
+        if torch.isnan(lm_loss):
+            print(f"🚨 NaN in lm_loss!")
+            return 0.0, 0.0
+        
         accum_steps = config["training"]["gradient_accumulation_steps"]
         if isinstance(lb_loss, float):
             total_loss = lm_loss / accum_steps
         else:
+            if torch.isnan(lb_loss):
+                print(f"🚨 NaN in lb_loss! Setting to 0")
+                lb_loss = 0.0
             lb_loss_coef = config["training"].get("lb_loss_coef", 0.01)
             total_loss = (lm_loss + lb_loss_coef * lb_loss) / accum_steps
 
-    # 🚨 DEBUG: Check for NaN in total loss
-    if torch.isnan(total_loss):
-        print(f"🚨 NaN in total_loss! lm_loss: {lm_loss.item():.4f}, lb_loss: {lb_loss}")
-        return 0.0, 0.0 
-
-    # Backward
+    # Backward with NaN check
     if scaler is not None:
         scaler.scale(total_loss).backward()
     else: