fix train

scripts/train.py

@@ -79,7 +79,7 @@ def main():
         save_steps=1000,
         save_total_limit=2,
         bf16=True,
-        gradient_checkpointing=False,  # Disabled for now to debug
+        gradient_checkpointing=False,  # Disabled for now
         report_to="tensorboard",
         optim="adamw_torch",
         lr_scheduler_type="cosine",
@@ -112,27 +112,50 @@ def main():
         batch["output_router_logits"] = True
         return batch
 
-    # Custom trainer class to handle gradient flow
+    # Fixed CustomTrainer class
     class CustomTrainer(Trainer):
         def compute_loss(self, model, inputs, return_outputs=False):
-            outputs = model(**inputs)
-            loss = outputs.loss
+            # Remove num_items_in_batch from inputs if present
+            inputs.pop('num_items_in_batch', None)
             
-            # Ensure we have gradients
-            if loss.requires_grad:
+            with torch.set_grad_enabled(True):  # Ensure gradients are enabled
+                outputs = model(**inputs)
+                loss = outputs.loss
+                
+                if not loss.requires_grad:
+                    raise RuntimeError("Loss doesn't require gradients. Check model parameters.")
+                
                 return (loss, outputs) if return_outputs else loss
-            else:
-                raise RuntimeError("Loss doesn't require gradients. Check model parameters.")
 
     # Initialize trainer
     trainer = CustomTrainer(
         model=model,
         args=training_args,
         train_dataset=tokenized_dataset,
-        tokenizer=tokenizer,
         data_collator=data_collator,
     )
     
+    # Test forward/backward pass before training
+    print("Testing gradient flow...")
+    test_batch = next(iter(DataLoader(tokenized_dataset, batch_size=1)))
+    test_batch = {k: v.to(model.device) for k, v in test_batch.items()}
+    
+    model.train()
+    outputs = model(**test_batch)
+    loss = outputs.loss
+    print(f"Initial loss: {loss.item()}")
+    
+    loss.backward()
+    print("Gradients computed successfully")
+    
+    # Check which parameters received gradients
+    for name, param in model.named_parameters():
+        if param.grad is not None:
+            print(f"Parameter {name} received gradients")
+    
+    # Reset gradients
+    model.zero_grad()
+    
     # Train
     print("Starting training...")
     trainer.train()