Fixes

Model_Architecture/model.py

@@ -505,6 +505,7 @@ class ismail(nn.Module):
         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)
 
@@ -527,13 +528,13 @@ class ismail(nn.Module):
             mask = torch.hstack([torch.zeros((seqlen, start_pos), device=tokens.device), mask]).type_as(h)
 
         total_lb_loss = 0.0
-        
+
         for layer in self.layers:
             layer.start_pos = start_pos
             layer.freqs_cis = freqs_cis
             layer.mask = mask
 
-            if self.training and True:  # Enable gradient checkpointing during training
+            if self.training and self.use_checkpointing:
                 from torch.utils.checkpoint import checkpoint
                 h, lb_loss = checkpoint(layer.checkpoint_forward, h)
             else:

Model_Architecture/train.py

@@ -136,28 +136,16 @@ def load_config(args):
 
 
 def setup_model(config, device):
-    """Initialize model and print size estimate"""
     args = ModelArgs(**config["model"])
-    
-    print("\n" + "="*70)
-    print("MODEL INITIALIZATION")
-    print("="*70 + "\n")
-    
-    # Estimate size
-    #size_info = estimate_model_size(args)
-
     model = ismail(args).to(device)
-
-    if config["training"].get("use_checkpointing", True):
-        for layer in model.layers:
-            layer.forward = lambda *args, layer=layer: checkpoint(layer._forward, *args)
-        print("✅ Gradient checkpointing enabled")
     
-    # Compile for speed (PyTorch 2.0+)
+    # Add this line to enable checkpointing
+    model.use_checkpointing = config["training"].get("use_checkpointing", True)
+    
     if config["training"]["compile"]:
         try:
             model = torch.compile(model)
-            print("✅ Model compiled with torch.compile()\n")
+            print("✅ Model compiled\n")
         except Exception as e:
             print(f"⚠️  Compilation failed: {e}\n")